chore: adjust agent code

agent/agent.go

@@ -0,0 +1,495 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"os"
+	"os/exec"
+	"runtime"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/robfig/cron/v3"
+)
+
+// WebSocket 消息类型
+const (
+	WSTypeHeartbeat    = "heartbeat"
+	WSTypeHeartbeatAck = "heartbeat_ack"
+	WSTypeTasks        = "tasks"
+	WSTypeTaskResult   = "task_result"
+	WSTypeUpdate       = "update"
+	WSTypeConnected    = "connected"
+	WSTypeDisabled     = "disabled"
+	WSTypeEnabled      = "enabled"
+	WSTypeFetchTasks   = "fetch_tasks"
+)
+
+type WSMessage struct {
+	Type string          `json:"type"`
+	Data json.RawMessage `json:"data,omitempty"`
+}
+
+type AgentTask struct {
+	ID       uint   `json:"id"`
+	Name     string `json:"name"`
+	Command  string `json:"command"`
+	Schedule string `json:"schedule"`
+	Cron     string `json:"cron"`
+	Timeout  int    `json:"timeout"`
+	WorkDir  string `json:"work_dir"`
+	Envs     string `json:"envs"`
+	Enabled  bool   `json:"enabled"`
+}
+
+type TaskResult struct {
+	TaskID    uint   `json:"task_id"`
+	Command   string `json:"command"`
+	Output    string `json:"output"`
+	Status    string `json:"status"`
+	Duration  int64  `json:"duration"`
+	ExitCode  int    `json:"exit_code"`
+	StartTime int64  `json:"start_time"`
+	EndTime   int64  `json:"end_time"`
+}
+
+type Agent struct {
+	config        *Config
+	configFile    string
+	machineID     string
+	cron          *cron.Cron
+	tasks         map[uint]*AgentTask
+	entryMap      map[uint]cron.EntryID
+	lastTaskCount int
+	mu            sync.RWMutex
+	client        *http.Client
+	wsConn        *websocket.Conn
+	wsMu          sync.Mutex
+	stopCh        chan struct{}
+}
+
+// generateMachineID 生成机器识别码
+func generateMachineID() string {
+	var parts []string
+
+	if hostname, err := os.Hostname(); err == nil {
+		parts = append(parts, hostname)
+	}
+
+	if interfaces, err := net.Interfaces(); err == nil {
+		var macs []string
+		for _, iface := range interfaces {
+			if iface.Flags&net.FlagLoopback != 0 || len(iface.HardwareAddr) == 0 {
+				continue
+			}
+			macs = append(macs, iface.HardwareAddr.String())
+		}
+		sort.Strings(macs)
+		parts = append(parts, macs...)
+	}
+
+	parts = append(parts, runtime.GOOS, runtime.GOARCH)
+
+	data := strings.Join(parts, "|")
+	hash := sha256.Sum256([]byte(data))
+	return hex.EncodeToString(hash[:])
+}
+
+func NewAgent(config *Config, configFile string) *Agent {
+	return &Agent{
+		config:     config,
+		configFile: configFile,
+		machineID:  generateMachineID(),
+		cron:       cron.New(cron.WithSeconds(), cron.WithLocation(cstZone)),
+		tasks:      make(map[uint]*AgentTask),
+		entryMap:   make(map[uint]cron.EntryID),
+		client:     &http.Client{Timeout: 30 * time.Second},
+		stopCh:     make(chan struct{}),
+	}
+}
+
+func (a *Agent) Start() error {
+	if a.config.Token == "" {
+		return fmt.Errorf("缺少令牌，请在配置文件中设置 token")
+	}
+
+	log.Infof("机器识别码: %s", a.machineID[:16]+"...")
+	a.cron.Start()
+
+	go a.wsLoop()
+
+	log.Info("Agent 已启动 (时区: Asia/Shanghai, 模式: WebSocket)")
+	return nil
+}
+
+func (a *Agent) Stop() {
+	close(a.stopCh)
+	a.closeWS()
+	ctx := a.cron.Stop()
+	<-ctx.Done()
+	log.Info("Agent 已停止")
+}
+
+// wsLoop WebSocket 连接循环
+func (a *Agent) wsLoop() {
+	for {
+		select {
+		case <-a.stopCh:
+			return
+		default:
+		}
+
+		if err := a.connectWS(); err != nil {
+			log.Warnf("WebSocket 连接失败: %v，5秒后重试...", err)
+			time.Sleep(5 * time.Second)
+			continue
+		}
+
+		a.readWS()
+
+		log.Warn("WebSocket 连接断开，5秒后重连...")
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func (a *Agent) connectWS() error {
+	serverURL := a.config.ServerURL
+	wsURL := strings.Replace(serverURL, "http://", "ws://", 1)
+	wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
+	wsURL = fmt.Sprintf("%s/api/agent/ws?token=%s&machine_id=%s", wsURL, url.QueryEscape(a.config.Token), url.QueryEscape(a.machineID))
+
+	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
+	conn, _, err := dialer.Dial(wsURL, nil)
+	if err != nil {
+		return err
+	}
+
+	a.wsMu.Lock()
+	a.wsConn = conn
+	a.wsMu.Unlock()
+
+	log.Info("WebSocket 已连接")
+	a.sendHeartbeat()
+	go a.heartbeatLoop()
+
+	return nil
+}
+
+func (a *Agent) closeWS() {
+	a.wsMu.Lock()
+	defer a.wsMu.Unlock()
+	if a.wsConn != nil {
+		a.wsConn.Close()
+		a.wsConn = nil
+	}
+}
+
+func (a *Agent) readWS() {
+	for {
+		a.wsMu.Lock()
+		conn := a.wsConn
+		a.wsMu.Unlock()
+
+		if conn == nil {
+			return
+		}
+
+		_, message, err := conn.ReadMessage()
+		if err != nil {
+			return
+		}
+
+		var msg WSMessage
+		if err := json.Unmarshal(message, &msg); err != nil {
+			continue
+		}
+
+		a.handleWSMessage(&msg)
+	}
+}
+
+func (a *Agent) handleWSMessage(msg *WSMessage) {
+	switch msg.Type {
+	case WSTypeConnected:
+		a.handleConnected(msg.Data)
+	case WSTypeHeartbeatAck:
+		a.handleHeartbeatAck(msg.Data)
+	case WSTypeTasks:
+		a.handleTasks(msg.Data)
+	case WSTypeUpdate:
+		log.Info("收到更新指令，开始更新...")
+		go a.selfUpdate()
+	case WSTypeDisabled:
+		log.Warn("Agent 已被禁用，清空所有任务")
+		a.clearAllTasks()
+	case WSTypeEnabled:
+		log.Info("Agent 已被启用，主动拉取任务")
+		a.fetchTasks()
+	}
+}
+
+func (a *Agent) fetchTasks() {
+	if err := a.sendWSMessage(WSTypeFetchTasks, map[string]interface{}{}); err != nil {
+		log.Warnf("请求任务列表失败: %v", err)
+	}
+}
+
+func (a *Agent) handleConnected(data json.RawMessage) {
+	var resp struct {
+		AgentID    uint   `json:"agent_id"`
+		Name       string `json:"name"`
+		IsNewAgent bool   `json:"is_new_agent"`
+		MachineID  string `json:"machine_id"`
+	}
+	json.Unmarshal(data, &resp)
+
+	if resp.IsNewAgent {
+		log.Infof("注册成功: Agent #%d, 机器码: %s", resp.AgentID, a.machineID[:16]+"...")
+	} else {
+		log.Infof("连接成功: Agent #%d (已存在), 机器码: %s", resp.AgentID, a.machineID[:16]+"...")
+	}
+
+	a.fetchTasks()
+}
+
+func (a *Agent) handleHeartbeatAck(data json.RawMessage) {
+	var resp struct {
+		AgentID       uint   `json:"agent_id"`
+		Name          string `json:"name"`
+		NeedUpdate    bool   `json:"need_update"`
+		ForceUpdate   bool   `json:"force_update"`
+		LatestVersion string `json:"latest_version"`
+	}
+	json.Unmarshal(data, &resp)
+
+	if resp.NeedUpdate && (a.config.AutoUpdate || resp.ForceUpdate) {
+		log.Infof("发现新版本 %s，开始更新...", resp.LatestVersion)
+		go a.selfUpdate()
+	}
+}
+
+func (a *Agent) handleTasks(data json.RawMessage) {
+	var resp struct {
+		Tasks []AgentTask `json:"tasks"`
+	}
+	json.Unmarshal(data, &resp)
+
+	newCount := len(resp.Tasks)
+	if newCount != a.lastTaskCount {
+		log.Infof("任务列表更新: %d -> %d 个任务", a.lastTaskCount, newCount)
+		a.lastTaskCount = newCount
+	}
+
+	a.updateTasks(resp.Tasks)
+}
+
+func (a *Agent) sendWSMessage(msgType string, data interface{}) error {
+	a.wsMu.Lock()
+	defer a.wsMu.Unlock()
+
+	if a.wsConn == nil {
+		return fmt.Errorf("WebSocket 未连接")
+	}
+
+	dataBytes, _ := json.Marshal(data)
+	msg := WSMessage{Type: msgType, Data: dataBytes}
+	msgBytes, _ := json.Marshal(msg)
+
+	a.wsConn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+	return a.wsConn.WriteMessage(websocket.TextMessage, msgBytes)
+}
+
+func (a *Agent) heartbeatLoop() {
+	ticker := time.NewTicker(time.Duration(a.config.Interval) * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-a.stopCh:
+			return
+		case <-ticker.C:
+			a.wsMu.Lock()
+			conn := a.wsConn
+			a.wsMu.Unlock()
+			if conn == nil {
+				return
+			}
+			a.sendHeartbeat()
+		}
+	}
+}
+
+func (a *Agent) sendHeartbeat() {
+	hostname, _ := os.Hostname()
+	data := map[string]interface{}{
+		"version":     Version,
+		"build_time":  BuildTime,
+		"hostname":    hostname,
+		"os":          runtime.GOOS,
+		"arch":        runtime.GOARCH,
+		"auto_update": a.config.AutoUpdate,
+	}
+	if err := a.sendWSMessage(WSTypeHeartbeat, data); err != nil {
+		log.Warnf("发送心跳失败: %v", err)
+	}
+}
+
+func (a *Agent) sendTaskResult(result *TaskResult) {
+	if err := a.sendWSMessage(WSTypeTaskResult, result); err != nil {
+		log.Warnf("发送任务结果失败: %v，尝试 HTTP 上报", err)
+		a.reportResultHTTP(result)
+	}
+}
+
+func (a *Agent) reportResultHTTP(result *TaskResult) error {
+	resp, err := a.doRequest("POST", "/api/agent/report", result)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	return nil
+}
+
+func (a *Agent) updateTasks(tasks []AgentTask) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	newTasks := make(map[uint]*AgentTask)
+	for i := range tasks {
+		newTasks[tasks[i].ID] = &tasks[i]
+	}
+
+	for id, entryID := range a.entryMap {
+		if _, exists := newTasks[id]; !exists {
+			a.cron.Remove(entryID)
+			delete(a.entryMap, id)
+			delete(a.tasks, id)
+			log.Infof("移除任务 #%d", id)
+		}
+	}
+
+	for id, task := range newTasks {
+		oldTask, exists := a.tasks[id]
+		if !exists || oldTask.Schedule != task.Schedule || oldTask.Command != task.Command {
+			if entryID, ok := a.entryMap[id]; ok {
+				a.cron.Remove(entryID)
+			}
+
+			taskCopy := *task
+			entryID, err := a.cron.AddFunc(task.Schedule, func() {
+				a.executeTask(&taskCopy)
+			})
+			if err != nil {
+				log.Errorf("添加任务 #%d 失败: %v", id, err)
+				continue
+			}
+
+			a.entryMap[id] = entryID
+			a.tasks[id] = task
+			log.Infof("调度任务 #%d %s (%s)", id, task.Name, task.Schedule)
+		}
+	}
+}
+
+func (a *Agent) clearAllTasks() {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	for id, entryID := range a.entryMap {
+		a.cron.Remove(entryID)
+		log.Infof("移除任务 #%d", id)
+	}
+
+	a.entryMap = make(map[uint]cron.EntryID)
+	a.tasks = make(map[uint]*AgentTask)
+	a.lastTaskCount = 0
+	log.Info("所有任务已清空")
+}
+
+func (a *Agent) executeTask(task *AgentTask) {
+	log.Infof("执行任务 #%d %s", task.ID, task.Name)
+
+	start := time.Now()
+	result := &TaskResult{
+		TaskID:    task.ID,
+		Command:   task.Command,
+		StartTime: start.Unix(),
+	}
+
+	timeout := task.Timeout
+	if timeout <= 0 {
+		timeout = 30
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Minute)
+	defer cancel()
+
+	var cmd *exec.Cmd
+	if runtime.GOOS == "windows" {
+		cmd = exec.CommandContext(ctx, "cmd", "/c", task.Command)
+	} else {
+		cmd = exec.CommandContext(ctx, "sh", "-c", task.Command)
+	}
+
+	if task.WorkDir != "" {
+		cmd.Dir = task.WorkDir
+	}
+
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	err := cmd.Run()
+	end := time.Now()
+
+	result.EndTime = end.Unix()
+	result.Duration = end.Sub(start).Milliseconds()
+	result.Output = stdout.String()
+
+	if err != nil {
+		result.Status = "failed"
+		result.Output += "\n[ERROR]\n" + stderr.String() + "\n" + err.Error()
+		if exitErr, ok := err.(*exec.ExitError); ok {
+			result.ExitCode = exitErr.ExitCode()
+		} else {
+			result.ExitCode = 1
+		}
+	} else {
+		result.Status = "success"
+		result.ExitCode = 0
+	}
+
+	a.sendTaskResult(result)
+	log.Infof("任务 #%d 执行完成 (%s)", result.TaskID, result.Status)
+}
+
+func (a *Agent) doRequest(method, path string, body interface{}) (*http.Response, error) {
+	var bodyReader io.Reader
+	if body != nil {
+		data, err := json.Marshal(body)
+		if err != nil {
+			return nil, err
+		}
+		bodyReader = bytes.NewReader(data)
+	}
+
+	req, err := http.NewRequest(method, a.config.ServerURL+path, bodyReader)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Authorization", "Bearer "+a.config.Token)
+	req.Header.Set("Content-Type", "application/json")
+
+	return a.client.Do(req)
+}

agent/config.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"os"
+	"path/filepath"
+	"strconv"
+
+	"gopkg.in/ini.v1"
+)
+
+// Config Agent 配置
+type Config struct {
+	ServerURL  string
+	Name       string
+	Token      string
+	Interval   int
+	AutoUpdate bool
+}
+
+func loadConfigFile(path string, config *Config) error {
+	cfg, err := ini.Load(path)
+	if err != nil {
+		return err
+	}
+
+	section := cfg.Section("agent")
+	if v := section.Key("server_url").String(); v != "" {
+		config.ServerURL = v
+	}
+	if v := section.Key("name").String(); v != "" {
+		config.Name = v
+	}
+	if v := section.Key("token").String(); v != "" {
+		config.Token = v
+	}
+	if v := section.Key("interval").String(); v != "" {
+		if i, err := strconv.Atoi(v); err == nil && i > 0 {
+			config.Interval = i
+		}
+	}
+	if v := section.Key("auto_update").String(); v != "" {
+		config.AutoUpdate = v == "true" || v == "1"
+	}
+	return nil
+}
+
+func saveConfigFile(path string, config *Config) error {
+	dir := filepath.Dir(path)
+	if dir != "" && dir != "." {
+		os.MkdirAll(dir, 0755)
+	}
+
+	cfg := ini.Empty()
+	section := cfg.Section("agent")
+	section.Key("server_url").SetValue(config.ServerURL)
+	section.Key("name").SetValue(config.Name)
+	section.Key("token").SetValue(config.Token)
+	section.Key("interval").SetValue(strconv.Itoa(config.Interval))
+	if config.AutoUpdate {
+		section.Key("auto_update").SetValue("true")
+	} else {
+		section.Key("auto_update").SetValue("false")
+	}
+
+	return cfg.SaveTo(path)
+}

agent/logger.go

@@ -0,0 +1,69 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/sirupsen/logrus"
+	"gopkg.in/natefinch/lumberjack.v2"
+)
+
+// 日志实例
+var log = logrus.New()
+
+// CustomFormatter 自定义日志格式
+type CustomFormatter struct{}
+
+// ANSI 颜色代码
+const (
+	colorReset  = "\033[0m"
+	colorRed    = "\033[31m"
+	colorYellow = "\033[33m"
+	colorBlue   = "\033[36m"
+	colorGray   = "\033[37m"
+)
+
+func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
+	timestamp := entry.Time.Format("2006-01-02 15:04:05")
+	level := strings.ToUpper(entry.Level.String())
+
+	var levelColor string
+	switch entry.Level {
+	case logrus.DebugLevel, logrus.TraceLevel:
+		levelColor = colorGray
+	case logrus.InfoLevel:
+		levelColor = colorBlue
+	case logrus.WarnLevel:
+		levelColor = colorYellow
+	case logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel:
+		levelColor = colorRed
+	default:
+		levelColor = colorBlue
+	}
+
+	msg := fmt.Sprintf("[%s]%s[%s]%s %s\n", timestamp, levelColor, level, colorReset, entry.Message)
+	return []byte(msg), nil
+}
+
+func initLogger(logFile string) {
+	logDir := filepath.Dir(logFile)
+	if logDir != "" && logDir != "." {
+		os.MkdirAll(logDir, 0755)
+	}
+
+	log.SetFormatter(&CustomFormatter{})
+	log.SetLevel(logrus.InfoLevel)
+
+	lumberjackLogger := &lumberjack.Logger{
+		Filename:   logFile,
+		MaxSize:    5,
+		MaxBackups: 3,
+		MaxAge:     0,
+		Compress:   false,
+	}
+
+	log.SetOutput(io.MultiWriter(os.Stdout, lumberjackLogger))
+}

agent/service.go

@@ -0,0 +1,207 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strconv"
+	"syscall"
+)
+
+// ========== PID 文件管理 ==========
+
+func getPidFile() string {
+	return filepath.Join(dataDir, "agent.pid")
+}
+
+func writePidFile() {
+	os.MkdirAll(dataDir, 0755)
+	pidFile := getPidFile()
+	os.WriteFile(pidFile, []byte(strconv.Itoa(os.Getpid())), 0644)
+}
+
+func readPidFile() int {
+	pidFile := getPidFile()
+	data, err := os.ReadFile(pidFile)
+	if err != nil {
+		return 0
+	}
+	pid, _ := strconv.Atoi(string(data))
+	return pid
+}
+
+func removePidFile() {
+	os.Remove(getPidFile())
+}
+
+// ========== 命令实现 ==========
+
+func cmdStop() {
+	pid := readPidFile()
+	if pid == 0 {
+		fmt.Println("Agent 未运行")
+		return
+	}
+
+	process, err := os.FindProcess(pid)
+	if err != nil {
+		fmt.Printf("找不到进程 %d\n", pid)
+		removePidFile()
+		return
+	}
+
+	if runtime.GOOS == "windows" {
+		err = process.Kill()
+	} else {
+		err = process.Signal(syscall.SIGTERM)
+	}
+
+	if err != nil {
+		fmt.Printf("停止失败: %v\n", err)
+		return
+	}
+
+	fmt.Println("Agent 已停止")
+	removePidFile()
+}
+
+func cmdStatus() {
+	pid := readPidFile()
+	if pid == 0 {
+		fmt.Println("状态: 未运行")
+		return
+	}
+
+	// 检查进程是否存在
+	process, err := os.FindProcess(pid)
+	if err != nil {
+		fmt.Println("状态: 未运行")
+		removePidFile()
+		return
+	}
+
+	// Unix 系统发送信号 0 检查进程
+	if runtime.GOOS != "windows" {
+		err = process.Signal(syscall.Signal(0))
+		if err != nil {
+			fmt.Println("状态: 未运行")
+			removePidFile()
+			return
+		}
+	}
+
+	fmt.Printf("状态: 运行中 (PID: %d)\n", pid)
+}
+
+func cmdInstall() {
+	exePath, _ := os.Executable()
+	exeDir := filepath.Dir(exePath)
+
+	if runtime.GOOS == "windows" {
+		installWindows(exePath, exeDir)
+	} else {
+		installLinux(exePath, exeDir)
+	}
+}
+
+func cmdUninstall() {
+	if runtime.GOOS == "windows" {
+		uninstallWindows()
+	} else {
+		uninstallLinux()
+	}
+}
+
+// ========== Linux systemd ==========
+
+func installLinux(exePath, exeDir string) {
+	serviceContent := fmt.Sprintf(`[Unit]
+Description=%s
+After=network.target
+
+[Service]
+Type=simple
+WorkingDirectory=%s
+ExecStart=%s start
+Restart=always
+RestartSec=5
+
+[Install]
+WantedBy=multi-user.target
+`, ServiceDesc, exeDir, exePath)
+
+	servicePath := fmt.Sprintf("/etc/systemd/system/%s.service", ServiceName)
+	if err := os.WriteFile(servicePath, []byte(serviceContent), 0644); err != nil {
+		fmt.Printf("创建服务文件失败: %v\n", err)
+		fmt.Println("请使用 sudo 运行")
+		return
+	}
+
+	// 重载 systemd
+	exec.Command("systemctl", "daemon-reload").Run()
+	exec.Command("systemctl", "enable", ServiceName).Run()
+
+	fmt.Printf("服务已安装: %s\n", servicePath)
+	fmt.Println("使用以下命令管理服务:")
+	fmt.Printf("  启动: sudo systemctl start %s\n", ServiceName)
+	fmt.Printf("  停止: sudo systemctl stop %s\n", ServiceName)
+	fmt.Printf("  状态: sudo systemctl status %s\n", ServiceName)
+}
+
+func uninstallLinux() {
+	// 停止服务
+	exec.Command("systemctl", "stop", ServiceName).Run()
+	exec.Command("systemctl", "disable", ServiceName).Run()
+
+	servicePath := fmt.Sprintf("/etc/systemd/system/%s.service", ServiceName)
+	if err := os.Remove(servicePath); err != nil {
+		fmt.Printf("删除服务文件失败: %v\n", err)
+		fmt.Println("请使用 sudo 运行")
+		return
+	}
+
+	exec.Command("systemctl", "daemon-reload").Run()
+	fmt.Println("服务已卸载")
+}
+
+// ========== Windows 服务 ==========
+
+func installWindows(exePath, exeDir string) {
+	// 使用 sc.exe 创建服务
+	cmd := exec.Command("sc", "create", ServiceName,
+		"binPath=", fmt.Sprintf(`"%s" start`, exePath),
+		"start=", "auto",
+		"DisplayName=", ServiceDesc)
+
+	if err := cmd.Run(); err != nil {
+		fmt.Printf("创建服务失败: %v\n", err)
+		fmt.Println("请以管理员身份运行")
+		return
+	}
+
+	// 设置服务描述
+	exec.Command("sc", "description", ServiceName, ServiceDesc).Run()
+
+	fmt.Println("服务已安装")
+	fmt.Println("使用以下命令管理服务:")
+	fmt.Printf("  启动: sc start %s\n", ServiceName)
+	fmt.Printf("  停止: sc stop %s\n", ServiceName)
+	fmt.Printf("  状态: sc query %s\n", ServiceName)
+}
+
+func uninstallWindows() {
+	// 停止服务
+	exec.Command("sc", "stop", ServiceName).Run()
+
+	// 删除服务
+	cmd := exec.Command("sc", "delete", ServiceName)
+	if err := cmd.Run(); err != nil {
+		fmt.Printf("删除服务失败: %v\n", err)
+		fmt.Println("请以管理员身份运行")
+		return
+	}
+
+	fmt.Println("服务已卸载")
+}

agent/update.go

@@ -0,0 +1,156 @@
+package main
+
+import (
+	"archive/tar"
+	"compress/gzip"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"syscall"
+	"time"
+)
+
+// selfUpdate 自动更新
+func (a *Agent) selfUpdate() {
+	// 获取当前可执行文件路径
+	exePath, err := os.Executable()
+	if err != nil {
+		log.Errorf("获取可执行文件路径失败: %v", err)
+		return
+	}
+	exePath, _ = filepath.Abs(exePath)
+
+	// 下载新版本 tar.gz
+	downloadURL := a.config.ServerURL + "/api/agent/download?os=" + runtime.GOOS + "&arch=" + runtime.GOARCH
+	req, err := http.NewRequest("GET", downloadURL, nil)
+	if err != nil {
+		log.Errorf("创建下载请求失败: %v", err)
+		return
+	}
+	req.Header.Set("Authorization", "Bearer "+a.config.Token)
+
+	client := &http.Client{Timeout: 5 * time.Minute}
+	resp, err := client.Do(req)
+	if err != nil {
+		log.Errorf("下载新版本失败: %v", err)
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		log.Errorf("下载新版本失败: HTTP %d", resp.StatusCode)
+		return
+	}
+
+	// 读取 tar.gz 内容
+	gzReader, err := gzip.NewReader(resp.Body)
+	if err != nil {
+		log.Errorf("解压 gzip 失败: %v", err)
+		return
+	}
+	defer gzReader.Close()
+
+	tarReader := tar.NewReader(gzReader)
+
+	// 解压并找到二进制文件
+	var newBinary []byte
+	binaryName := "baihu-agent"
+	if runtime.GOOS == "windows" {
+		binaryName = "baihu-agent.exe"
+	}
+
+	for {
+		header, err := tarReader.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			log.Errorf("读取 tar 失败: %v", err)
+			return
+		}
+
+		if header.Typeflag == tar.TypeReg && header.Name == binaryName {
+			newBinary, err = io.ReadAll(tarReader)
+			if err != nil {
+				log.Errorf("读取二进制文件失败: %v", err)
+				return
+			}
+			break
+		}
+	}
+
+	if newBinary == nil {
+		log.Errorf("tar.gz 中未找到 %s", binaryName)
+		return
+	}
+
+	// 保存到临时文件（放到 data 目录）
+	os.MkdirAll(dataDir, 0755)
+	tmpFile := filepath.Join(dataDir, binaryName+".new")
+	if err := os.WriteFile(tmpFile, newBinary, 0755); err != nil {
+		log.Errorf("保存新版本失败: %v", err)
+		return
+	}
+
+	// 计算基础路径（去掉所有 .bak 后缀）
+	basePath := exePath
+	for strings.HasSuffix(basePath, ".bak") {
+		basePath = strings.TrimSuffix(basePath, ".bak")
+	}
+	backupFile := basePath + ".bak"
+
+	// 如果当前运行的就是 .bak 文件，直接删除它（更新后会用新版本）
+	// 否则需要备份当前文件
+	if exePath != backupFile {
+		os.Remove(backupFile)
+		if err := os.Rename(exePath, backupFile); err != nil {
+			log.Errorf("备份旧版本失败: %v", err)
+			os.Remove(tmpFile)
+			return
+		}
+	}
+
+	// 替换为新版本（放到 basePath，即不带 .bak 的路径）
+	if err := os.Rename(tmpFile, basePath); err != nil {
+		log.Errorf("替换新版本失败: %v", err)
+		if exePath != backupFile {
+			os.Rename(backupFile, exePath) // 恢复旧版本
+		}
+		return
+	}
+
+	// 如果之前运行的是 .bak 文件，现在可以删除它了
+	if exePath == backupFile {
+		os.Remove(exePath)
+	}
+
+	log.Info("更新完成，正在重启...")
+
+	// 重启服务
+	a.restart()
+}
+
+// restart 重启服务
+func (a *Agent) restart() {
+	exePath, _ := os.Executable()
+
+	// 计算基础路径（去掉所有 .bak 后缀），确保启动的是正确的可执行文件
+	basePath := exePath
+	for strings.HasSuffix(basePath, ".bak") {
+		basePath = strings.TrimSuffix(basePath, ".bak")
+	}
+
+	if runtime.GOOS == "windows" {
+		// Windows: 启动新进程后退出
+		cmd := exec.Command(basePath, "start")
+		cmd.Start()
+		os.Exit(0)
+	} else {
+		// Linux/macOS: 使用 exec 替换当前进程
+		syscall.Exec(basePath, []string{basePath, "start"}, os.Environ())
+	}
+}