package main import ( "crypto/tls" "encoding/json" "fmt" "io" "log" "net/http" "strings" "time" "github.com/gorilla/websocket" "github.com/komari-monitor/komari-agent/config" "github.com/komari-monitor/komari-agent/monitoring" "github.com/komari-monitor/komari-agent/update" ) var ( CurrentVersion string = "0.0.1" repo = "komari-monitor/komari-agent" ) func main() { log.Printf("Komari Agent %s\n", CurrentVersion) localConfig, err := config.LoadConfig() if err != nil { log.Fatalln("Failed to load local config:", err) } if localConfig.IgnoreUnsafeCert { http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true} } go func() { err = uploadBasicInfo(localConfig.Endpoint, localConfig.Token) ticker := time.NewTicker(time.Duration(time.Minute * 15)) for range ticker.C { err = uploadBasicInfo(localConfig.Endpoint, localConfig.Token) if err != nil { log.Fatalln("Failed to upload basic info:", err) } } }() websocketEndpoint := strings.TrimSuffix(localConfig.Endpoint, "/") + "/api/clients/report?token=" + localConfig.Token websocketEndpoint = "ws" + strings.TrimPrefix(websocketEndpoint, "http") var conn *websocket.Conn defer func() { if conn != nil { conn.Close() } }() ticker := time.NewTicker(time.Duration(localConfig.Interval * float64(time.Second))) defer ticker.Stop() if localConfig.AutoUpdate { go func() { ticker_ := time.NewTicker(time.Duration(6) * time.Hour) update_komari() for range ticker_.C { update_komari() } }() } for range ticker.C { // If no connection, attempt to connect if conn == nil { log.Println("Attempting to connect to WebSocket...") retry := 0 for retry < localConfig.MaxRetries { conn, err = connectWebSocket(websocketEndpoint) if err == nil { log.Println("WebSocket connected") go handleWebSocketMessages(localConfig, conn, make(chan struct{})) break } retry++ time.Sleep(time.Duration(localConfig.ReconnectInterval) * time.Second) } if retry >= localConfig.MaxRetries { log.Println("Max retries reached, falling back to POST") // Send report via POST and continue data := report(localConfig) if err := reportWithPOST(localConfig.Endpoint, data); err != nil { log.Println("Failed to send POST report:", err) } continue } } // Send report via WebSocket data := report(localConfig) err = conn.WriteMessage(websocket.TextMessage, data) if err != nil { log.Println("Failed to send WebSocket message:", err) conn.Close() conn = nil // Mark connection as dead continue } } } func update_komari() { // 初始化 Updater updater := update.NewUpdater(CurrentVersion, repo) // 检查并更新 err := updater.CheckAndUpdate() if err != nil { log.Println("Update Failed: %v", err) } } // connectWebSocket attempts to establish a WebSocket connection and upload basic info func connectWebSocket(websocketEndpoint string) (*websocket.Conn, error) { dialer := &websocket.Dialer{ HandshakeTimeout: 5 * time.Second, } conn, _, err := dialer.Dial(websocketEndpoint, nil) if err != nil { return nil, err } return conn, nil } func handleWebSocketMessages(localConfig config.LocalConfig, conn *websocket.Conn, done chan<- struct{}) { defer close(done) for { _, message_raw, err := conn.ReadMessage() if err != nil { log.Println("WebSocket read error:", err) return } // TODO: Remote config update // TODO: Handle incoming messages log.Println("Received message:", string(message_raw)) message := make(map[string]interface{}) err = json.Unmarshal(message_raw, &message) if err != nil { log.Println("Bad ws message:", err) continue } } } func reportWithPOST(endpoint string, data []byte) error { url := strings.TrimSuffix(endpoint, "/") + "/api/clients/report" req, err := http.NewRequest("POST", url, strings.NewReader(string(data))) if err != nil { return err } req.Header.Set("Content-Type", "application/json") client := &http.Client{} resp, err := client.Do(req) if err != nil { return err } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return err } return nil } func uploadBasicInfo(endpoint string, token string) error { log.Println("Uploading basic info...") defer log.Println("Upload complete") cpu := monitoring.Cpu() osname := monitoring.OSName() ipv4, ipv6, _ := monitoring.GetIPAddress() data := map[string]interface{}{ "cpu_name": cpu.CPUName, "cpu_cores": cpu.CPUCores, "arch": cpu.CPUArchitecture, "os": osname, "ipv4": ipv4, "ipv6": ipv6, "mem_total": monitoring.Ram().Total, "swap_total": monitoring.Swap().Total, "disk_total": monitoring.Disk().Total, "gpu_name": "Unknown", "version": CurrentVersion, } endpoint = strings.TrimSuffix(endpoint, "/") + "/api/clients/uploadBasicInfo?token=" + token payload, err := json.Marshal(data) if err != nil { return err } req, err := http.NewRequest("POST", endpoint, strings.NewReader(string(payload))) if err != nil { return err } req.Header.Set("Content-Type", "application/json") client := &http.Client{} resp, err := client.Do(req) if err != nil { return err } defer resp.Body.Close() body, err := io.ReadAll(resp.Body) if err != nil { return err } message := string(body) if resp.StatusCode != http.StatusOK { return fmt.Errorf("status code: %d,%s", resp.StatusCode, message) } return nil } func report(localConfig config.LocalConfig) []byte { message := "" data := map[string]interface{}{} cpu := monitoring.Cpu() data["cpu"] = map[string]interface{}{ "usage": cpu.CPUUsage, } ram := monitoring.Ram() data["ram"] = map[string]interface{}{ "total": ram.Total, "used": ram.Used, } swap := monitoring.Swap() data["swap"] = map[string]interface{}{ "total": swap.Total, "used": swap.Used, } load := monitoring.Load() data["load"] = map[string]interface{}{ "load1": load.Load1, "load5": load.Load5, "load15": load.Load15, } disk := monitoring.Disk() data["disk"] = map[string]interface{}{ "total": disk.Total, "used": disk.Used, } totalUp, totalDown, networkUp, networkDown, err := monitoring.NetworkSpeed(int(localConfig.Interval)) if err != nil { message += fmt.Sprintf("failed to get network speed: %v\n", err) } data["network"] = map[string]interface{}{ "up": networkUp, "down": networkDown, "totalUp": totalUp, "totalDown": totalDown, } tcpCount, udpCount, err := monitoring.ConnectionsCount() if err != nil { message += fmt.Sprintf("failed to get connections: %v\n", err) } data["connections"] = map[string]interface{}{ "tcp": tcpCount, "udp": udpCount, } uptime, err := monitoring.Uptime() if err != nil { message += fmt.Sprintf("failed to get uptime: %v\n", err) } data["uptime"] = uptime processcount := monitoring.ProcessCount() data["process"] = processcount data["message"] = message s, err := json.Marshal(data) if err != nil { log.Println("Failed to marshal data:", err) } return s }