ok-1

internal/system/ssh.go

@@ -0,0 +1,187 @@
+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()
+}