pm-4G/src/main/java/com/nokia/common/ssh/SSHUtil.java

package com.nokia.common.ssh;

import com.jcraft.jsch.*;
import com.nokia.common.ssh.entity.SSHServer;
import com.nokia.common.ssh.entity.UserInfoImpl;
import com.nokia.common.ssh.exception.SSHUtilException;
import com.nokia.common.ssh.exception.ScpAckErrorException;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.io.*;
import java.nio.file.NoSuchFileException;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.Vector;

/**
 * 使用jsch库实现的ssh的工具类
 * <p>
 * todo: scpTo和scpFrom 在本机和targetServer默认编码不一致的时候,文件名中的中文会乱码,但是不会影响到文件内容,
 */

@Slf4j
public class SSHUtil {

    @Getter
    @Setter
    private SSHServer targetServer = new SSHServer();
    private Session session = null;
    private Channel channel = null;
    private JSch jSch = null;
    private FileInputStream fileInputStream = null;
    private FileOutputStream fileOutputStream = null;
    private OutputStream outputStream = null;
    private InputStream inputStream = null;
    private ChannelSftp channelSftp = null;

    public SSHUtil() {
    }

    public SSHUtil(String host, String user, String password) {
        targetServer = new SSHServer(host, 22, user, password);
    }

    public SSHUtil(String host, Integer port, String user, String password) {
        targetServer = new SSHServer(host, port, user, password);
    }

    /**
     * 获取文件列表
     */
    public List<String> ls(String path) throws JSchException, SftpException {
        session = getConnectSession();
        channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();
        List<String> fileNameList = new ArrayList<>();
        Vector fileList = channelSftp.ls(path);
        for (Object o : fileList) {
            String fileName = ((ChannelSftp.LsEntry) o).getFilename();
            if (".".equals(fileName) || "..".equals(fileName)) {
                continue;
            }

            fileNameList.add(fileName);
            channelSftp.quit();
            session.disconnect();
        }
        return fileNameList;
    }

    /**
     * 删除文件
     */
    public void delete(String fileName) throws JSchException, SftpException {
        session = getConnectSession();
        channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();
        channelSftp.rm(fileName);
        channelSftp.quit();
        session.disconnect();
    }

    /**
     * 远程执行指令
     */
    public String exec(String command) throws JSchException, IOException {
        StringBuilder stringBuilder = new StringBuilder();
        try {
            session = getConnectSession();
            channel = session.openChannel("exec");
            // jsch的登陆是无环境登陆即非login状态登陆,因此是没有环境变量的,
            String execCommand;
            // 在命令前添加 bash --login -c "command"以获取环境变量
            // source .bashrc && command 也可以解决问题, 但是可能环境加载不全
            if (command.startsWith("bash --login -c")) {
                execCommand = command;
            } else {
                execCommand = String.format("bash --login -c \"%s\"", command);
            }
            ((ChannelExec) channel).setCommand(execCommand);
            channel.setInputStream(null);
            ((ChannelExec) channel).setErrStream(System.err);
            InputStream in = channel.getInputStream();
            channel.connect();
            byte[] tmp = new byte[1024];
            while (true) {
                while (in.available() > 0) {
                    int i = in.read(tmp, 0, 1024);
                    if (i < 0) {
                        break;
                    }
                    stringBuilder.append(new String(tmp, 0, i));
                }
                if (channel.isClosed()) {
                    if (in.available() > 0) {
                        continue;
                    }
                    break;
                }
            }
        } finally {
            disconnect();
        }
        return stringBuilder.toString();
    }

    /**
     * 使用SCP把本地文件推送到targetServer目录下
     * <p>
     * 注意,文件名不能包含中文
     */
    public boolean scpTo(String sourceFilePath, String targetPath) throws JSchException, IOException, SSHUtilException {
        try {
            session = getConnectSession();
            // scp内置了两个参数 -t 和 -f ，这两个参数是隐藏的，不会被用户显式提供，
            // 两个scp进程之间传输数据时，远端机器上的scp进程被本地scp进程启动起来时提供上去。
            // 需要说明的是，这是通过本地scp进程经ssh远程过去开启远端机器的scp进程来实现的。
            // -t 指定为to 也就是目的端模式 指定的对象就是session对应的连接对象targetServer
            String command = "scp " + "-t " + targetPath;
            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);
            outputStream = channel.getOutputStream();
            inputStream = channel.getInputStream();
            channel.connect();
            if (checkAck(inputStream) != 0) {
                log.error("scpTo 执行失败");
                return false;
            }
            File sourceFile = new File(sourceFilePath);
            if (sourceFile.isDirectory()) {
                log.error("sourceFilePath 必须是文件");
                return false;
            }
            long fileSize = sourceFile.length();
            command = "C0644 " + fileSize + " " + sourceFile.getName() + "\n";
            outputStream.write(command.getBytes());
            outputStream.flush();
            if (checkAck(inputStream) != 0) {
                log.error("scpTo 执行失败");
                return false;
            }
            fileInputStream = new FileInputStream(sourceFile);
            byte[] buffer = new byte[1024];
            while (true) {
                int len = fileInputStream.read(buffer, 0, buffer.length);
                if (len <= 0) {
                    break;
                }
                outputStream.write(buffer, 0, len);
            }
            buffer[0] = 0;
            outputStream.write(buffer, 0, 1);
            outputStream.flush();
            return checkAck(inputStream) == 0;
        } finally {
            disconnect();
        }
    }

    /**
     * 使用scp把targetServer目录下的文件复制到本地
     */
    public boolean scpFrom(String sourceFilePath, String targetPath) throws JSchException, IOException, SSHUtilException {
        try {
            log.debug(sourceFilePath);
            session = getConnectSession();
            // scp内置了两个参数 -t 和 -f ，这两个参数是隐藏的，不会被用户显式提供，
            // 两个scp进程之间传输数据时，远端机器上的scp进程被本地scp进程启动起来时提供上去。
            // 需要说明的是，这是通过本地scp进程经ssh远程过去开启远端机器的scp进程来实现的。
            // -f 指定对端为from 也就是源端模式 指定的对象就是session对应的连接对象targetServer
            String command = "scp -f " + sourceFilePath;
            Channel channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);
            outputStream = channel.getOutputStream();
            inputStream = channel.getInputStream();
            channel.connect();
            byte[] buf = new byte[1024];
            // 发送指令 '0'
            // 源端会一直等宿端的回应, 直到等到回应才会传输下一条协议文本.
            // 在送出最后一条协议文本后, 源端会传出一个大小为零的字符'0'来表示真正文件传输的开始.
            // 当文件接收完成后, 宿端会给源端发送一个'0'
            buf[0] = 0;
            outputStream.write(buf, 0, 1);
            outputStream.flush();
            // 接收C0644 这条消息携带了文件的信息
            while (true) {
                int c = checkAck(inputStream);
                // 遇到C时跳出循环
                if (c == 'C') {
                    break;
                }
            }
            // 接收 '0644 ' 这段字符表示文件的权限
            inputStream.read(buf, 0, 5);
            // 获取filesize
            long filesize = 0L;
            while (true) {
                if (inputStream.read(buf, 0, 1) < 0) {
                    break;
                }
                if (buf[0] == ' ') {
                    break;
                }
                filesize = filesize * 10L + (long) (buf[0] - '0');
            }
            // 从 C0644命令读取文件名,命令中的文件名是不带路径的
            String file = null;
            for (int i = 0; ; i++) {
                inputStream.read(buf, i, 1);
                // 0x0a 是LF 换行符
                if (buf[i] == (byte) 0x0a) {
                    file = new String(buf, 0, i);
                    break;
                }
            }
            log.debug("filesize={}, file={}", filesize, file);
            // 发送 '0'
            buf[0] = 0;
            outputStream.write(buf, 0, 1);
            outputStream.flush();
            // 如果目标是目录,则需要加上文件名
            File target = new File(targetPath);
            if (target.isDirectory()) {
                log.debug("{} 是目录,需要添加文件名", target.getAbsolutePath());
                target = new File(targetPath + File.separator + file);
            }

            fileOutputStream = new FileOutputStream(target);
            int foo;
            while (true) {
                if (buf.length < filesize) {
                    foo = buf.length;
                } else {
                    foo = (int) filesize;
                }
                foo = inputStream.read(buf, 0, foo);
                if (foo < 0) {
                    break;
                }
                fileOutputStream.write(buf, 0, foo);
                filesize -= foo;
                if (filesize == 0L) {
                    break;
                }
            }
            if (checkAck(inputStream) != 0) {
                return false;
            }
            // 发送 '0'
            buf[0] = 0;
            outputStream.write(buf, 0, 1);
            outputStream.flush();
            log.debug("scp from {}@{}:{}{} to {} 完成", targetServer.getUser(), targetServer.getHost(), targetServer.getPort(), sourceFilePath, target.getAbsolutePath());
            return true;
        } finally {
            disconnect();
        }
    }

    private Session getConnectSession() throws JSchException {
        jSch = new JSch();
        session = jSch.getSession(targetServer.getUser(), targetServer.getHost(), targetServer.getPort());
        session.setPassword(targetServer.getPassword());
        session.setUserInfo(new UserInfoImpl());
        // 不需要输入保存ssh安全密钥的yes或no
        Properties properties = new Properties();
        properties.put("StrictHostKeyChecking", "no");
        session.setConfig(properties);
        session.connect();
        log.debug("已连接到{}@{}:{}", targetServer.getUser(), targetServer.getHost(), targetServer.getPort());
        return session;
    }

    private void disconnect() throws IOException {
        if (fileOutputStream != null) {
            fileOutputStream.close();
            fileOutputStream = null;
        }
        if (fileInputStream != null) {
            fileInputStream.close();
            fileInputStream = null;
        }
        if (outputStream != null) {
            outputStream.close();
            outputStream = null;
        }
        if (channel != null) {
            channel.disconnect();
            channel = null;
        }
        if (session != null) {
            session.disconnect();
            session = null;
        }
        jSch = null;
    }

    /**
     * 来自源端的每条消息和每个传输完毕的文件都需要宿端的确认和响应.
     * 宿端会返回三种确认消息: 0(正常), 1(警告)或2(严重错误, 将中断连接).
     * 消息1和2可以跟一个字符串和一个换行符, 这个字符串将显示在scp的源端. 无论这个字符串是否为空, 换行符都是不可缺少的.
     */
    private static int checkAck(InputStream in) throws IOException, SSHUtilException {
        int b = in.read();
        // b 取值为0表示成功
        if (b == 0) {
            return b;
        }
        if (b == -1) {
            return b;
        }

        // 1表示警告 2表示严重错误,将中断连接
        // 1和2 后面会携带一条错误信息,以\n结尾
        if (b == 1 || b == 2) {
            // 打印消息后面跟的字符串
            StringBuilder sb = new StringBuilder();
            int c;
            do {
                // 读取字符串直到遇到换行符
                c = in.read();
                sb.append((char) c);
            } while (c != '\n');
            log.debug("checkAck发现错误消息: ack={}-msg={}", b, sb);
            if (b == 1 && sb.toString().endsWith("No such file or directory")) {
                throw new NoSuchFileException(sb.toString());
            } else {
                throw new ScpAckErrorException(sb.toString());
            }
        }
        return b;
    }
}