package cn.reghao.bnt.web.devops.machine.util;

import cn.reghao.jutil.jdk.serializer.JsonConverter;
import cn.reghao.jutil.jdk.thread.ThreadPoolWrapper;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;

/**
 * @author reghao
 * @date 2021-07-07 13:19:42
 */
@Slf4j
@Service
public class WebSsh {
    private static final Map<String, SshConnInfo> sshConnMap = new ConcurrentHashMap<>();
    private final ExecutorService threadPool = ThreadPoolWrapper.threadPool("webssh-service");

    public void init(WebSocketSession session) {
        JSch jSch = new JSch();
        SshConnInfo connInfo = new SshConnInfo();
        connInfo.setWebSocketSession(session);
        connInfo.setJSch(jSch);
        String uuid = String.valueOf(session.getAttributes().get(Keys.USER_UUID));
        sshConnMap.put(uuid, connInfo);
    }

    /**
     * 处理客户端发送的消息
     *
     * @param
     * @return
     * @date 2021-07-07 下午4:45
     */
    public void handleMessage(String buffer, WebSocketSession session) {
        SshConnData sshConnData = JsonConverter.jsonToObject(buffer, SshConnData.class);
        String uuid = String.valueOf(session.getAttributes().get(Keys.USER_UUID));
        SshConnInfo sshConnInfo = sshConnMap.get(uuid);
        if (sshConnData.getOps().equals(Keys.OPS_CONNECT)) {
            threadPool.submit(() -> {
                try {
                    establishSshConn(sshConnInfo, sshConnData, session);
                } catch (JSchException | IOException e) {
                    log.error("SSH 连接建立失败 -> {}", e.getMessage());
                }
            });
        } else if (sshConnData.getOps().equals(Keys.OPS_COMMAND)) {
            String cmd = sshConnData.getCommand();
            if (cmd == null) {
                log.error("command 操作的参数为 null");
                return;
            }
            sendToSsh(sshConnInfo.getChannel(), sshConnData.getCommand());
        } else {
            log.error("无法处理 {} 操作", sshConnData.getOps());
        }
    }

    /**
     * 返回消息给客户端
     *
     * @param
     * @return
     * @date 2021-07-07 下午4:45
     */
    public void sendBack(WebSocketSession session, byte[] buffer) throws IOException {
        session.sendMessage(new TextMessage(buffer));
    }

    public void close(WebSocketSession session) {
        String uuid = String.valueOf(session.getAttributes().get(Keys.USER_UUID));
        SshConnInfo connInfo = sshConnMap.get(uuid);
        if (connInfo != null) {
            if (connInfo.getChannel() != null) {
                connInfo.getChannel().disconnect();
            }
            sshConnMap.remove(uuid);
        }
    }

    private void establishSshConn(SshConnInfo connInfo, SshConnData connData, WebSocketSession session) throws JSchException, IOException {
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");

        Session sshSession = connInfo.getJSch()
                .getSession(connData.getUsername(), connData.getHost(), connData.getPort());
        sshSession.setConfig(config);
        sshSession.setPassword(connData.getPassword());
        sshSession.connect(30_000);
        Channel channel = sshSession.openChannel("shell");
        channel.connect(30_000);
        connInfo.setChannel(channel);
        sendToSsh(channel, "\r");

        try (InputStream input = channel.getInputStream()) {
            byte[] buffer = new byte[1024];
            int i;
            while ((i = input.read(buffer)) != -1) {
                sendBack(session, Arrays.copyOfRange(buffer, 0, i));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            sshSession.disconnect();
            channel.disconnect();
        }
    }

    private void sendToSsh(Channel channel, String cmd) {
        if (channel != null && channel.isConnected()) {
            try {
                OutputStream out = channel.getOutputStream();
                out.write(cmd.getBytes(StandardCharsets.UTF_8));
                out.flush();
            } catch (IOException e) {
                log.error("向 SSH 发送数据失败 -> {}", e.getMessage());
            }
        }
    }
}
