feat: 添加自动发现功能及相关配置

.gitignore

@@ -4,4 +4,5 @@ agent.json
 komari-agent.exe
 komari-agent
 .komari-agent.exe.old
-build/
+build/
+auto-discovery.json

cmd/autodiscovery.go

@@ -0,0 +1,185 @@
+package cmd
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"net/url"
+	"os"
+	"path/filepath"
+
+	"github.com/komari-monitor/komari-agent/cmd/flags"
+)
+
+// AutoDiscoveryConfig 自动发现配置结构体
+type AutoDiscoveryConfig struct {
+	UUID  string `json:"uuid"`
+	Token string `json:"token"`
+}
+
+// RegisterRequest 注册请求结构体
+type RegisterRequest struct {
+	Key string `json:"key"`
+}
+
+// RegisterResponse 注册响应结构体
+type RegisterResponse struct {
+	Status  string `json:"status"`
+	Message string `json:"message"`
+	Data    struct {
+		UUID  string `json:"uuid"`
+		Token string `json:"token"`
+	} `json:"data"`
+}
+
+// getAutoDiscoveryFilePath 获取自动发现配置文件路径
+func getAutoDiscoveryFilePath() string {
+	// 获取程序运行目录
+	execPath, err := os.Executable()
+	if err != nil {
+		log.Println("Failed to get executable path:", err)
+		return "auto-discovery.json"
+	}
+	execDir := filepath.Dir(execPath)
+	return filepath.Join(execDir, "auto-discovery.json")
+}
+
+// loadAutoDiscoveryConfig 加载自动发现配置
+func loadAutoDiscoveryConfig() (*AutoDiscoveryConfig, error) {
+	configPath := getAutoDiscoveryFilePath()
+
+	// 检查文件是否存在
+	if _, err := os.Stat(configPath); os.IsNotExist(err) {
+		return nil, nil // 文件不存在，返回nil
+	}
+
+	// 读取文件内容
+	data, err := os.ReadFile(configPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read auto-discovery config: %v", err)
+	}
+
+	// 解析JSON
+	var config AutoDiscoveryConfig
+	if err := json.Unmarshal(data, &config); err != nil {
+		return nil, fmt.Errorf("failed to parse auto-discovery config: %v", err)
+	}
+
+	return &config, nil
+}
+
+// saveAutoDiscoveryConfig 保存自动发现配置
+func saveAutoDiscoveryConfig(config *AutoDiscoveryConfig) error {
+	configPath := getAutoDiscoveryFilePath()
+
+	// 序列化为JSON
+	data, err := json.MarshalIndent(config, "", "  ")
+	if err != nil {
+		return fmt.Errorf("failed to marshal auto-discovery config: %v", err)
+	}
+
+	// 写入文件
+	if err := os.WriteFile(configPath, data, 0644); err != nil {
+		return fmt.Errorf("failed to write auto-discovery config: %v", err)
+	}
+
+	log.Printf("Auto-discovery config saved to: %s", configPath)
+	return nil
+}
+
+// registerWithAutoDiscovery 使用自动发现key注册
+func registerWithAutoDiscovery() error {
+	// 构造注册请求
+	requestData := RegisterRequest{
+		Key: flags.AutoDiscoveryKey,
+	}
+
+	hostname, _ := os.Hostname()
+
+	jsonData, err := json.Marshal(requestData)
+	if err != nil {
+		return fmt.Errorf("failed to marshal register request: %v", err)
+	}
+
+	// 构造请求URL
+	endpoint := flags.Endpoint
+	if len(endpoint) > 0 && endpoint[len(endpoint)-1] == '/' {
+		endpoint = endpoint[:len(endpoint)-1]
+	}
+	registerURL := fmt.Sprintf("%s/api/clients/register?name=%s", endpoint, url.QueryEscape(hostname))
+
+	// 创建HTTP请求
+	req, err := http.NewRequest("POST", registerURL, bytes.NewBuffer(jsonData))
+	if err != nil {
+		return fmt.Errorf("failed to create register request: %v", err)
+	}
+
+	// 设置请求头
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", flags.AutoDiscoveryKey))
+
+	// 发送请求
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("failed to send register request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	// 检查响应状态
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("register request failed with status %d: %s", resp.StatusCode, string(body))
+	}
+
+	// 解析响应
+	var registerResp RegisterResponse
+	if err := json.NewDecoder(resp.Body).Decode(&registerResp); err != nil {
+		return fmt.Errorf("failed to parse register response: %v", err)
+	}
+
+	// 检查响应状态
+	if registerResp.Status != "success" {
+		return fmt.Errorf("register request failed: %s", registerResp.Message)
+	}
+
+	// 保存配置
+	config := &AutoDiscoveryConfig{
+		UUID:  registerResp.Data.UUID,
+		Token: registerResp.Data.Token,
+	}
+
+	if err := saveAutoDiscoveryConfig(config); err != nil {
+		return fmt.Errorf("failed to save auto-discovery config: %v", err)
+	}
+
+	// 设置token
+	flags.Token = registerResp.Data.Token
+	log.Printf("Successfully registered with auto-discovery. UUID: %s", registerResp.Data.UUID)
+
+	return nil
+}
+
+// handleAutoDiscovery 处理自动发现逻辑
+func handleAutoDiscovery() error {
+	// 尝试加载现有配置
+	config, err := loadAutoDiscoveryConfig()
+	if err != nil {
+		log.Printf("Failed to load auto-discovery config: %v", err)
+		// 继续尝试注册
+	}
+
+	if config != nil {
+		// 配置文件存在，使用现有token
+		flags.Token = config.Token
+		log.Printf("Using existing auto-discovery token for UUID: %s", config.UUID)
+		return nil
+	}
+
+	// 配置文件不存在，进行注册
+	log.Println("Auto-discovery config not found, registering with server...")
+	return registerWithAutoDiscovery()
+}

cmd/flags/flag.go

@@ -1,6 +1,7 @@
 package flags
 
 var (
+	AutoDiscoveryKey     string
 	DisableAutoUpdate   bool
 	DisableWebSsh       bool
 	MemoryModeAvailable bool

cmd/root.go

@@ -20,6 +20,14 @@ var RootCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		log.Println("Komari Agent", update.CurrentVersion)
 		log.Println("Github Repo:", update.Repo)
+		// Auto discovery
+		if flags.AutoDiscoveryKey != "" {
+			err := handleAutoDiscovery()
+			if err != nil {
+				log.Printf("Auto-discovery failed: %v", err)
+				os.Exit(1)
+			}
+		}
 		diskList, err := monitoring.DiskList()
 		if err != nil {
 			log.Println("Failed to get disk list:", err)
@@ -68,9 +76,10 @@ func Execute() {
 
 func init() {
 	RootCmd.PersistentFlags().StringVarP(&flags.Token, "token", "t", "", "API token")
-	RootCmd.MarkPersistentFlagRequired("token")
+	//RootCmd.MarkPersistentFlagRequired("token")
 	RootCmd.PersistentFlags().StringVarP(&flags.Endpoint, "endpoint", "e", "", "API endpoint")
 	RootCmd.MarkPersistentFlagRequired("endpoint")
+	RootCmd.PersistentFlags().StringVar(&flags.AutoDiscoveryKey, "auto-discovery", "", "Auto discovery key for the agent")
 	RootCmd.PersistentFlags().BoolVar(&flags.DisableAutoUpdate, "disable-auto-update", false, "Disable automatic updates")
 	RootCmd.PersistentFlags().BoolVar(&flags.DisableWebSsh, "disable-web-ssh", false, "Disable remote control(web ssh and rce)")
 	RootCmd.PersistentFlags().BoolVar(&flags.MemoryModeAvailable, "memory-mode-available", false, "Report memory as available instead of used.")

update/update.go

@@ -60,7 +60,7 @@ func CheckAndUpdate() error {
 
 	// Determine if update is needed
 	if latest.Version.Equals(currentSemVer) {
-		fmt.Println("Current version is the latest:", CurrentVersion)
+		log.Println("Current version is the latest:", CurrentVersion)
 		return nil
 	}
 	// Default is installed as a service, so don't automatically restart
@@ -75,7 +75,7 @@ func CheckAndUpdate() error {
 	// if err != nil {
 	// 	return fmt.Errorf("failed to restart program: %v", err)
 	// }
-	fmt.Printf("Successfully updated to version %s\n", latest.Version)
+	log.Printf("Successfully updated to version %s\n", latest.Version)
 	os.Exit(42)
 	return nil
 }