sunhpc-go/internal/system/ssh.go

package system

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"sunhpc/internal/config"
)

// ConfigureSSH 配置 SSH 服务
// 参数: cfg - config.SSHConfig 结构体
// 返回: error - 配置错误
func ConfigureSSH(cfg config.SSHConfig) error {
	const sshdConfig = "/etc/ssh/sshd_config"

	// 读取现有配置
	content, err := os.ReadFile(sshdConfig)
	if err != nil {
		return fmt.Errorf("读取 sshd_config 失败: %v", err)
	}

	// 备份原始配置
	backupPath := sshdConfig + ".sunhpc.bak"
	if _, err := os.Stat(backupPath); os.IsNotExist(err) {
		if err := os.WriteFile(backupPath, content, 0644); err != nil {
			fmt.Printf("警告: 无法创建备份文件 %s: %v\n", backupPath, err)
		}
	}

	// 解析和修改配置
	lines := strings.Split(string(content), "\n")
	newLines := make([]string, 0, len(lines))
	configMap := make(map[string]bool)

	// 处理每一行
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.HasPrefix(trimmed, "#") {
			newLines = append(newLines, line)
			continue
		}

		parts := strings.Fields(trimmed)
		if len(parts) < 2 {
			newLines = append(newLines, line)
			continue
		}

		key := parts[0]
		configMap[key] = true

		// 根据配置更新
		switch key {
		case "PermitRootLogin":
			if cfg.PermitRootLogin != "" {
				newLines = append(newLines, fmt.Sprintf("PermitRootLogin %s", cfg.PermitRootLogin))
			} else {
				newLines = append(newLines, line)
			}
		case "PasswordAuthentication":
			if cfg.PasswordAuth != "" {
				newLines = append(newLines, fmt.Sprintf("PasswordAuthentication %s", cfg.PasswordAuth))
			} else {
				newLines = append(newLines, line)
			}
		default:
			newLines = append(newLines, line)
		}
	}

	// 添加缺失的配置项
	if cfg.PermitRootLogin != "" && !configMap["PermitRootLogin"] {
		newLines = append(newLines, fmt.Sprintf("PermitRootLogin %s", cfg.PermitRootLogin))
	}
	if cfg.PasswordAuth != "" && !configMap["PasswordAuthentication"] {
		newLines = append(newLines, fmt.Sprintf("PasswordAuthentication %s", cfg.PasswordAuth))
	}

	// 写入新配置
	newContent := strings.Join(newLines, "\n")
	if err := os.WriteFile(sshdConfig, []byte(newContent), 0644); err != nil {
		return fmt.Errorf("写入 sshd_config 失败: %v", err)
	}

	// 测试配置语法
	if err := testSSHDConfig(); err != nil {
		// 恢复备份
		if backup, err := os.ReadFile(backupPath); err == nil {
			os.WriteFile(sshdConfig, backup, 0644)
		}
		return fmt.Errorf("SSH 配置语法错误: %v", err)
	}

	// 重启 SSH 服务
	return restartSSHD()
}

// testSSHDConfig 测试 sshd 配置语法
func testSSHDConfig() error {
	cmd := exec.Command("sshd", "-t")
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// restartSSHD 重启 SSH 服务
func restartSSHD() error {
	// 尝试不同的服务管理器
	serviceMgrs := []struct {
		name string
		args []string
	}{
		{"systemctl", []string{"restart", "sshd"}},
		{"systemctl", []string{"restart", "ssh"}},
		{"service", []string{"sshd", "restart"}},
		{"service", []string{"ssh", "restart"}},
	}

	for _, mgr := range serviceMgrs {
		if _, err := exec.LookPath(mgr.name); err == nil {
			cmd := exec.Command(mgr.name, mgr.args...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			if err := cmd.Run(); err == nil {
				return nil
			}
		}
	}

	return fmt.Errorf("无法重启 SSH 服务，请手动重启")
}

// AddSSHKey 添加 SSH 公钥到指定用户
func AddSSHKey(username, pubkey string) error {
	// 获取用户主目录
	homeDir, err := getUserHomeDir(username)
	if err != nil {
		return err
	}

	sshDir := homeDir + "/.ssh"
	authKeys := sshDir + "/authorized_keys"

	// 创建 .ssh 目录
	if err := os.MkdirAll(sshDir, 0700); err != nil {
		return err
	}

	// 追加公钥
	f, err := os.OpenFile(authKeys, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.WriteString(pubkey + "\n")
	if err != nil {
		return err
	}

	// 修改所有权
	return chownRecursive(sshDir, username)
}

// getUserHomeDir 获取用户主目录
func getUserHomeDir(username string) (string, error) {
	cmd := exec.Command("getent", "passwd", username)
	output, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("用户 %s 不存在", username)
	}

	parts := strings.Split(strings.TrimSpace(string(output)), ":")
	if len(parts) >= 6 {
		return parts[5], nil
	}
	return "", fmt.Errorf("无法获取用户 %s 的主目录", username)
}

// chownRecursive 递归修改文件所有者
func chownRecursive(path, username string) error {
	cmd := exec.Command("chown", "-R", username+":"+username, path)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}