From 78c3707cce8872e16f6fdb79e7aad6bb928583c7 Mon Sep 17 00:00:00 2001 From: engigu Date: Sun, 28 Dec 2025 11:42:58 +0800 Subject: [PATCH] feat: add agent start code --- .gitignore | 3 + agent/config.example.ini | 8 +- agent/go.mod | 1 + agent/go.sum | 2 + agent/main.go | 499 ++++++++++++++--------- internal/controllers/agent_controller.go | 374 +++++++++++++---- internal/controllers/task_controller.go | 65 ++- internal/database/migrate.go | 20 + internal/models/agent.go | 32 +- internal/router/router.go | 10 +- internal/services/agent_service.go | 235 +++++++---- internal/services/agent_ws_service.go | 293 +++++++++++++ internal/services/task_service.go | 6 +- web/src/api/index.ts | 25 +- web/src/views/agents/Agents.vue | 263 ++++++------ web/src/views/history/History.vue | 15 +- web/src/views/tasks/RepoDialog.vue | 40 +- web/src/views/tasks/TaskDialog.vue | 47 ++- web/src/views/tasks/Tasks.vue | 54 ++- 19 files changed, 1491 insertions(+), 501 deletions(-) create mode 100644 internal/services/agent_ws_service.go diff --git a/.gitignore b/.gitignore index 5efbfe4..ca83406 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,9 @@ envs/ # logs/ # scripts/ configs/config.ini +agent/config.ini +agent/agent.pid +agent/baihu-agent web/dist/* !web/dist/.gitkeep diff --git a/agent/config.example.ini b/agent/config.example.ini index ce2a1cf..98ebf18 100644 --- a/agent/config.example.ini +++ b/agent/config.example.ini @@ -1,11 +1,11 @@ [agent] -# 主服务器地址 +# 主服务器地址（http/https，Agent 会自动转换为 WebSocket 连接） server_url = http://192.168.1.100:8052 # Agent 名称（留空则使用主机名） -name = agent-01 -# Token（由服务器下发，首次运行留空） +name = +# 注册令牌（首次注册时填写，注册成功后会自动替换为认证 Token） token = -# 心跳间隔（秒） +# 心跳间隔（秒），默认 30 interval = 30 # 自动更新（true/false） auto_update = true diff --git a/agent/go.mod b/agent/go.mod index 20bf62d..4146af2 100644 --- a/agent/go.mod +++ b/agent/go.mod @@ -3,6 +3,7 @@ module baihu-agent go 1.24 require ( + github.com/gorilla/websocket v1.5.3 github.com/robfig/cron/v3 v3.0.1 github.com/sirupsen/logrus v1.9.3 gopkg.in/ini.v1 v1.67.0 diff --git a/agent/go.sum b/agent/go.sum index 737c46b..ce54287 100644 --- a/agent/go.sum +++ b/agent/go.sum @@ -1,6 +1,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= diff --git a/agent/main.go b/agent/main.go index f313ced..e4dee43 100644 --- a/agent/main.go +++ b/agent/main.go @@ -5,21 +5,27 @@ import ( "bytes" "compress/gzip" "context" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "io" + "net" "net/http" + "net/url" "os" "os/exec" "os/signal" "path/filepath" "runtime" + "sort" "strconv" "strings" "sync" "syscall" "time" + "github.com/gorilla/websocket" "github.com/robfig/cron/v3" "github.com/sirupsen/logrus" "gopkg.in/ini.v1" @@ -158,7 +164,7 @@ func cmdStart() { config.Name = hostname } - log.Infof("Baihu Agent v%s", Version) + log.Infof("Baihu Agent Version: %s", Version) if BuildTime != "" { log.Infof("构建时间: %s", BuildTime) } @@ -320,7 +326,7 @@ func installWindows(exePath, exeDir string) { "binPath=", fmt.Sprintf(`"%s" start`, exePath), "start=", "auto", "DisplayName=", ServiceDesc) - + if err := cmd.Run(); err != nil { fmt.Printf("创建服务失败: %v\n", err) fmt.Println("请以管理员身份运行") @@ -340,7 +346,7 @@ func installWindows(exePath, exeDir string) { func uninstallWindows() { // 停止服务 exec.Command("sc", "stop", ServiceName).Run() - + // 删除服务 cmd := exec.Command("sc", "delete", ServiceName) if err := cmd.Run(); err != nil { @@ -433,7 +439,6 @@ func initLogger(logFile string) { log.SetOutput(io.MultiWriter(os.Stdout, lumberjackLogger)) } - // ========== 配置相关 ========== type Config struct { @@ -494,6 +499,24 @@ func saveConfigFile(path string, config *Config) error { // ========== Agent 结构 ========== +// WebSocket 消息类型 +const ( + WSTypeHeartbeat = "heartbeat" + WSTypeHeartbeatAck = "heartbeat_ack" + WSTypeTasks = "tasks" + WSTypeTaskResult = "task_result" + WSTypeUpdate = "update" + WSTypeConnected = "connected" + WSTypeDisabled = "disabled" // Agent 被禁用 + WSTypeEnabled = "enabled" // Agent 被启用 + WSTypeFetchTasks = "fetch_tasks" // Agent 请求任务列表 +) + +type WSMessage struct { + Type string `json:"type"` + Data json.RawMessage `json:"data,omitempty"` +} + type AgentTask struct { ID uint `json:"id"` Name string `json:"name"` @@ -517,147 +540,309 @@ type TaskResult struct { } type Agent struct { - config *Config - configFile string - cron *cron.Cron - tasks map[uint]*AgentTask - entryMap map[uint]cron.EntryID - mu sync.RWMutex - client *http.Client + 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 生成机器识别码（基于 hostname + MAC 地址） +func generateMachineID() string { + var parts []string + + // 1. Hostname + if hostname, err := os.Hostname(); err == nil { + parts = append(parts, hostname) + } + + // 2. MAC 地址（取所有非回环网卡的 MAC） + if interfaces, err := net.Interfaces(); err == nil { + var macs []string + for _, iface := range interfaces { + // 跳过回环和无 MAC 的接口 + if iface.Flags&net.FlagLoopback != 0 || len(iface.HardwareAddr) == 0 { + continue + } + macs = append(macs, iface.HardwareAddr.String()) + } + // 排序确保顺序一致 + sort.Strings(macs) + parts = append(parts, macs...) + } + + // 3. 操作系统和架构 + parts = append(parts, runtime.GOOS, runtime.GOARCH) + + // 生成 SHA256 哈希 + 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 == "" { - log.Info("未找到 Token，开始注册流程...") - if err := a.registerAndWait(); err != nil { - return err - } - } - - if err := a.heartbeat(); err != nil { - log.Warnf("首次心跳失败: %v（将继续重试）", err) - } - - if err := a.syncTasks(); err != nil { - log.Warnf("同步任务失败: %v（将继续重试）", err) + return fmt.Errorf("缺少令牌，请在配置文件中设置 token") } + log.Infof("机器识别码: %s", a.machineID[:16]+"...") a.cron.Start() - go a.heartbeatLoop() - go a.syncTasksLoop() - log.Info("Agent 已启动 (时区: Asia/Shanghai)") + // 启动 WebSocket 连接 + 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 已停止") } -func (a *Agent) registerAndWait() error { - hostname, _ := os.Hostname() - - body := map[string]string{ - "name": a.config.Name, - "hostname": hostname, - "version": Version, - } - - resp, err := a.doRequestNoAuth("POST", "/api/agent/register", body) - if err != nil { - return fmt.Errorf("注册失败: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - data, _ := io.ReadAll(resp.Body) - return fmt.Errorf("注册失败: %s", string(data)) - } - - var result struct { - Code int `json:"code"` - Data struct { - AgentID uint `json:"agent_id"` - Status string `json:"status"` - Message string `json:"message"` - } `json:"data"` - } - json.NewDecoder(resp.Body).Decode(&result) - - log.Infof("注册成功 (ID: %d)，等待管理员审核...", result.Data.AgentID) - - ticker := time.NewTicker(5 * time.Second) - defer ticker.Stop() - +// wsLoop WebSocket 连接循环（自动重连） +func (a *Agent) wsLoop() { for { - <-ticker.C + select { + case <-a.stopCh: + return + default: + } - statusResp, err := a.doRequestNoAuth("POST", "/api/agent/status", map[string]string{ - "name": a.config.Name, - }) - if err != nil { - log.Warnf("检查状态失败: %v", err) + if err := a.connectWS(); err != nil { + log.Warnf("WebSocket 连接失败: %v，5秒后重试...", err) + time.Sleep(5 * time.Second) continue } - if statusResp.StatusCode != http.StatusOK { - statusResp.Body.Close() - continue - } + // 连接成功，开始读取消息 + a.readWS() - var statusResult struct { - Code int `json:"code"` - Data struct { - AgentID uint `json:"agent_id"` - Status string `json:"status"` - Token string `json:"token"` - } `json:"data"` - } - json.NewDecoder(statusResp.Body).Decode(&statusResult) - statusResp.Body.Close() - - if statusResult.Data.Status != "pending" && statusResult.Data.Token != "" { - a.config.Token = statusResult.Data.Token - if err := saveConfigFile(a.configFile, a.config); err != nil { - log.Warnf("保存配置文件失败: %v", err) - } else { - log.Infof("Token 已保存到 %s", a.configFile) - } - log.Info("审核通过，开始工作...") - return nil - } - - log.Debug("等待审核中...") + // 连接断开，等待后重连 + log.Warn("WebSocket 连接断开，5秒后重连...") + time.Sleep(5 * time.Second) } } +// connectWS 建立 WebSocket 连接 +func (a *Agent) connectWS() error { + // 构建 WebSocket URL + 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 +} + +// closeWS 关闭 WebSocket 连接 +func (a *Agent) closeWS() { + a.wsMu.Lock() + defer a.wsMu.Unlock() + if a.wsConn != nil { + a.wsConn.Close() + a.wsConn = nil + } +} + +// readWS 读取 WebSocket 消息 +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) + } +} + +// handleWSMessage 处理 WebSocket 消息 +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() + } +} + +// fetchTasks 主动请求任务列表 +func (a *Agent) fetchTasks() { + if err := a.sendWSMessage(WSTypeFetchTasks, map[string]interface{}{}); err != nil { + log.Warnf("请求任务列表失败: %v", err) + } +} + +// handleConnected 处理连接成功消息 +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() +} + +// handleHeartbeatAck 处理心跳响应 +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() + } +} + +// handleTasks 处理任务列表 +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) +} + +// sendWSMessage 发送 WebSocket 消息 +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) +} + +// heartbeatLoop 心跳循环 func (a *Agent) heartbeatLoop() { ticker := time.NewTicker(time.Duration(a.config.Interval) * time.Second) defer ticker.Stop() - for range ticker.C { - if err := a.heartbeat(); err != nil { - log.Warnf("心跳失败: %v", err) + 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) heartbeat() error { +// sendHeartbeat 发送心跳 +func (a *Agent) sendHeartbeat() { hostname, _ := os.Hostname() - body := map[string]interface{}{ + data := map[string]interface{}{ "version": Version, "build_time": BuildTime, "hostname": hostname, @@ -665,76 +850,27 @@ func (a *Agent) heartbeat() error { "arch": runtime.GOARCH, "auto_update": a.config.AutoUpdate, } + if err := a.sendWSMessage(WSTypeHeartbeat, data); err != nil { + log.Warnf("发送心跳失败: %v", err) + } +} - resp, err := a.doRequest("POST", "/api/agent/heartbeat", body) +// sendTaskResult 发送任务结果 +func (a *Agent) sendTaskResult(result *TaskResult) { + if err := a.sendWSMessage(WSTypeTaskResult, result); err != nil { + log.Warnf("发送任务结果失败: %v，尝试 HTTP 上报", err) + // 降级到 HTTP + a.reportResultHTTP(result) + } +} + +// reportResultHTTP HTTP 方式上报结果（降级方案） +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() - - if resp.StatusCode != http.StatusOK { - data, _ := io.ReadAll(resp.Body) - return fmt.Errorf("心跳失败: %s", string(data)) - } - - var result struct { - Code int `json:"code"` - Data 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:"data"` - } - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return nil // 忽略解析错误 - } - - // 检查是否需要更新 - if result.Data.NeedUpdate && (a.config.AutoUpdate || result.Data.ForceUpdate) { - log.Infof("发现新版本 %s，开始更新...", result.Data.LatestVersion) - go a.selfUpdate() - } - - return nil -} - -func (a *Agent) syncTasksLoop() { - ticker := time.NewTicker(time.Duration(a.config.Interval) * time.Second) - defer ticker.Stop() - - for range ticker.C { - if err := a.syncTasks(); err != nil { - log.Warnf("同步任务失败: %v", err) - } - } -} - -func (a *Agent) syncTasks() error { - resp, err := a.doRequest("GET", "/api/agent/tasks", nil) - if err != nil { - return err - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - data, _ := io.ReadAll(resp.Body) - return fmt.Errorf("获取任务失败: %s", string(data)) - } - - var result struct { - Code int `json:"code"` - Data struct { - AgentID uint `json:"agent_id"` - Tasks []AgentTask `json:"tasks"` - } `json:"data"` - } - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return err - } - - a.updateTasks(result.Data.Tasks) return nil } @@ -779,6 +915,21 @@ func (a *Agent) updateTasks(tasks []AgentTask) { } } +// clearAllTasks 清空所有任务（Agent 被禁用时调用） +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) + log.Info("所有任务已清空") +} + func (a *Agent) executeTask(task *AgentTask) { log.Infof("执行任务 #%d %s", task.ID, task.Name) @@ -831,25 +982,9 @@ func (a *Agent) executeTask(task *AgentTask) { result.ExitCode = 0 } - if err := a.reportResult(result); err != nil { - log.Errorf("上报结果失败: %v", err) - } -} - -func (a *Agent) reportResult(result *TaskResult) error { - resp, err := a.doRequest("POST", "/api/agent/report", result) - if err != nil { - return err - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - data, _ := io.ReadAll(resp.Body) - return fmt.Errorf("上报失败: %s", string(data)) - } - + // 使用 WebSocket 上报结果 + a.sendTaskResult(result) log.Infof("任务 #%d 执行完成 (%s)", result.TaskID, result.Status) - return nil } func (a *Agent) doRequest(method, path string, body interface{}) (*http.Response, error) { diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go index a304fb6..17e9643 100644 --- a/internal/controllers/agent_controller.go +++ b/internal/controllers/agent_controller.go @@ -1,72 +1,46 @@ package controllers import ( + "baihu/internal/logger" "baihu/internal/models" "baihu/internal/services" "baihu/internal/utils" + "encoding/json" + "net/http" "strconv" "strings" + "time" "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" ) +var agentUpgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + return true + }, +} + // AgentController Agent 控制器 type AgentController struct { agentService *services.AgentService + wsManager *services.AgentWSManager } // NewAgentController 创建 Agent 控制器 func NewAgentController() *AgentController { return &AgentController{ agentService: services.NewAgentService(), + wsManager: services.GetAgentWSManager(), } } -// List 获取已审核的 Agent 列表 +// List 获取 Agent 列表 func (c *AgentController) List(ctx *gin.Context) { agents := c.agentService.List() utils.Success(ctx, agents) } -// ListPending 获取待审核的 Agent 列表 -func (c *AgentController) ListPending(ctx *gin.Context) { - agents := c.agentService.ListPending() - utils.Success(ctx, agents) -} - -// Approve 审核通过 Agent -func (c *AgentController) Approve(ctx *gin.Context) { - id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) - if err != nil { - utils.BadRequest(ctx, "无效的 ID") - return - } - - agent, err := c.agentService.Approve(uint(id)) - if err != nil { - utils.BadRequest(ctx, err.Error()) - return - } - - utils.Success(ctx, agent) -} - -// Reject 拒绝 Agent -func (c *AgentController) Reject(ctx *gin.Context) { - id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) - if err != nil { - utils.BadRequest(ctx, "无效的 ID") - return - } - - if err := c.agentService.Reject(uint(id)); err != nil { - utils.ServerError(ctx, err.Error()) - return - } - - utils.SuccessMsg(ctx, "已拒绝") -} - // Update 更新 Agent func (c *AgentController) Update(ctx *gin.Context) { id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) @@ -86,11 +60,36 @@ func (c *AgentController) Update(ctx *gin.Context) { return } + // 获取旧状态 + oldAgent := c.agentService.GetByID(uint(id)) + if oldAgent == nil { + utils.NotFound(ctx, "Agent 不存在") + return + } + wasEnabled := oldAgent.Enabled + if err := c.agentService.Update(uint(id), req.Name, req.Description, req.Enabled); err != nil { utils.ServerError(ctx, err.Error()) return } + // 如果启用状态发生变化，通知 Agent + if wasEnabled != req.Enabled { + if req.Enabled { + // 启用：发送任务列表 + c.wsManager.SendToAgent(uint(id), services.WSTypeEnabled, map[string]interface{}{ + "message": "Agent 已启用", + }) + // 发送任务列表 + c.wsManager.BroadcastTasks(uint(id)) + } else { + // 禁用：发送禁用消息，Agent 收到后清空任务 + c.wsManager.SendToAgent(uint(id), services.WSTypeDisabled, map[string]interface{}{ + "message": "Agent 已禁用", + }) + } + } + utils.SuccessMsg(ctx, "更新成功") } @@ -143,49 +142,19 @@ func (c *AgentController) Register(ctx *gin.Context) { } ip := ctx.ClientIP() - agent, err := c.agentService.Register(&req, ip) + agent, token, err := c.agentService.Register(&req, ip) if err != nil { - utils.ServerError(ctx, err.Error()) + utils.BadRequest(ctx, err.Error()) return } utils.Success(ctx, gin.H{ "agent_id": agent.ID, - "status": agent.Status, - "message": "注册成功，等待审核", + "token": token, + "message": "注册成功", }) } -// CheckStatus Agent 检查状态（用于轮询等待审核结果） -func (c *AgentController) CheckStatus(ctx *gin.Context) { - var req struct { - Name string `json:"name"` - } - if err := ctx.ShouldBindJSON(&req); err != nil { - utils.BadRequest(ctx, "参数错误") - return - } - - ip := ctx.ClientIP() - agent, err := c.agentService.CheckPendingAgent(req.Name, ip) - if err != nil { - utils.NotFound(ctx, err.Error()) - return - } - - response := gin.H{ - "agent_id": agent.ID, - "status": agent.Status, - } - - // 如果已审核通过，返回 Token - if agent.Status != "pending" && agent.Token != "" { - response["token"] = agent.Token - } - - utils.Success(ctx, response) -} - // Heartbeat Agent 心跳 func (c *AgentController) Heartbeat(ctx *gin.Context) { token := c.getAgentToken(ctx) @@ -348,3 +317,260 @@ func (c *AgentController) ForceUpdate(ctx *gin.Context) { utils.SuccessMsg(ctx, "已标记强制更新，Agent 下次心跳时将自动更新") } + + +// ========== WebSocket ========== + +// WSConnect Agent WebSocket 连接 +func (c *AgentController) WSConnect(ctx *gin.Context) { + ip := ctx.ClientIP() + + // 检查 IP 限流 + if allowed, reason := c.wsManager.CheckRateLimit(ip); !allowed { + logger.Warnf("[AgentWS] IP %s 被限流: %s", ip, reason) + ctx.JSON(http.StatusTooManyRequests, gin.H{"error": reason}) + return + } + + token := ctx.Query("token") + if token == "" { + c.wsManager.RecordConnectFail(ip) + ctx.JSON(http.StatusUnauthorized, gin.H{"error": "缺少 token"}) + return + } + + machineID := ctx.Query("machine_id") + isNewAgent := false + + // 先尝试用 token 查找已有 Agent + agent := c.agentService.GetByToken(token) + + // 如果没找到，尝试用令牌注册（会检查 machine_id 是否已存在） + if agent == nil { + var err error + agent, isNewAgent, err = c.agentService.RegisterByToken(token, machineID, ip) + if err != nil { + c.wsManager.RecordConnectFail(ip) + ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()}) + return + } + } + + if !agent.Enabled { + c.wsManager.RecordConnectFail(ip) + ctx.JSON(http.StatusForbidden, gin.H{"error": "Agent 已禁用"}) + return + } + + conn, err := agentUpgrader.Upgrade(ctx.Writer, ctx.Request, nil) + if err != nil { + logger.Errorf("[AgentWS] 升级连接失败: %v", err) + return + } + + // 连接成功，重置失败计数 + c.wsManager.RecordConnectSuccess(ip) + + // 注册连接 + ac := c.wsManager.Register(agent.ID, conn, ip) + + // 更新 Agent 状态 + c.agentService.Heartbeat(token, ip, "", "", "", "", "") + + // 发送连接成功消息（包含注册状态） + c.wsManager.SendToAgent(agent.ID, services.WSTypeConnected, map[string]interface{}{ + "agent_id": agent.ID, + "name": agent.Name, + "is_new_agent": isNewAgent, + "machine_id": machineID, + }) + + // 启动读写协程 + go c.wsWritePump(ac) + go c.wsReadPump(ac, agent) +} + +// wsReadPump 读取消息 +func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) { + defer func() { + c.wsManager.Unregister(agent.ID) + }() + + ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second)) + ac.Conn.SetPongHandler(func(string) error { + ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second)) + return nil + }) + + for { + _, message, err := ac.Conn.ReadMessage() + if err != nil { + if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { + logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err) + } + break + } + + var msg services.WSMessage + if err := json.Unmarshal(message, &msg); err != nil { + continue + } + + c.handleWSMessage(ac, agent, &msg) + } +} + +// wsWritePump 写入消息 +func (c *AgentController) wsWritePump(ac *services.AgentConnection) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + for { + select { + case message, ok := <-ac.Send: + if !ok { + return + } + if err := ac.WriteMessage(message); err != nil { + return + } + case <-ticker.C: + ac.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if err := ac.Conn.WriteMessage(websocket.PingMessage, nil); err != nil { + return + } + } + } +} + +// handleWSMessage 处理 WebSocket 消息 +func (c *AgentController) handleWSMessage(ac *services.AgentConnection, agent *models.Agent, msg *services.WSMessage) { + switch msg.Type { + case services.WSTypeHeartbeat: + c.handleHeartbeat(ac, agent, msg.Data) + + case services.WSTypeTaskResult: + c.handleTaskResult(agent, msg.Data) + + case services.WSTypeFetchTasks: + c.handleFetchTasks(agent) + } +} + +// handleFetchTasks 处理 Agent 请求任务列表 +func (c *AgentController) handleFetchTasks(agent *models.Agent) { + tasks := c.agentService.GetTasks(agent.ID) + c.wsManager.SendToAgent(agent.ID, services.WSTypeTasks, map[string]interface{}{ + "tasks": tasks, + }) + logger.Infof("[AgentWS] Agent #%d 请求任务列表，返回 %d 个任务", agent.ID, len(tasks)) +} + +// handleHeartbeat 处理心跳 +func (c *AgentController) handleHeartbeat(ac *services.AgentConnection, agent *models.Agent, data json.RawMessage) { + var req struct { + Version string `json:"version"` + BuildTime string `json:"build_time"` + Hostname string `json:"hostname"` + OS string `json:"os"` + Arch string `json:"arch"` + AutoUpdate bool `json:"auto_update"` + } + json.Unmarshal(data, &req) + + ac.UpdatePing() + + // 更新 Agent 信息（使用连接时保存的 IP） + c.agentService.Heartbeat(agent.Token, ac.IP, req.Version, req.BuildTime, req.Hostname, req.OS, req.Arch) + + // 检查是否需要更新 + latestVersion := c.agentService.GetLatestVersion() + needUpdate := latestVersion != "" && req.Version != "" && req.Version != latestVersion + forceUpdate := agent.ForceUpdate + + if forceUpdate && needUpdate { + c.agentService.ClearForceUpdate(agent.ID) + } + + // 发送心跳响应 + response := map[string]interface{}{ + "agent_id": agent.ID, + "name": agent.Name, + "need_update": needUpdate, + "force_update": forceUpdate, + "latest_version": latestVersion, + } + c.wsManager.SendToAgent(agent.ID, services.WSTypeHeartbeatAck, response) +} + +// handleTaskResult 处理任务结果 +func (c *AgentController) handleTaskResult(agent *models.Agent, data json.RawMessage) { + var result models.AgentTaskResult + if err := json.Unmarshal(data, &result); err != nil { + return + } + + result.AgentID = agent.ID + c.agentService.ReportResult(&result) +} + +// NotifyTaskUpdate 通知 Agent 任务更新 +func (c *AgentController) NotifyTaskUpdate(agentID uint) { + c.wsManager.BroadcastTasks(agentID) +} + +// ========== 注册码管理 ========== + +// ListRegCodes 获取注册码列表 +func (c *AgentController) ListRegCodes(ctx *gin.Context) { + codes := c.agentService.ListRegCodes() + utils.Success(ctx, codes) +} + +// CreateRegCode 创建注册码 +func (c *AgentController) CreateRegCode(ctx *gin.Context) { + var req struct { + Remark string `json:"remark"` + MaxUses int `json:"max_uses"` + ExpiresAt string `json:"expires_at"` // 格式: 2006-01-02 15:04:05 + } + + if err := ctx.ShouldBindJSON(&req); err != nil { + utils.BadRequest(ctx, "参数错误") + return + } + + var expiresAt *time.Time + if req.ExpiresAt != "" { + t, err := time.ParseInLocation("2006-01-02 15:04:05", req.ExpiresAt, time.Local) + if err != nil { + utils.BadRequest(ctx, "过期时间格式错误") + return + } + expiresAt = &t + } + + code, err := c.agentService.CreateRegCode(req.Remark, req.MaxUses, expiresAt) + if err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.Success(ctx, code) +} + +// DeleteRegCode 删除注册码 +func (c *AgentController) DeleteRegCode(ctx *gin.Context) { + id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) + if err != nil { + utils.BadRequest(ctx, "无效的 ID") + return + } + + if err := c.agentService.DeleteRegCode(uint(id)); err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.SuccessMsg(ctx, "删除成功") +} diff --git a/internal/controllers/task_controller.go b/internal/controllers/task_controller.go index 366d0a6..376986b 100644 --- a/internal/controllers/task_controller.go +++ b/internal/controllers/task_controller.go @@ -12,14 +12,16 @@ import ( ) type TaskController struct { - taskService *services.TaskService - cronService *services.CronService + taskService *services.TaskService + cronService *services.CronService + agentWSManager *services.AgentWSManager } func NewTaskController(taskService *services.TaskService, cronService *services.CronService) *TaskController { return &TaskController{ - taskService: taskService, - cronService: cronService, + taskService: taskService, + cronService: cronService, + agentWSManager: services.GetAgentWSManager(), } } @@ -57,6 +59,7 @@ func (tc *TaskController) CreateTask(c *gin.Context) { WorkDir string `json:"work_dir"` CleanConfig string `json:"clean_config"` Envs string `json:"envs"` + AgentID *uint `json:"agent_id"` } if err := c.ShouldBindJSON(&req); err != nil { @@ -78,8 +81,14 @@ func (tc *TaskController) CreateTask(c *gin.Context) { // 转换为绝对路径 workDir := resolveWorkDir(req.WorkDir) - task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config) - tc.cronService.AddTask(task) + task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config, req.AgentID) + + // 如果是 Agent 任务，通知 Agent；否则添加到本地 cron + if task.AgentID != nil && *task.AgentID > 0 { + tc.agentWSManager.BroadcastTasks(*task.AgentID) + } else { + tc.cronService.AddTask(task) + } utils.Success(c, task) } @@ -115,6 +124,13 @@ func (tc *TaskController) UpdateTask(c *gin.Context) { return } + // 获取旧任务信息（用于判断 agent 变更） + oldTask := tc.taskService.GetTaskByID(id) + var oldAgentID *uint + if oldTask != nil { + oldAgentID = oldTask.AgentID + } + var req struct { Name string `json:"name"` Command string `json:"command"` @@ -126,6 +142,7 @@ func (tc *TaskController) UpdateTask(c *gin.Context) { CleanConfig string `json:"clean_config"` Envs string `json:"envs"` Enabled bool `json:"enabled"` + AgentID *uint `json:"agent_id"` } if err := c.ShouldBindJSON(&req); err != nil { @@ -140,16 +157,32 @@ func (tc *TaskController) UpdateTask(c *gin.Context) { } } - task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config) + task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config, req.AgentID) if task == nil { utils.NotFound(c, "任务不存在") return } - if task.Enabled { - tc.cronService.AddTask(task) - } else { + // 处理任务调度 + if task.AgentID != nil && *task.AgentID > 0 { + // Agent 任务：从本地 cron 移除，通知 Agent tc.cronService.RemoveTask(task.ID) + tc.agentWSManager.BroadcastTasks(*task.AgentID) + // 如果 agent 变更了，也通知旧 agent + if oldAgentID != nil && *oldAgentID > 0 && *oldAgentID != *task.AgentID { + tc.agentWSManager.BroadcastTasks(*oldAgentID) + } + } else { + // 本地任务 + if task.Enabled { + tc.cronService.AddTask(task) + } else { + tc.cronService.RemoveTask(task.ID) + } + // 如果之前是 agent 任务，通知旧 agent 移除 + if oldAgentID != nil && *oldAgentID > 0 { + tc.agentWSManager.BroadcastTasks(*oldAgentID) + } } utils.Success(c, task) @@ -162,6 +195,13 @@ func (tc *TaskController) DeleteTask(c *gin.Context) { return } + // 获取任务信息（用于通知 agent） + task := tc.taskService.GetTaskByID(id) + var agentID *uint + if task != nil { + agentID = task.AgentID + } + tc.cronService.RemoveTask(uint(id)) success := tc.taskService.DeleteTask(id) @@ -170,5 +210,10 @@ func (tc *TaskController) DeleteTask(c *gin.Context) { return } + // 如果是 agent 任务，通知 agent + if agentID != nil && *agentID > 0 { + tc.agentWSManager.BroadcastTasks(*agentID) + } + utils.SuccessMsg(c, "删除成功") } diff --git a/internal/database/migrate.go b/internal/database/migrate.go index 9ee9da7..c27eced 100644 --- a/internal/database/migrate.go +++ b/internal/database/migrate.go @@ -1,10 +1,16 @@ package database import ( + "baihu/internal/logger" "baihu/internal/models" ) func Migrate() error { + // 先执行自定义迁移 + if err := customMigrations(); err != nil { + logger.Warnf("[Database] 自定义迁移警告: %v", err) + } + return AutoMigrate( &models.User{}, &models.Task{}, @@ -16,5 +22,19 @@ func Migrate() error { &models.SendStats{}, &models.Dependency{}, &models.Agent{}, + &models.AgentRegCode{}, ) } + +// customMigrations 自定义迁移（处理 AutoMigrate 无法自动完成的变更） +func customMigrations() error { + // 检查 ql_tokens 表是否存在，如果存在则修改 code 列大小为 64 + if DB.Migrator().HasTable("ql_tokens") { + // MySQL: 修改 code 列大小 + if err := DB.Exec("ALTER TABLE ql_tokens MODIFY COLUMN code VARCHAR(64)").Error; err != nil { + // 忽略错误（可能是 SQLite 或列已经是正确大小） + logger.Debugf("[Database] 修改 ql_tokens.code 列: %v", err) + } + } + return nil +} diff --git a/internal/models/agent.go b/internal/models/agent.go index b1ddcb6..722b7b8 100644 --- a/internal/models/agent.go +++ b/internal/models/agent.go @@ -10,9 +10,10 @@ import ( type Agent struct { ID uint `json:"id" gorm:"primaryKey"` Name string `json:"name" gorm:"size:100;not null"` // Agent 名称 - Token string `json:"token" gorm:"size:64;uniqueIndex"` // 认证 Token + Token string `json:"token" gorm:"size:64;index"` // 认证 Token（可重复使用） + MachineID string `json:"machine_id" gorm:"size:64;uniqueIndex"` // 机器识别码（唯一） Description string `json:"description" gorm:"size:255"` // 描述 - Status string `json:"status" gorm:"size:20;default:'pending'"` // 状态: pending(待审核), online, offline + Status string `json:"status" gorm:"size:20;default:'pending'"` // 状态: pending(待审核), online, offline, blocked(拉黑) LastSeen *LocalTime `json:"last_seen"` // 最后心跳时间 IP string `json:"ip" gorm:"size:45"` // Agent IP 地址 Version string `json:"version" gorm:"size:20"` // Agent 版本 @@ -31,6 +32,24 @@ func (Agent) TableName() string { return constant.TablePrefix + "agents" } +// AgentRegCode 注册码 +type AgentRegCode struct { + ID uint `json:"id" gorm:"primaryKey"` + Code string `json:"code" gorm:"size:64;uniqueIndex;not null"` // 令牌 + Remark string `json:"remark" gorm:"size:255"` // 备注 + MaxUses int `json:"max_uses" gorm:"default:0"` // 最大使用次数，0 表示无限制 + UsedCount int `json:"used_count" gorm:"default:0"` // 已使用次数 + ExpiresAt *LocalTime `json:"expires_at"` // 过期时间，null 表示永不过期 + Enabled bool `json:"enabled" gorm:"default:true"` // 是否启用 + CreatedAt LocalTime `json:"created_at"` + UpdatedAt LocalTime `json:"updated_at"` + DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` +} + +func (AgentRegCode) TableName() string { + return constant.TablePrefix + "tokens" +} + // AgentTask Agent 任务配置（用于下发给 Agent） type AgentTask struct { ID uint `json:"id"` @@ -58,7 +77,10 @@ type AgentTaskResult struct { // AgentRegisterRequest Agent 注册请求 type AgentRegisterRequest struct { - Name string `json:"name"` - Hostname string `json:"hostname"` - Version string `json:"version"` + Name string `json:"name"` + Hostname string `json:"hostname"` + Version string `json:"version"` + BuildTime string `json:"build_time"` + Token string `json:"token"` // 注册令牌 + MachineID string `json:"machine_id"` // 机器识别码 } diff --git a/internal/router/router.go b/internal/router/router.go index 6c6c9dd..1b6110c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -203,26 +203,26 @@ func Setup(c *Controllers) *gin.Engine { agents := authorized.Group("/agents") { agents.GET("", c.Agent.List) - agents.GET("/pending", c.Agent.ListPending) agents.GET("/version", c.Agent.GetVersion) - agents.POST("/:id/approve", c.Agent.Approve) - agents.POST("/:id/reject", c.Agent.Reject) agents.PUT("/:id", c.Agent.Update) agents.DELETE("/:id", c.Agent.Delete) agents.POST("/:id/token", c.Agent.RegenerateToken) agents.POST("/:id/update", c.Agent.ForceUpdate) + // 令牌管理 + agents.GET("/regcodes", c.Agent.ListRegCodes) + agents.POST("/regcodes", c.Agent.CreateRegCode) + agents.DELETE("/regcodes/:id", c.Agent.DeleteRegCode) } } // Agent API（供远程 Agent 调用） agentAPI := api.Group("/agent") { - agentAPI.POST("/register", c.Agent.Register) - agentAPI.POST("/status", c.Agent.CheckStatus) agentAPI.POST("/heartbeat", c.Agent.Heartbeat) agentAPI.GET("/tasks", c.Agent.GetTasks) agentAPI.POST("/report", c.Agent.ReportResult) agentAPI.GET("/download", c.Agent.Download) + agentAPI.GET("/ws", c.Agent.WSConnect) // WebSocket 连接 } } diff --git a/internal/services/agent_service.go b/internal/services/agent_service.go index 7be7b40..1aec1d0 100644 --- a/internal/services/agent_service.go +++ b/internal/services/agent_service.go @@ -12,6 +12,8 @@ import ( "path/filepath" "strings" "time" + + "gorm.io/gorm" ) // AgentService Agent 服务 @@ -29,75 +31,171 @@ func generateToken() string { return hex.EncodeToString(bytes) } -// Register Agent 注册（进入待审核状态） -func (s *AgentService) Register(req *models.AgentRegisterRequest, ip string) (*models.Agent, error) { - // 检查是否已存在同名待审核的 Agent - var existing models.Agent - if err := database.DB.Where("name = ? AND status = ?", req.Name, "pending").First(&existing).Error; err == nil { - // 更新现有记录 - now := models.LocalTime(time.Now()) - database.DB.Model(&existing).Updates(map[string]interface{}{ - "hostname": req.Hostname, - "version": req.Version, - "ip": ip, - "last_seen": now, - }) - return &existing, nil +// generateRegCode 生成令牌（64位，与认证 Token 相同） +func generateRegCode() string { + bytes := make([]byte, 32) + rand.Read(bytes) + return hex.EncodeToString(bytes) +} + +// ========== 注册码管理 ========== + +// CreateRegCode 创建令牌（同时创建 Agent 记录） +func (s *AgentService) CreateRegCode(remark string, maxUses int, expiresAt *time.Time) (*models.AgentRegCode, error) { + var expires *models.LocalTime + if expiresAt != nil { + t := models.LocalTime(*expiresAt) + expires = &t } - // 创建新的待审核 Agent + token := generateRegCode() + + regCode := &models.AgentRegCode{ + Code: token, + Remark: remark, + MaxUses: maxUses, + ExpiresAt: expires, + Enabled: true, + } + + if err := database.DB.Create(regCode).Error; err != nil { + return nil, err + } + + logger.Infof("[Agent] 创建令牌: %s (max_uses=%d)", token[:8]+"...", maxUses) + return regCode, nil +} + +// ListRegCodes 获取注册码列表 +func (s *AgentService) ListRegCodes() []models.AgentRegCode { + var codes []models.AgentRegCode + database.DB.Order("id DESC").Find(&codes) + return codes +} + +// DeleteRegCode 删除注册码 +func (s *AgentService) DeleteRegCode(id uint) error { + return database.DB.Delete(&models.AgentRegCode{}, id).Error +} + +// ValidateRegCode 验证注册码 +func (s *AgentService) ValidateRegCode(code string) (*models.AgentRegCode, error) { + var regCode models.AgentRegCode + if err := database.DB.Where("code = ?", code).First(®Code).Error; err != nil { + return nil, &ServiceError{Message: "无效的注册码"} + } + + if !regCode.Enabled { + return nil, &ServiceError{Message: "注册码已禁用"} + } + + // 检查使用次数 + if regCode.MaxUses > 0 && regCode.UsedCount >= regCode.MaxUses { + return nil, &ServiceError{Message: "注册码已达到使用上限"} + } + + // 检查过期时间 + if regCode.ExpiresAt != nil && time.Time(*regCode.ExpiresAt).Before(time.Now()) { + return nil, &ServiceError{Message: "注册码已过期"} + } + + return ®Code, nil +} + +// UseRegCode 使用注册码（增加使用计数） +func (s *AgentService) UseRegCode(id uint) { + database.DB.Model(&models.AgentRegCode{}).Where("id = ?", id).UpdateColumn("used_count", gorm.Expr("used_count + 1")) +} + +// ========== Agent 注册 ========== + +// RegisterByToken 通过令牌注册 Agent（首次 WebSocket 连接时调用） +// 返回: agent, isNewAgent, error +func (s *AgentService) RegisterByToken(token string, machineID string, ip string) (*models.Agent, bool, error) { + // 验证令牌 + regCode, err := s.ValidateRegCode(token) + if err != nil { + return nil, false, err + } + + // 如果提供了 machine_id，先检查是否已存在 + if machineID != "" { + var existing models.Agent + if err := database.DB.Where("machine_id = ?", machineID).First(&existing).Error; err == nil { + // 已存在，更新 token 和状态，复用已有 Agent + now := models.LocalTime(time.Now()) + database.DB.Model(&existing).Updates(map[string]interface{}{ + "token": token, + "ip": ip, + "status": "online", + "last_seen": now, + }) + s.UseRegCode(regCode.ID) + logger.Infof("[Agent] Agent #%d 通过 machine_id 复用 (%s)", existing.ID, machineID[:8]+"...") + return &existing, false, nil + } + } + + // 创建 Agent，使用令牌作为认证 Token now := models.LocalTime(time.Now()) agent := &models.Agent{ - Name: req.Name, - Hostname: req.Hostname, - Version: req.Version, - IP: ip, - Status: "pending", - LastSeen: &now, - Enabled: true, + Name: fmt.Sprintf("agent-%d", time.Now().Unix()), + Token: token, + MachineID: machineID, + IP: ip, + Status: "online", + LastSeen: &now, + Enabled: true, } if err := database.DB.Create(agent).Error; err != nil { - return nil, err + return nil, false, err } - logger.Infof("[Agent] 新 Agent 注册: %s (%s)", req.Name, ip) - return agent, nil + s.UseRegCode(regCode.ID) + logger.Infof("[Agent] Agent 通过令牌注册: #%d (%s)", agent.ID, ip) + return agent, true, nil } -// Approve 审核通过 Agent，生成 Token -func (s *AgentService) Approve(id uint) (*models.Agent, error) { - agent := s.GetByID(id) - if agent == nil { - return nil, &ServiceError{Message: "Agent 不存在"} +// Register Agent 注册（必须使用令牌）- 保留兼容旧版本 +func (s *AgentService) Register(req *models.AgentRegisterRequest, ip string) (*models.Agent, string, error) { + // 必须提供令牌 + if req.Token == "" { + return nil, "", &ServiceError{Message: "缺少注册令牌"} } - if agent.Status != "pending" { - return nil, &ServiceError{Message: "Agent 状态不是待审核"} + regCode, err := s.ValidateRegCode(req.Token) + if err != nil { + return nil, "", err } - token := generateToken() + // 检查是否已存在同名 Agent + var existing models.Agent + if err := database.DB.Where("name = ?", req.Name).First(&existing).Error; err == nil { + return nil, "", &ServiceError{Message: "Agent 名称已存在"} + } + + // 创建新 Agent，使用令牌作为认证 Token now := models.LocalTime(time.Now()) - - if err := database.DB.Model(agent).Updates(map[string]interface{}{ - "token": token, - "status": "online", - "last_seen": now, - }).Error; err != nil { - return nil, err + agent := &models.Agent{ + Name: req.Name, + Token: req.Token, + Hostname: req.Hostname, + Version: req.Version, + BuildTime: req.BuildTime, + IP: ip, + Status: "online", + LastSeen: &now, + Enabled: true, } - agent.Token = token - agent.Status = "online" - agent.LastSeen = &now + if err := database.DB.Create(agent).Error; err != nil { + return nil, "", err + } - logger.Infof("[Agent] Agent 已审核通过: %s (#%d)", agent.Name, agent.ID) - return agent, nil -} - -// Reject 拒绝 Agent -func (s *AgentService) Reject(id uint) error { - return database.DB.Delete(&models.Agent{}, id).Error + s.UseRegCode(regCode.ID) + logger.Infof("[Agent] Agent 注册成功: %s (%s)", req.Name, ip) + return agent, req.Token, nil } // Update 更新 Agent @@ -109,7 +207,7 @@ func (s *AgentService) Update(id uint, name, description string, enabled bool) e }).Error } -// Delete 删除 Agent +// Delete 删除 Agent（物理删除） func (s *AgentService) Delete(id uint) error { // 检查是否有关联任务 var count int64 @@ -118,7 +216,7 @@ func (s *AgentService) Delete(id uint) error { return &ServiceError{Message: "该 Agent 下还有关联任务，无法删除"} } - return database.DB.Delete(&models.Agent{}, id).Error + return database.DB.Unscoped().Delete(&models.Agent{}, id).Error } // GetByID 根据 ID 获取 Agent @@ -139,27 +237,16 @@ func (s *AgentService) GetByToken(token string) *models.Agent { return &agent } -// List 获取已审核的 Agent 列表 +// List 获取 Agent 列表 func (s *AgentService) List() []models.Agent { var agents []models.Agent - database.DB.Where("status != ?", "pending").Order("id DESC").Find(&agents) + database.DB.Order("id DESC").Find(&agents) return agents } -// ListPending 获取待审核的 Agent 列表 -func (s *AgentService) ListPending() []models.Agent { - var agents []models.Agent - database.DB.Where("status = ?", "pending").Order("id DESC").Find(&agents) - return agents -} - -// RegenerateToken 重新生成 Token +// RegenerateToken 重新生成 Token - 已废弃，保留空实现避免路由错误 func (s *AgentService) RegenerateToken(id uint) (string, error) { - newToken := generateToken() - if err := database.DB.Model(&models.Agent{}).Where("id = ?", id).Update("token", newToken).Error; err != nil { - return "", err - } - return newToken, nil + return "", &ServiceError{Message: "此功能已禁用"} } // Heartbeat Agent 心跳 @@ -209,20 +296,6 @@ func (s *AgentService) Heartbeat(token, ip, version, buildTime, hostname, osType return agent, nil } -// CheckPendingAgent 检查待审核 Agent 的状态（用于 Agent 轮询） -func (s *AgentService) CheckPendingAgent(name, ip string) (*models.Agent, error) { - var agent models.Agent - if err := database.DB.Where("name = ? AND ip = ?", name, ip).First(&agent).Error; err != nil { - return nil, &ServiceError{Message: "Agent 未注册"} - } - - // 更新最后心跳时间 - now := models.LocalTime(time.Now()) - database.DB.Model(&agent).Update("last_seen", now) - - return &agent, nil -} - // GetTasks 获取 Agent 的任务列表 func (s *AgentService) GetTasks(agentID uint) []models.AgentTask { var tasks []models.Task diff --git a/internal/services/agent_ws_service.go b/internal/services/agent_ws_service.go new file mode 100644 index 0000000..9faca90 --- /dev/null +++ b/internal/services/agent_ws_service.go @@ -0,0 +1,293 @@ +package services + +import ( + "baihu/internal/database" + "baihu/internal/logger" + "baihu/internal/models" + "encoding/json" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// AgentWSManager WebSocket 连接管理器 +type AgentWSManager struct { + connections map[uint]*AgentConnection // agentID -> connection + ipConnections map[string]int // IP -> 连接数 + ipLastAttempt map[string]time.Time // IP -> 最后连接尝试时间 + ipFailCount map[string]int // IP -> 连续失败次数 + mu sync.RWMutex +} + +// 限流配置 +const ( + maxConnectionsPerIP = 10 // 每个 IP 最大连接数 + minConnectInterval = 5 * time.Second // 同一 IP 最小连接间隔 + maxFailCount = 5 // 最大连续失败次数 + failBlockDuration = 5 * time.Minute // 失败后封禁时长 +) + +// AgentConnection Agent WebSocket 连接 +type AgentConnection struct { + AgentID uint + IP string + Conn *websocket.Conn + Send chan []byte + LastPing time.Time + mu sync.Mutex +} + +// WSMessage WebSocket 消息结构 +type WSMessage struct { + Type string `json:"type"` + Data json.RawMessage `json:"data,omitempty"` +} + +// 消息类型常量 +const ( + WSTypeHeartbeat = "heartbeat" + WSTypeHeartbeatAck = "heartbeat_ack" + WSTypeTasks = "tasks" + WSTypeTaskResult = "task_result" + WSTypeUpdate = "update" + WSTypeDisconnect = "disconnect" + WSTypeConnected = "connected" // 连接成功，包含注册状态 + WSTypeDisabled = "disabled" // Agent 被禁用 + WSTypeEnabled = "enabled" // Agent 被启用 + WSTypeFetchTasks = "fetch_tasks" // Agent 请求任务列表 +) + +var agentWSManager *AgentWSManager +var agentWSOnce sync.Once + +// GetAgentWSManager 获取单例 +func GetAgentWSManager() *AgentWSManager { + agentWSOnce.Do(func() { + agentWSManager = &AgentWSManager{ + connections: make(map[uint]*AgentConnection), + ipConnections: make(map[string]int), + ipLastAttempt: make(map[string]time.Time), + ipFailCount: make(map[string]int), + } + go agentWSManager.cleanupLoop() + }) + return agentWSManager +} + +// CheckRateLimit 检查 IP 限流，返回是否允许连接 +func (m *AgentWSManager) CheckRateLimit(ip string) (bool, string) { + m.mu.Lock() + defer m.mu.Unlock() + + now := time.Now() + + // 检查是否被封禁（连续失败过多） + if failCount, exists := m.ipFailCount[ip]; exists && failCount >= maxFailCount { + if lastAttempt, ok := m.ipLastAttempt[ip]; ok { + if now.Sub(lastAttempt) < failBlockDuration { + remaining := failBlockDuration - now.Sub(lastAttempt) + return false, "连接失败次数过多，请 " + remaining.Round(time.Second).String() + " 后重试" + } + // 封禁时间已过，重置计数 + delete(m.ipFailCount, ip) + } + } + + // 检查连接频率 + if lastAttempt, exists := m.ipLastAttempt[ip]; exists { + if now.Sub(lastAttempt) < minConnectInterval { + return false, "连接过于频繁，请稍后重试" + } + } + + // 检查 IP 连接数 + if count, exists := m.ipConnections[ip]; exists && count >= maxConnectionsPerIP { + return false, "该 IP 连接数已达上限" + } + + m.ipLastAttempt[ip] = now + return true, "" +} + +// RecordConnectFail 记录连接失败 +func (m *AgentWSManager) RecordConnectFail(ip string) { + m.mu.Lock() + defer m.mu.Unlock() + m.ipFailCount[ip]++ + m.ipLastAttempt[ip] = time.Now() + if m.ipFailCount[ip] >= maxFailCount { + logger.Warnf("[AgentWS] IP %s 连续失败 %d 次，已封禁 %v", ip, m.ipFailCount[ip], failBlockDuration) + } +} + +// RecordConnectSuccess 记录连接成功，重置失败计数 +func (m *AgentWSManager) RecordConnectSuccess(ip string) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.ipFailCount, ip) +} + +// Register 注册连接 +func (m *AgentWSManager) Register(agentID uint, conn *websocket.Conn, ip string) *AgentConnection { + m.mu.Lock() + defer m.mu.Unlock() + + // 关闭旧连接 + if old, exists := m.connections[agentID]; exists { + // 减少旧 IP 的连接计数 + if old.IP != "" { + if count, ok := m.ipConnections[old.IP]; ok && count > 0 { + m.ipConnections[old.IP] = count - 1 + } + } + old.Close() + } + + ac := &AgentConnection{ + AgentID: agentID, + IP: ip, + Conn: conn, + Send: make(chan []byte, 256), + LastPing: time.Now(), + } + m.connections[agentID] = ac + + // 增加 IP 连接计数 + m.ipConnections[ip]++ + + logger.Infof("[AgentWS] Agent #%d 已连接 (%s)", agentID, ip) + return ac +} + +// Unregister 注销连接 +func (m *AgentWSManager) Unregister(agentID uint) { + m.mu.Lock() + defer m.mu.Unlock() + + if conn, exists := m.connections[agentID]; exists { + // 减少 IP 连接计数 + if conn.IP != "" { + if count, ok := m.ipConnections[conn.IP]; ok && count > 0 { + m.ipConnections[conn.IP] = count - 1 + } + } + conn.Close() + delete(m.connections, agentID) + logger.Infof("[AgentWS] Agent #%d 已断开", agentID) + } +} + +// GetConnection 获取连接 +func (m *AgentWSManager) GetConnection(agentID uint) *AgentConnection { + m.mu.RLock() + defer m.mu.RUnlock() + return m.connections[agentID] +} + +// SendToAgent 发送消息给指定 Agent +func (m *AgentWSManager) SendToAgent(agentID uint, msgType string, data interface{}) error { + conn := m.GetConnection(agentID) + if conn == nil { + return nil // Agent 不在线 + } + + dataBytes, _ := json.Marshal(data) + msg := WSMessage{Type: msgType, Data: dataBytes} + msgBytes, _ := json.Marshal(msg) + + select { + case conn.Send <- msgBytes: + return nil + default: + return nil // 缓冲区满，丢弃 + } +} + +// BroadcastTasks 广播任务更新给指定 Agent +func (m *AgentWSManager) BroadcastTasks(agentID uint) { + agentService := NewAgentService() + tasks := agentService.GetTasks(agentID) + m.SendToAgent(agentID, WSTypeTasks, map[string]interface{}{ + "tasks": tasks, + }) +} + +// OnlineCount 在线 Agent 数量 +func (m *AgentWSManager) OnlineCount() int { + m.mu.RLock() + defer m.mu.RUnlock() + return len(m.connections) +} + +// cleanupLoop 清理超时连接 +func (m *AgentWSManager) cleanupLoop() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + + for range ticker.C { + m.mu.Lock() + now := time.Now() + + // 清理超时连接 + for agentID, conn := range m.connections { + if now.Sub(conn.LastPing) > 2*time.Minute { + // 减少 IP 连接计数 + if conn.IP != "" { + if count, ok := m.ipConnections[conn.IP]; ok && count > 0 { + m.ipConnections[conn.IP] = count - 1 + } + } + conn.Close() + delete(m.connections, agentID) + // 更新数据库状态 + database.DB.Model(&models.Agent{}).Where("id = ?", agentID).Update("status", "offline") + logger.Infof("[AgentWS] Agent #%d 心跳超时，已断开", agentID) + } + } + + // 清理过期的限流记录（超过 10 分钟未活动） + for ip, lastAttempt := range m.ipLastAttempt { + if now.Sub(lastAttempt) > 10*time.Minute { + delete(m.ipLastAttempt, ip) + delete(m.ipFailCount, ip) + // 只清理没有活跃连接的 IP 计数 + if m.ipConnections[ip] == 0 { + delete(m.ipConnections, ip) + } + } + } + + m.mu.Unlock() + } +} + +// Close 关闭连接 +func (c *AgentConnection) Close() { + c.mu.Lock() + defer c.mu.Unlock() + if c.Conn != nil { + c.Conn.Close() + c.Conn = nil + } + if c.Send != nil { + close(c.Send) + c.Send = nil + } +} + +// WriteMessage 写入消息 +func (c *AgentConnection) WriteMessage(data []byte) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.Conn == nil { + return nil + } + c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) + return c.Conn.WriteMessage(websocket.TextMessage, data) +} + +// UpdatePing 更新心跳时间 +func (c *AgentConnection) UpdatePing() { + c.LastPing = time.Now() +} diff --git a/internal/services/task_service.go b/internal/services/task_service.go index b9aea18..845c1e1 100644 --- a/internal/services/task_service.go +++ b/internal/services/task_service.go @@ -11,7 +11,7 @@ func NewTaskService() *TaskService { return &TaskService{} } -func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string) *models.Task { +func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string, agentID *uint) *models.Task { if taskType == "" { taskType = "task" } @@ -25,6 +25,7 @@ func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, w WorkDir: workDir, CleanConfig: cleanConfig, Envs: envs, + AgentID: agentID, Enabled: true, } database.DB.Create(task) @@ -61,7 +62,7 @@ func (ts *TaskService) GetTaskByID(id int) *models.Task { return &task } -func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string) *models.Task { +func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string, agentID *uint) *models.Task { var task models.Task if err := database.DB.First(&task, id).Error; err != nil { return nil @@ -74,6 +75,7 @@ func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeou task.CleanConfig = cleanConfig task.Envs = envs task.Enabled = enabled + task.AgentID = agentID if taskType != "" { task.Type = taskType } diff --git a/web/src/api/index.ts b/web/src/api/index.ts index e3f24f7..c3fd3f7 100644 --- a/web/src/api/index.ts +++ b/web/src/api/index.ts @@ -209,16 +209,17 @@ export const api = { }, agents: { list: () => request('/agents'), - listPending: () => request('/agents/pending'), getVersion: () => request<{ version: string; platforms: { os: string; arch: string; filename: string }[] }>('/agents/version'), - approve: (id: number) => request('/agents/' + id + '/approve', { method: 'POST' }), - reject: (id: number) => request('/agents/' + id + '/reject', { method: 'POST' }), update: (id: number, data: { name: string; description?: string; enabled: boolean }) => request('/agents/' + id, { method: 'PUT', body: JSON.stringify(data) }), delete: (id: number) => request('/agents/' + id, { method: 'DELETE' }), - regenerateToken: (id: number) => request<{ token: string }>('/agents/' + id + '/token', { method: 'POST' }), forceUpdate: (id: number) => request('/agents/' + id + '/update', { method: 'POST' }), - downloadUrl: (os: string, arch: string) => `${BASE_URL}/agent/download?os=${os}&arch=${arch}` + downloadUrl: (os: string, arch: string) => `${BASE_URL}/agent/download?os=${os}&arch=${arch}`, + // 令牌管理 + listRegCodes: () => request('/agents/regcodes'), + createRegCode: (data: { remark?: string; max_uses?: number; expires_at?: string }) => + request('/agents/regcodes', { method: 'POST', body: JSON.stringify(data) }), + deleteRegCode: (id: number) => request('/agents/regcodes/' + id, { method: 'DELETE' }) } } @@ -394,6 +395,7 @@ export interface Agent { id: number name: string token: string + machine_id: string description: string status: string last_seen: string @@ -401,7 +403,20 @@ export interface Agent { version: string build_time: string hostname: string + os: string + arch: string enabled: boolean created_at: string updated_at: string } + +export interface AgentRegCode { + id: number + code: string + remark: string + max_uses: number + used_count: number + expires_at: string | null + enabled: boolean + created_at: string +} diff --git a/web/src/views/agents/Agents.vue b/web/src/views/agents/Agents.vue index 96af6c9..52b8ae5 100644 --- a/web/src/views/agents/Agents.vue +++ b/web/src/views/agents/Agents.vue @@ -3,31 +3,29 @@ import { ref, onMounted, computed, onUnmounted } from 'vue' import { Button } from '@/components/ui/button' import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' -import { Badge } from '@/components/ui/badge' -import { Switch } from '@/components/ui/switch' import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogFooter, DialogDescription } from '@/components/ui/dialog' import { AlertDialog, AlertDialogAction, AlertDialogCancel, AlertDialogContent, AlertDialogDescription, AlertDialogFooter, AlertDialogHeader, AlertDialogTitle } from '@/components/ui/alert-dialog' import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs' -import { RefreshCw, Trash2, Edit, Copy, Key, Server, Search, Check, X, Download, RotateCw } from 'lucide-vue-next' -import { api, type Agent } from '@/api' +import { RefreshCw, Trash2, Edit, Copy, Server, Search, Download, RotateCw, Plus, Ticket, Power, PowerOff } from 'lucide-vue-next' +import { api, type Agent, type AgentRegCode } from '@/api' import { toast } from 'vue-sonner' import TextOverflow from '@/components/TextOverflow.vue' const agents = ref([]) -const pendingAgents = ref([]) +const regCodes = ref([]) const loading = ref(false) const searchQuery = ref('') -const activeTab = ref('approved') +const activeTab = ref('agents') const agentVersion = ref('') const platforms = ref<{ os: string; arch: string; filename: string }[]>([]) const showEditDialog = ref(false) const showDeleteDialog = ref(false) -const showTokenDialog = ref(false) const showDownloadDialog = ref(false) +const showRegCodeDialog = ref(false) const formData = ref({ name: '', description: '' }) +const regCodeForm = ref({ remark: '', max_uses: 0, expires_at: '' }) const editingAgent = ref(null) const deletingAgent = ref(null) -const currentToken = ref('') let refreshTimer: ReturnType | null = null const filteredAgents = computed(() => { @@ -40,18 +38,27 @@ const filteredAgents = computed(() => { ) }) +// 判断 Agent 是否在线（last_seen 在 2 分钟内） +function isOnline(agent: Agent): boolean { + if (!agent.last_seen) return false + const lastSeen = new Date(agent.last_seen) + const now = new Date() + const diffMs = now.getTime() - lastSeen.getTime() + return diffMs < 2 * 60 * 1000 // 2 分钟 +} + async function loadAgents() { loading.value = true try { - const [agentList, pendingList, versionInfo] = await Promise.all([ + const [agentList, versionInfo, codeList] = await Promise.all([ api.agents.list(), - api.agents.listPending(), - api.agents.getVersion() + api.agents.getVersion(), + api.agents.listRegCodes() ]) agents.value = agentList - pendingAgents.value = pendingList agentVersion.value = versionInfo.version || '' platforms.value = versionInfo.platforms || [] + regCodes.value = codeList } catch { toast.error('加载失败') } finally { @@ -59,28 +66,6 @@ async function loadAgents() { } } -async function approveAgent(agent: Agent) { - try { - const approved = await api.agents.approve(agent.id) - currentToken.value = approved.token - showTokenDialog.value = true - await loadAgents() - toast.success('已通过审核') - } catch (e: unknown) { - toast.error((e as Error).message || '操作失败') - } -} - -async function rejectAgent(agent: Agent) { - try { - await api.agents.reject(agent.id) - await loadAgents() - toast.success('已拒绝') - } catch (e: unknown) { - toast.error((e as Error).message || '操作失败') - } -} - function openEditDialog(agent: Agent) { editingAgent.value = agent formData.value = { name: agent.name, description: agent.description } @@ -101,8 +86,10 @@ async function updateAgent() { async function toggleEnabled(agent: Agent) { try { - await api.agents.update(agent.id, { name: agent.name, description: agent.description, enabled: !agent.enabled }) + const newEnabled = !agent.enabled + await api.agents.update(agent.id, { name: agent.name, description: agent.description, enabled: newEnabled }) await loadAgents() + toast.success(`${agent.name} 已${newEnabled ? '启用' : '禁用'}`) } catch (e: unknown) { toast.error((e as Error).message || '操作失败') } @@ -125,17 +112,6 @@ async function deleteAgent() { } } -async function regenerateToken(agent: Agent) { - try { - const res = await api.agents.regenerateToken(agent.id) - currentToken.value = res.token - showTokenDialog.value = true - toast.success('Token 已重新生成') - } catch (e: unknown) { - toast.error((e as Error).message || '操作失败') - } -} - async function forceUpdate(agent: Agent) { try { await api.agents.forceUpdate(agent.id) @@ -145,11 +121,46 @@ async function forceUpdate(agent: Agent) { } } -function copyToken() { - navigator.clipboard.writeText(currentToken.value) +function copyRegCode(code: string) { + navigator.clipboard.writeText(code) toast.success('已复制') } +async function createRegCode() { + try { + await api.agents.createRegCode({ + remark: regCodeForm.value.remark, + max_uses: regCodeForm.value.max_uses, + expires_at: regCodeForm.value.expires_at || undefined + }) + showRegCodeDialog.value = false + regCodeForm.value = { remark: '', max_uses: 0, expires_at: '' } + await loadAgents() + toast.success('创建成功') + } catch (e: unknown) { + toast.error((e as Error).message || '创建失败') + } +} + +async function deleteRegCode(id: number) { + try { + await api.agents.deleteRegCode(id) + await loadAgents() + toast.success('删除成功') + } catch (e: unknown) { + toast.error((e as Error).message || '删除失败') + } +} + +function isRegCodeExpired(code: AgentRegCode) { + if (!code.expires_at) return false + return new Date(code.expires_at) < new Date() +} + +function isRegCodeExhausted(code: AgentRegCode) { + return code.max_uses > 0 && code.used_count >= code.max_uses +} + function downloadAgent(os: string, arch: string) { window.open(api.agents.downloadUrl(os, arch), '_blank') } @@ -194,57 +205,61 @@ onUnmounted(() => { - 已注册 - - 未注册 - {{ pendingAgents.length }} + Agent 列表 + + 令牌 - +
-
+
+ 名称 - 状态 IP 主机名 版本 构建时间 + 心跳时间 描述 - 启用 操作
-
+
{{ searchQuery ? '无匹配结果' : '暂无 Agent' }}
- {{ agent.name }} - - {{ agent.status === 'online' ? '在线' : '离线' }} + + + + + + {{ agent.name }} {{ agent.ip || '-' }} {{ agent.hostname || '-' }} {{ agent.version || '-' }} {{ agent.build_time || '-' }} + {{ agent.last_seen || '-' }} - - - + - - - @@ -253,32 +268,45 @@ onUnmounted(() => {
- +
-
- 名称 - IP - 主机名 - 版本 - 注册时间 - 操作 +
+ + 令牌 + 备注 + 使用次数 + 过期时间 + + +
-
-
- 暂无未注册的 Agent +
+
+ 暂无令牌
-
- {{ agent.name }} - {{ agent.ip || '-' }} - {{ agent.hostname || '-' }} - {{ agent.version || '-' }} - {{ agent.created_at }} - - -
@@ -324,27 +352,6 @@ onUnmounted(() => { - - - - - Agent Token - Token 已下发给 Agent，此处仅供查看 - -
-
- {{ currentToken }} - -
-
- - - - - - @@ -354,7 +361,7 @@ onUnmounted(() => {
- 暂无可用的 Agent 程序，请先构建并上传到 data/agent 目录 + 暂无可用的 Agent 程序
@@ -377,5 +384,33 @@ onUnmounted(() => {
+ + + + + + 生成令牌 + Agent 使用令牌可直接注册，无需审核 + +
+
+ + +
+
+ + +
+
+ + +
+
+ + + + + +
diff --git a/web/src/views/history/History.vue b/web/src/views/history/History.vue index 3f0b633..513530f 100644 --- a/web/src/views/history/History.vue +++ b/web/src/views/history/History.vue @@ -184,7 +184,10 @@ watch(() => route.query.task_id, (newTaskId) => { {{ log.task_name }} - + + + + {{ formatDuration(log.duration) }}
@@ -200,7 +203,10 @@ watch(() => route.query.task_id, (newTaskId) => { - + + + + {{ formatDuration(log.duration) }} @@ -230,7 +236,10 @@ watch(() => route.query.task_id, (newTaskId) => {
状态 - + + + + {{ selectedLog.status }}
diff --git a/web/src/views/tasks/RepoDialog.vue b/web/src/views/tasks/RepoDialog.vue index 0d3ef8c..e49321c 100644 --- a/web/src/views/tasks/RepoDialog.vue +++ b/web/src/views/tasks/RepoDialog.vue @@ -7,7 +7,7 @@ import { Label } from '@/components/ui/label' import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select' import { Checkbox } from '@/components/ui/checkbox' import DirTreeSelect from '@/components/DirTreeSelect.vue' -import { api, type Task, type RepoConfig } from '@/api' +import { api, type Task, type RepoConfig, type Agent } from '@/api' import { toast } from 'vue-sonner' const props = defineProps<{ @@ -54,13 +54,19 @@ const repoConfig = ref({ }) const cleanType = ref('none') const cleanKeep = ref(30) +const allAgents = ref([]) +const selectedAgentId = ref('local') const cleanConfig = computed(() => { if (!cleanType.value || cleanType.value === 'none' || cleanKeep.value <= 0) return '' return JSON.stringify({ type: cleanType.value, keep: cleanKeep.value }) }) -watch(() => props.open, (val) => { +const onlineAgents = computed(() => { + return allAgents.value.filter(a => a.enabled) +}) + +watch(() => props.open, async (val) => { if (val) { form.value = { ...props.task } // 解析清理配置 @@ -87,15 +93,26 @@ watch(() => props.open, (val) => { } else { repoConfig.value = { source_type: 'git', source_url: '', target_path: '', branch: '', sparse_path: '', single_file: false, proxy: 'none', proxy_url: '', auth_token: '' } } + // 解析 Agent + selectedAgentId.value = props.task?.agent_id ? String(props.task.agent_id) : 'local' + // 加载 Agent 列表 + await loadAgents() } }) +async function loadAgents() { + try { + allAgents.value = await api.agents.list() + } catch { /* ignore */ } +} + async function save() { try { form.value.clean_config = cleanConfig.value form.value.type = 'repo' form.value.config = JSON.stringify(repoConfig.value) form.value.command = `[${repoConfig.value.source_type}] ${repoConfig.value.source_url}` + form.value.agent_id = selectedAgentId.value === 'local' ? null : Number(selectedAgentId.value) if (props.isEdit && form.value.id) { await api.tasks.update(form.value.id, form.value) toast.success('同步任务已更新') @@ -139,7 +156,24 @@ async function save() {
- + + +
+
+
+ +
+
diff --git a/web/src/views/tasks/TaskDialog.vue b/web/src/views/tasks/TaskDialog.vue index 083e712..0b43785 100644 --- a/web/src/views/tasks/TaskDialog.vue +++ b/web/src/views/tasks/TaskDialog.vue @@ -9,7 +9,7 @@ import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover import { Badge } from '@/components/ui/badge' import DirTreeSelect from '@/components/DirTreeSelect.vue' import { Plus, ChevronDown, X } from 'lucide-vue-next' -import { api, type Task, type EnvVar } from '@/api' +import { api, type Task, type EnvVar, type Agent } from '@/api' import { toast } from 'vue-sonner' const props = defineProps<{ @@ -39,7 +39,9 @@ const form = ref>({}) const cleanType = ref('none') const cleanKeep = ref(30) const allEnvVars = ref([]) +const allAgents = ref([]) const selectedEnvIds = ref([]) +const selectedAgentId = ref('local') const envSearchQuery = ref('') const cleanConfig = computed(() => { @@ -61,6 +63,10 @@ const selectedEnvs = computed(() => { .filter((e): e is EnvVar => e !== undefined) }) +const onlineAgents = computed(() => { + return allAgents.value.filter(a => a.enabled) +}) + watch(() => props.open, async (val) => { if (val) { form.value = { ...props.task } @@ -84,14 +90,25 @@ watch(() => props.open, async (val) => { } else { selectedEnvIds.value = [] } + // 解析 Agent + selectedAgentId.value = props.task?.agent_id ? String(props.task.agent_id) : 'local' envSearchQuery.value = '' - // 加载环境变量 - try { - allEnvVars.value = await api.env.all() - } catch { /* ignore */ } + // 加载数据 + await loadData() } }) +async function loadData() { + try { + const [envs, agents] = await Promise.all([ + api.env.all(), + api.agents.list() + ]) + allEnvVars.value = envs + allAgents.value = agents + } catch { /* ignore */ } +} + function addEnv(id: number) { if (!selectedEnvIds.value.includes(id)) { selectedEnvIds.value.push(id) @@ -109,6 +126,7 @@ async function save() { form.value.clean_config = cleanConfig.value form.value.envs = selectedEnvIds.value.join(',') form.value.type = 'task' + form.value.agent_id = selectedAgentId.value === 'local' ? null : Number(selectedAgentId.value) if (props.isEdit && form.value.id) { await api.tasks.update(form.value.id, form.value) toast.success('任务已更新') @@ -140,7 +158,24 @@ async function save() {
- + + +
+
+
+ +
+
diff --git a/web/src/views/tasks/Tasks.vue b/web/src/views/tasks/Tasks.vue index 7338526..3da7e5d 100644 --- a/web/src/views/tasks/Tasks.vue +++ b/web/src/views/tasks/Tasks.vue @@ -1,13 +1,13 @@