From 0d6d42a4304b25a04399ae7f2cf74a53b2511bb0 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Thu, 5 Jun 2025 19:46:19 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20sshsvc=E6=9C=8D=E5=8A=A1=E7=9A=84?=
 =?UTF-8?q?=E7=A7=81=E9=92=A5=E6=96=87=E4=BB=B6=E5=88=A4=E6=96=AD=E7=94=9F?=
 =?UTF-8?q?=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 sshsvc/sshsvc.go | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/sshsvc/sshsvc.go b/sshsvc/sshsvc.go
index 52d3802a..d0526ff8 100644
--- a/sshsvc/sshsvc.go
+++ b/sshsvc/sshsvc.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
@@ -50,24 +51,41 @@ func init() {
 	logmml.InitMmlLogger(conf.Logmml.File, conf.Logmml.Duration, conf.Logmml.Count, "omc", config.GetLogMmlLevel())
 }
 
-func main() {
-	// 生成SSH密钥对
+// readPrivateKey 读取SSH私钥，如果不存在则生成新的密钥对
+func readPrivateKey() ssh.Signer {
+	// 检查私钥文件是否存在
+	if _, err := os.Stat(conf.Sshd.PrivateKey); os.IsNotExist(err) {
+		// 如果文件不存在，创建目录并生成密钥
+		dir := filepath.Dir(conf.Sshd.PrivateKey)
+		if err := os.MkdirAll(dir, 0700); err != nil {
+			log.Fatal("Failed to create .ssh directory:", err)
+			os.Exit(2)
+		}
+
+		// 使用ssh-keygen命令生成密钥对
+		cmd := exec.Command("ssh-keygen", "-t", "rsa", "-P", "", "-f", conf.Sshd.PrivateKey)
+		if err := cmd.Run(); err != nil {
+			log.Fatal("Failed to generate SSH key:", err)
+			os.Exit(2)
+		}
+	}
+
+	// 读取SSH密钥对
 	privateKeyBytes, err := os.ReadFile(conf.Sshd.PrivateKey)
 	if err != nil {
-		// ssh-keygen -t rsa -P "" -f /root/.ssh/id_rsa
-		// ssh-keygen -t rsa -f /etc/ssh/ssh_host_rsa_key -N "" -y
-		exec.Command("ssh-keygen", "-t", "rsa", "-P", "", "-f", conf.Sshd.PrivateKey, "-N", "", "-y").Run()
 		log.Fatal("Failed to ReadFile", err)
 		os.Exit(2)
 	}
 
 	privateKey, err := ssh.ParsePrivateKey(privateKeyBytes)
 	if err != nil {
-		exec.Command("ssh-keygen", "-t", "rsa", "-P", "", "-f", conf.Sshd.PrivateKey, "-N", "", "-y").Run()
 		log.Fatal("Failed to ParsePrivateKey", err)
 		os.Exit(3)
 	}
+	return privateKey
+}
 
+func main() {
 	// 配置SSH服务器
 	serverConfig := &ssh.ServerConfig{
 		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
@@ -106,6 +124,7 @@ func main() {
 		},
 	}
 
+	privateKey := readPrivateKey()
 	serverConfig.AddHostKey(privateKey)
 
 	// 启动SSH服务器