diff --git a/.gitignore b/.gitignore index 8767c86..67fc407 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ agent.json komari-agent.exe komari-agent .komari-agent.exe.old -build/ \ No newline at end of file +build/ +auto-discovery.json \ No newline at end of file diff --git a/cmd/autodiscovery.go b/cmd/autodiscovery.go new file mode 100644 index 0000000..0403999 --- /dev/null +++ b/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(®isterResp); 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() +} diff --git a/cmd/flags/flag.go b/cmd/flags/flag.go index 98c0949..14d2a43 100644 --- a/cmd/flags/flag.go +++ b/cmd/flags/flag.go @@ -1,6 +1,7 @@ package flags var ( + AutoDiscoveryKey string DisableAutoUpdate bool DisableWebSsh bool MemoryModeAvailable bool diff --git a/cmd/root.go b/cmd/root.go index 34233d6..8efcb77 100644 --- a/cmd/root.go +++ b/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.") diff --git a/update/update.go b/update/update.go index 21a2208..1966305 100644 --- a/update/update.go +++ b/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 }