ssh-manager/backend/src/main/java/com/sshmanager/service/SshBootstrapService.java

package com.sshmanager.service;

import com.jcraft.jsch.JSch;
import com.jcraft.jsch.KeyPair;
import com.sshmanager.dto.ConnectionCreateRequest;
import com.sshmanager.entity.Connection;
import com.sshmanager.exception.InvalidOperationException;
import org.springframework.stereotype.Service;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.time.Instant;

@Service
public class SshBootstrapService {

    private final SshService sshService;

    public SshBootstrapService(SshService sshService) {
        this.sshService = sshService;
    }

    public BootstrapResult bootstrapWithPassword(ConnectionCreateRequest request, Long userId) {
        String bootstrapPassword = trimToNull(request.getBootstrapPassword());
        if (bootstrapPassword == null) {
            throw new InvalidOperationException("启用一键免密配置时必须填写初始登录密码");
        }

        GeneratedKeyPair keyPair = generateKeyPair(buildKeyComment(userId, request.getName()));
        Connection passwordConnection = buildConnection(request, Connection.AuthType.PASSWORD);
        Connection privateKeyConnection = buildConnection(request, Connection.AuthType.PRIVATE_KEY);

        authorizePublicKey(passwordConnection, bootstrapPassword, keyPair.getPublicKey());
        verifyPrivateKeyLogin(privateKeyConnection, keyPair.getPrivateKey());

        return new BootstrapResult(keyPair.getPrivateKey());
    }

    private GeneratedKeyPair generateKeyPair(String comment) {
        KeyPair keyPair = null;
        try {
            JSch jsch = new JSch();
            keyPair = KeyPair.genKeyPair(jsch, KeyPair.RSA, 2048);

            ByteArrayOutputStream privateKeyOutput = new ByteArrayOutputStream();
            ByteArrayOutputStream publicKeyOutput = new ByteArrayOutputStream();
            keyPair.writePrivateKey(privateKeyOutput);
            keyPair.writePublicKey(publicKeyOutput, comment);

            return new GeneratedKeyPair(
                privateKeyOutput.toString(StandardCharsets.UTF_8.name()),
                publicKeyOutput.toString(StandardCharsets.UTF_8.name()).trim()
            );
        } catch (Exception e) {
            throw new InvalidOperationException("免密初始化失败：无法生成 SSH 密钥");
        } finally {
            if (keyPair != null) {
                keyPair.dispose();
            }
        }
    }

    private void authorizePublicKey(Connection connection, String bootstrapPassword, String publicKey) {
        String command = buildAuthorizeCommand(publicKey);
        try {
            SshService.CommandResult result = sshService.executeCommandWithResult(
                connection,
                bootstrapPassword,
                null,
                null,
                command
            );
            if (result.getExitStatus() != 0) {
                throw new InvalidOperationException(buildRemoteFailureMessage(
                    "免密初始化失败：无法写入远端 authorized_keys",
                    result
                ));
            }
        } catch (InvalidOperationException e) {
            throw e;
        } catch (Exception e) {
            throw new InvalidOperationException("免密初始化失败：密码登录或公钥下发失败" + formatCauseMessage(e));
        }
    }

    private void verifyPrivateKeyLogin(Connection connection, String privateKey) {
        try {
            SshService.CommandResult result = sshService.executeCommandWithResult(
                connection,
                null,
                privateKey,
                null,
                "printf 'ssh-manager bootstrap ok'"
            );
            if (result.getExitStatus() != 0) {
                throw new InvalidOperationException(buildRemoteFailureMessage(
                    "免密初始化失败：公钥已下发，但私钥验证失败",
                    result
                ));
            }
        } catch (InvalidOperationException e) {
            throw e;
        } catch (Exception e) {
            throw new InvalidOperationException("免密初始化失败：公钥已下发，但私钥验证失败" + formatCauseMessage(e));
        }
    }

    private String buildAuthorizeCommand(String publicKey) {
        String escapedPublicKey = shellQuote(publicKey);
        String innerCommand =
            "umask 077 && " +
            "mkdir -p ~/.ssh && chmod 700 ~/.ssh && " +
            "touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && " +
            "{ grep -Fqx " + escapedPublicKey + " ~/.ssh/authorized_keys || " +
            "printf '%s\\n' " + escapedPublicKey + " >> ~/.ssh/authorized_keys; }";
        return "sh -lc " + shellQuote(innerCommand);
    }

    private String buildRemoteFailureMessage(String prefix, SshService.CommandResult result) {
        String stderr = trimToNull(result.getStderr());
        String stdout = trimToNull(result.getStdout());
        String detail = stderr != null ? stderr : stdout;
        if (detail == null) {
            return prefix;
        }
        return prefix + "：" + detail;
    }

    private String buildKeyComment(Long userId, String connectionName) {
        String sanitizedName = sanitizeForComment(connectionName);
        long timestamp = Instant.now().getEpochSecond();
        if (sanitizedName == null) {
            return "ssh-manager-" + userId + "-" + timestamp;
        }
        return "ssh-manager-" + userId + "-" + timestamp + "-" + sanitizedName;
    }

    private String sanitizeForComment(String value) {
        String trimmed = trimToNull(value);
        if (trimmed == null) {
            return null;
        }

        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < trimmed.length(); i++) {
            char ch = trimmed.charAt(i);
            if ((ch >= 'a' && ch <= 'z')
                || (ch >= 'A' && ch <= 'Z')
                || (ch >= '0' && ch <= '9')
                || ch == '.'
                || ch == '_'
                || ch == '-') {
                builder.append(ch);
            } else if (builder.length() == 0 || builder.charAt(builder.length() - 1) != '-') {
                builder.append('-');
            }
        }

        String sanitized = builder.toString().replaceAll("^-+|-+$", "");
        return sanitized.isEmpty() ? null : sanitized;
    }

    private Connection buildConnection(ConnectionCreateRequest request, Connection.AuthType authType) {
        Connection connection = new Connection();
        connection.setHost(trimToNull(request.getHost()));
        connection.setPort(request.getPort() != null ? request.getPort() : 22);
        connection.setUsername(trimToNull(request.getUsername()));
        connection.setAuthType(authType);
        return connection;
    }

    private String shellQuote(String value) {
        return "'" + value.replace("'", "'\"'\"'") + "'";
    }

    private String formatCauseMessage(Exception e) {
        String message = trimToNull(e.getMessage());
        if (message == null) {
            return "";
        }
        return "：" + message;
    }

    private String trimToNull(String value) {
        if (value == null) {
            return null;
        }
        String trimmed = value.trim();
        return trimmed.isEmpty() ? null : trimmed;
    }

    public static class BootstrapResult {
        private final String privateKey;

        public BootstrapResult(String privateKey) {
            this.privateKey = privateKey;
        }

        public String getPrivateKey() {
            return privateKey;
        }
    }

    private static class GeneratedKeyPair {
        private final String privateKey;
        private final String publicKey;

        private GeneratedKeyPair(String privateKey, String publicKey) {
            this.privateKey = privateKey;
            this.publicKey = publicKey;
        }

        public String getPrivateKey() {
            return privateKey;
        }

        public String getPublicKey() {
            return publicKey;
        }
    }
}