diff --git a/.gitignore b/.gitignore index 3e1789e..5efbfe4 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,9 @@ baihu # Data & Logs data/ +!data/agent/ +data/agent/* +!data/agent/version.txt envs/ .kiro/ # logs/ diff --git a/Dockerfile b/Dockerfile index 62918bd..c1e188d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,7 +45,34 @@ RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \ -o baihu . # ================================ -# Stage 3: Final image +# Stage 3: Build Agent (all platforms) +# ================================ +FROM --platform=$BUILDPLATFORM golang:1.24-alpine AS agent-builder + +ARG VERSION=dev +ARG BUILD_TIME + +WORKDIR /app + +# Copy agent source +COPY agent/ ./agent/ + +# Download dependencies +WORKDIR /app/agent +RUN go env -w GOPROXY=https://goproxy.cn,direct && go mod download + +# Build agent for all platforms (parallel) +RUN mkdir -p /opt/agent && \ + echo "${VERSION}" > /opt/agent/version.txt && \ + CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.BuildTime=${BUILD_TIME}'" -o /opt/agent/baihu-agent-linux-amd64 . & \ + # CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.BuildTime=${BUILD_TIME}'" -o /opt/agent/baihu-agent-linux-arm64 . & \ + # CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.BuildTime=${BUILD_TIME}'" -o /opt/agent/baihu-agent-windows-amd64.exe . & \ + # CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.BuildTime=${BUILD_TIME}'" -o /opt/agent/baihu-agent-darwin-amd64 . & \ + # CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 go build -ldflags="-s -w -X 'main.Version=${VERSION}' -X 'main.BuildTime=${BUILD_TIME}'" -o /opt/agent/baihu-agent-darwin-arm64 . & \ + wait + +# ================================ +# Stage 4: Final image # ================================ FROM debian:bookworm-slim @@ -77,6 +104,10 @@ COPY docker-entrypoint.sh . # Copy sync.py to /opt COPY custom/sync.py /opt/sync.py + +# Copy agent binaries to /opt/agent +COPY --from=agent-builder /opt/agent /opt/agent + RUN chmod +x /opt/sync.py \ && chmod +x docker-entrypoint.sh \ && touch "dont-not-delete-anythings" \ diff --git a/Makefile b/Makefile index 513ca40..361f3c4 100644 --- a/Makefile +++ b/Makefile @@ -24,6 +24,18 @@ build: # Build all (frontend + backend) build-all: build-web build +# Build agent for all platforms (local development) +build-agent: + @mkdir -p data/agent + @echo "$(VERSION)" > data/agent/version.txt + cd agent && CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=$(VERSION)' -X 'main.BuildTime=$(BUILD_TIME)'" -o ../data/agent/baihu-agent-linux-amd64 . + cd agent && CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w -X 'main.Version=$(VERSION)' -X 'main.BuildTime=$(BUILD_TIME)'" -o ../data/agent/baihu-agent-linux-arm64 . + cd agent && CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=$(VERSION)' -X 'main.BuildTime=$(BUILD_TIME)'" -o ../data/agent/baihu-agent-windows-amd64.exe . + cd agent && CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -ldflags="-s -w -X 'main.Version=$(VERSION)' -X 'main.BuildTime=$(BUILD_TIME)'" -o ../data/agent/baihu-agent-darwin-amd64 . + cd agent && CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 go build -ldflags="-s -w -X 'main.Version=$(VERSION)' -X 'main.BuildTime=$(BUILD_TIME)'" -o ../data/agent/baihu-agent-darwin-arm64 . + @echo "Agent binaries built in data/agent/ (for local dev)" + @echo "In Docker, agents are built to /opt/agent/" + # Clean built files clean: $(GOCLEAN) @@ -67,6 +79,7 @@ help: @echo "Available targets:" @echo " all - Build the application (default)" @echo " build - Build the application" + @echo " build-agent - Build agent for all platforms" @echo " clean - Clean built files" @echo " run - Run the application" @echo " deps - Install dependencies" diff --git a/agent/config.example.ini b/agent/config.example.ini new file mode 100644 index 0000000..ce2a1cf --- /dev/null +++ b/agent/config.example.ini @@ -0,0 +1,11 @@ +[agent] +# 主服务器地址 +server_url = http://192.168.1.100:8052 +# Agent 名称（留空则使用主机名） +name = agent-01 +# Token（由服务器下发，首次运行留空） +token = +# 心跳间隔（秒） +interval = 30 +# 自动更新（true/false） +auto_update = true diff --git a/agent/go.mod b/agent/go.mod new file mode 100644 index 0000000..20bf62d --- /dev/null +++ b/agent/go.mod @@ -0,0 +1,12 @@ +module baihu-agent + +go 1.24 + +require ( + github.com/robfig/cron/v3 v3.0.1 + github.com/sirupsen/logrus v1.9.3 + gopkg.in/ini.v1 v1.67.0 + gopkg.in/natefinch/lumberjack.v2 v2.2.1 +) + +require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect diff --git a/agent/go.sum b/agent/go.sum new file mode 100644 index 0000000..737c46b --- /dev/null +++ b/agent/go.sum @@ -0,0 +1,21 @@ +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/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= +github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/ini.v1 v1.67.0 h1:Dgnx+6+nfE+IfzjUEISNeydPJh9AXNNsWbGP9KzCsOA= +gopkg.in/ini.v1 v1.67.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= +gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= +gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/agent/main.go b/agent/main.go new file mode 100644 index 0000000..80d0076 --- /dev/null +++ b/agent/main.go @@ -0,0 +1,976 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "os/exec" + "os/signal" + "path/filepath" + "runtime" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "github.com/robfig/cron/v3" + "github.com/sirupsen/logrus" + "gopkg.in/ini.v1" + "gopkg.in/natefinch/lumberjack.v2" +) + +const ServiceName = "baihu-agent" +const ServiceDesc = "Baihu Agent Service" + +// 版本信息（通过 ldflags 注入） +var ( + Version = "dev" + BuildTime = "" +) + +// 东八区时区 +var cstZone = time.FixedZone("CST", 8*3600) + +// 日志实例 +var log = logrus.New() + +// 全局配置 +var ( + configFile = "config.ini" + logFile = "logs/agent.log" +) + +func main() { + // 获取程序所在目录 + exePath, _ := os.Executable() + exeDir := filepath.Dir(exePath) + os.Chdir(exeDir) + + if len(os.Args) < 2 { + printUsage() + os.Exit(1) + } + + cmd := os.Args[1] + + // 解析额外参数 + for i := 2; i < len(os.Args); i++ { + switch os.Args[i] { + case "-c", "--config": + if i+1 < len(os.Args) { + configFile = os.Args[i+1] + i++ + } + case "-l", "--log": + if i+1 < len(os.Args) { + logFile = os.Args[i+1] + i++ + } + } + } + + switch cmd { + case "start": + cmdStart() + case "stop": + cmdStop() + case "status": + cmdStatus() + case "install": + cmdInstall() + case "uninstall": + cmdUninstall() + case "version", "-v", "--version": + fmt.Printf("Baihu Agent v%s\n", Version) + if BuildTime != "" { + fmt.Printf("Build Time: %s\n", BuildTime) + } + case "help", "-h", "--help": + printUsage() + default: + fmt.Printf("未知命令: %s\n", cmd) + printUsage() + os.Exit(1) + } +} + +func printUsage() { + fmt.Printf(`Baihu Agent v%s + +用法: baihu-agent <命令> [选项] + +命令: + start 启动 Agent + stop 停止 Agent + status 查看运行状态 + install 安装为系统服务（开机自启） + uninstall 卸载系统服务 + version 显示版本信息 + help 显示帮助信息 + +选项: + -c, --config 配置文件路径 (默认: config.ini) + -l, --log 日志文件路径 (默认: logs/agent.log) + +示例: + baihu-agent start + baihu-agent start -c /etc/baihu/config.ini + baihu-agent install + baihu-agent status +`, Version) +} + +// ========== 命令实现 ========== + +func cmdStart() { + // 初始化日志 + initLogger(logFile) + + // 加载配置 + config := &Config{Interval: 30} + if err := loadConfigFile(configFile, config); err != nil { + if !os.IsNotExist(err) { + log.Warnf("加载配置文件失败: %v", err) + } + } + + // 从环境变量加载 + if v := os.Getenv("AGENT_SERVER"); v != "" { + config.ServerURL = v + } + if v := os.Getenv("AGENT_NAME"); v != "" { + config.Name = v + } + + // 验证配置 + if config.ServerURL == "" { + log.Fatal("请在配置文件中设置 server_url") + } + if config.Name == "" { + hostname, _ := os.Hostname() + config.Name = hostname + } + + log.Infof("Baihu Agent v%s", Version) + if BuildTime != "" { + log.Infof("构建时间: %s", BuildTime) + } + log.Infof("服务器: %s", config.ServerURL) + log.Infof("名称: %s", config.Name) + + // 写入 PID 文件 + writePidFile() + + // 创建并启动 Agent + agent := NewAgent(config, configFile) + if err := agent.Start(); err != nil { + log.Fatalf("启动失败: %v", err) + } + + // 等待退出信号 + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + + log.Info("正在停止...") + agent.Stop() + removePidFile() +} + +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("服务已卸载") +} + +// ========== PID 文件管理 ========== + +func getPidFile() string { + return filepath.Join(filepath.Dir(configFile), "agent.pid") +} + +func writePidFile() { + 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()) +} + +// ========== 日志初始化 ========== + +// 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)) +} + + +// ========== 配置相关 ========== + +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 结构 ========== + +type AgentTask struct { + ID uint `json:"id"` + Name string `json:"name"` + Command string `json:"command"` + Schedule string `json:"schedule"` + 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 + cron *cron.Cron + tasks map[uint]*AgentTask + entryMap map[uint]cron.EntryID + mu sync.RWMutex + client *http.Client +} + +func NewAgent(config *Config, configFile string) *Agent { + return &Agent{ + config: config, + configFile: configFile, + 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}, + } +} + +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) + } + + a.cron.Start() + go a.heartbeatLoop() + go a.syncTasksLoop() + + log.Info("Agent 已启动 (时区: Asia/Shanghai)") + return nil +} + +func (a *Agent) Stop() { + 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() + + for { + <-ticker.C + + statusResp, err := a.doRequestNoAuth("POST", "/api/agent/status", map[string]string{ + "name": a.config.Name, + }) + if err != nil { + log.Warnf("检查状态失败: %v", err) + continue + } + + if statusResp.StatusCode != http.StatusOK { + statusResp.Body.Close() + continue + } + + 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("等待审核中...") + } +} + +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) + } + } +} + +func (a *Agent) heartbeat() error { + hostname, _ := os.Hostname() + body := map[string]interface{}{ + "version": Version, + "build_time": BuildTime, + "hostname": hostname, + "os": runtime.GOOS, + "arch": runtime.GOARCH, + "auto_update": a.config.AutoUpdate, + } + + resp, err := a.doRequest("POST", "/api/agent/heartbeat", body) + 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 +} + +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) 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 + } + + 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)) + } + + log.Infof("任务 #%d 执行完成 (%s)", result.TaskID, result.Status) + return nil +} + +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) +} + +func (a *Agent) doRequestNoAuth(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("Content-Type", "application/json") + + return a.client.Do(req) +} + +// selfUpdate 自动更新 +func (a *Agent) selfUpdate() { + // 获取当前可执行文件路径 + exePath, err := os.Executable() + if err != nil { + log.Errorf("获取可执行文件路径失败: %v", err) + return + } + + // 下载新版本 + downloadURL := fmt.Sprintf("%s/api/agent/download?os=%s&arch=%s", a.config.ServerURL, runtime.GOOS, 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 + } + + // 保存到临时文件 + tmpFile := exePath + ".new" + f, err := os.OpenFile(tmpFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755) + if err != nil { + log.Errorf("创建临时文件失败: %v", err) + return + } + + _, err = io.Copy(f, resp.Body) + f.Close() + if err != nil { + log.Errorf("保存新版本失败: %v", err) + os.Remove(tmpFile) + return + } + + // 备份旧版本 + backupFile := exePath + ".bak" + os.Remove(backupFile) + if err := os.Rename(exePath, backupFile); err != nil { + log.Errorf("备份旧版本失败: %v", err) + os.Remove(tmpFile) + return + } + + // 替换为新版本 + if err := os.Rename(tmpFile, exePath); err != nil { + log.Errorf("替换新版本失败: %v", err) + os.Rename(backupFile, exePath) // 恢复旧版本 + return + } + + log.Info("更新完成，正在重启...") + + // 重启服务 + a.restart() +} + +// restart 重启服务 +func (a *Agent) restart() { + exePath, _ := os.Executable() + + if runtime.GOOS == "windows" { + // Windows: 启动新进程后退出 + cmd := exec.Command(exePath, "start") + cmd.Start() + os.Exit(0) + } else { + // Linux/macOS: 使用 exec 替换当前进程 + syscall.Exec(exePath, []string{exePath, "start"}, os.Environ()) + } +} diff --git a/internal/controllers/agent_controller.go b/internal/controllers/agent_controller.go new file mode 100644 index 0000000..8e9306d --- /dev/null +++ b/internal/controllers/agent_controller.go @@ -0,0 +1,350 @@ +package controllers + +import ( + "baihu/internal/models" + "baihu/internal/services" + "baihu/internal/utils" + "strconv" + "strings" + + "github.com/gin-gonic/gin" +) + +// AgentController Agent 控制器 +type AgentController struct { + agentService *services.AgentService +} + +// NewAgentController 创建 Agent 控制器 +func NewAgentController() *AgentController { + return &AgentController{ + agentService: services.NewAgentService(), + } +} + +// 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) + if err != nil { + utils.BadRequest(ctx, "无效的 ID") + return + } + + var req struct { + Name string `json:"name" binding:"required"` + Description string `json:"description"` + Enabled bool `json:"enabled"` + } + + if err := ctx.ShouldBindJSON(&req); err != nil { + utils.BadRequest(ctx, "参数错误") + return + } + + if err := c.agentService.Update(uint(id), req.Name, req.Description, req.Enabled); err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.SuccessMsg(ctx, "更新成功") +} + +// Delete 删除 Agent +func (c *AgentController) Delete(ctx *gin.Context) { + id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) + if err != nil { + utils.BadRequest(ctx, "无效的 ID") + return + } + + if err := c.agentService.Delete(uint(id)); err != nil { + utils.BadRequest(ctx, err.Error()) + return + } + + utils.SuccessMsg(ctx, "删除成功") +} + +// RegenerateToken 重新生成 Token +func (c *AgentController) RegenerateToken(ctx *gin.Context) { + id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) + if err != nil { + utils.BadRequest(ctx, "无效的 ID") + return + } + + token, err := c.agentService.RegenerateToken(uint(id)) + if err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.Success(ctx, gin.H{"token": token}) +} + +// ========== Agent API（供 Agent 调用）========== + +// Register Agent 注册（无需认证） +func (c *AgentController) Register(ctx *gin.Context) { + var req models.AgentRegisterRequest + if err := ctx.ShouldBindJSON(&req); err != nil { + utils.BadRequest(ctx, "参数错误") + return + } + + if req.Name == "" { + utils.BadRequest(ctx, "名称不能为空") + return + } + + ip := ctx.ClientIP() + agent, err := c.agentService.Register(&req, ip) + if err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.Success(ctx, gin.H{ + "agent_id": agent.ID, + "status": agent.Status, + "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) + if token == "" { + utils.Unauthorized(ctx, "缺少认证 Token") + return + } + + 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"` + } + ctx.ShouldBindJSON(&req) + + ip := ctx.ClientIP() + agent, err := c.agentService.Heartbeat(token, ip, req.Version, req.BuildTime, req.Hostname, req.OS, req.Arch) + if err != nil { + utils.Unauthorized(ctx, err.Error()) + return + } + + // 检查是否需要更新 + latestVersion := c.agentService.GetLatestVersion() + needUpdate := latestVersion != "" && req.Version != "" && req.Version != latestVersion + forceUpdate := agent.ForceUpdate + + // 如果强制更新已触发，重置标志 + if forceUpdate && needUpdate { + c.agentService.ClearForceUpdate(agent.ID) + } + + utils.Success(ctx, gin.H{ + "agent_id": agent.ID, + "name": agent.Name, + "need_update": needUpdate, + "force_update": forceUpdate, + "latest_version": latestVersion, + }) +} + +// GetTasks Agent 获取任务列表 +func (c *AgentController) GetTasks(ctx *gin.Context) { + token := c.getAgentToken(ctx) + if token == "" { + utils.Unauthorized(ctx, "缺少认证 Token") + return + } + + agent := c.agentService.GetByToken(token) + if agent == nil { + utils.Unauthorized(ctx, "无效的 Token") + return + } + + if !agent.Enabled { + utils.Forbidden(ctx, "Agent 已禁用") + return + } + + tasks := c.agentService.GetTasks(agent.ID) + utils.Success(ctx, gin.H{ + "agent_id": agent.ID, + "tasks": tasks, + }) +} + +// ReportResult Agent 上报执行结果 +func (c *AgentController) ReportResult(ctx *gin.Context) { + token := c.getAgentToken(ctx) + if token == "" { + utils.Unauthorized(ctx, "缺少认证 Token") + return + } + + agent := c.agentService.GetByToken(token) + if agent == nil { + utils.Unauthorized(ctx, "无效的 Token") + return + } + + if !agent.Enabled { + utils.Forbidden(ctx, "Agent 已禁用") + return + } + + var result models.AgentTaskResult + if err := ctx.ShouldBindJSON(&result); err != nil { + utils.BadRequest(ctx, "参数错误") + return + } + + result.AgentID = agent.ID + + if err := c.agentService.ReportResult(&result); err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.SuccessMsg(ctx, "上报成功") +} + +// getAgentToken 从请求头获取 Agent Token +func (c *AgentController) getAgentToken(ctx *gin.Context) string { + auth := ctx.GetHeader("Authorization") + if auth == "" { + return "" + } + // Bearer + parts := strings.SplitN(auth, " ", 2) + if len(parts) == 2 && parts[0] == "Bearer" { + return parts[1] + } + return auth +} + +// Download 下载 Agent 程序 +func (c *AgentController) Download(ctx *gin.Context) { + osType := ctx.DefaultQuery("os", "linux") + arch := ctx.DefaultQuery("arch", "amd64") + + data, filename, err := c.agentService.GetAgentBinary(osType, arch) + if err != nil { + utils.NotFound(ctx, err.Error()) + return + } + + ctx.Header("Content-Disposition", "attachment; filename="+filename) + ctx.Header("Content-Type", "application/octet-stream") + ctx.Header("Content-Length", strconv.Itoa(len(data))) + ctx.Data(200, "application/octet-stream", data) +} + +// GetVersion 获取 Agent 最新版本信息 +func (c *AgentController) GetVersion(ctx *gin.Context) { + version := c.agentService.GetLatestVersion() + platforms := c.agentService.GetAvailablePlatforms() + + utils.Success(ctx, gin.H{ + "version": version, + "platforms": platforms, + }) +} + +// ForceUpdate 强制更新指定 Agent +func (c *AgentController) ForceUpdate(ctx *gin.Context) { + id, err := strconv.ParseUint(ctx.Param("id"), 10, 32) + if err != nil { + utils.BadRequest(ctx, "无效的 ID") + return + } + + if err := c.agentService.SetForceUpdate(uint(id)); err != nil { + utils.ServerError(ctx, err.Error()) + return + } + + utils.SuccessMsg(ctx, "已标记强制更新，Agent 下次心跳时将自动更新") +} diff --git a/internal/database/migrate.go b/internal/database/migrate.go index 41e351f..9ee9da7 100644 --- a/internal/database/migrate.go +++ b/internal/database/migrate.go @@ -15,5 +15,6 @@ func Migrate() error { &models.LoginLog{}, &models.SendStats{}, &models.Dependency{}, + &models.Agent{}, ) } diff --git a/internal/models/agent.go b/internal/models/agent.go new file mode 100644 index 0000000..b1ddcb6 --- /dev/null +++ b/internal/models/agent.go @@ -0,0 +1,64 @@ +package models + +import ( + "baihu/internal/constant" + + "gorm.io/gorm" +) + +// Agent 远程执行代理 +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 + Description string `json:"description" gorm:"size:255"` // 描述 + Status string `json:"status" gorm:"size:20;default:'pending'"` // 状态: pending(待审核), online, offline + LastSeen *LocalTime `json:"last_seen"` // 最后心跳时间 + IP string `json:"ip" gorm:"size:45"` // Agent IP 地址 + Version string `json:"version" gorm:"size:20"` // Agent 版本 + BuildTime string `json:"build_time" gorm:"size:30"` // Agent 构建时间 + Hostname string `json:"hostname" gorm:"size:100"` // Agent 主机名 + OS string `json:"os" gorm:"size:20"` // 操作系统 + Arch string `json:"arch" gorm:"size:20"` // 架构 + ForceUpdate bool `json:"force_update" gorm:"default:false"` // 强制更新标志 + Enabled bool `json:"enabled" gorm:"default:true"` // 是否启用 + CreatedAt LocalTime `json:"created_at"` + UpdatedAt LocalTime `json:"updated_at"` + DeletedAt gorm.DeletedAt `json:"-" gorm:"index"` +} + +func (Agent) TableName() string { + return constant.TablePrefix + "agents" +} + +// AgentTask Agent 任务配置（用于下发给 Agent） +type AgentTask struct { + ID uint `json:"id"` + Name string `json:"name"` + Command string `json:"command"` + Schedule string `json:"schedule"` + Timeout int `json:"timeout"` + WorkDir string `json:"work_dir"` + Envs string `json:"envs"` + Enabled bool `json:"enabled"` +} + +// AgentTaskResult Agent 上报的任务执行结果 +type AgentTaskResult struct { + TaskID uint `json:"task_id"` + AgentID uint `json:"agent_id"` + Command string `json:"command"` + Output string `json:"output"` + Status string `json:"status"` // success, failed + Duration int64 `json:"duration"` // milliseconds + ExitCode int `json:"exit_code"` + StartTime int64 `json:"start_time"` // unix timestamp + EndTime int64 `json:"end_time"` // unix timestamp +} + +// AgentRegisterRequest Agent 注册请求 +type AgentRegisterRequest struct { + Name string `json:"name"` + Hostname string `json:"hostname"` + Version string `json:"version"` +} diff --git a/internal/models/task.go b/internal/models/task.go index 5485493..0ad5393 100644 --- a/internal/models/task.go +++ b/internal/models/task.go @@ -37,6 +37,7 @@ type Task struct { WorkDir string `json:"work_dir" gorm:"size:255;default:''"` // 工作目录，为空则使用 scripts 目录 CleanConfig string `json:"clean_config" gorm:"size:255;default:''"` // 清理配置 JSON Envs string `json:"envs" gorm:"size:255;default:''"` // 环境变量ID列表，逗号分隔 + AgentID *uint `json:"agent_id" gorm:"index"` // Agent ID，为空表示本地执行 Enabled bool `json:"enabled" gorm:"default:true"` LastRun *LocalTime `json:"last_run"` NextRun *LocalTime `json:"next_run"` @@ -53,6 +54,7 @@ func (Task) TableName() string { type TaskLog struct { ID uint `json:"id" gorm:"primaryKey"` TaskID uint `json:"task_id" gorm:"index"` + AgentID *uint `json:"agent_id" gorm:"index"` // Agent ID，为空表示本地执行 Command string `json:"command" gorm:"type:text"` Output string `json:"-" gorm:"type:longtext"` // gzip+base64 compressed Status string `json:"status" gorm:"size:20"` // success, failed diff --git a/internal/router/register.go b/internal/router/register.go index 3a0c0a7..9f1f798 100644 --- a/internal/router/register.go +++ b/internal/router/register.go @@ -39,6 +39,7 @@ func RegisterControllers() *Controllers { Terminal: controllers.NewTerminalController(), Settings: controllers.NewSettingsController(userService, loginLogService, executorService), Dependency: controllers.NewDependencyController(), + Agent: controllers.NewAgentController(), } } diff --git a/internal/router/router.go b/internal/router/router.go index fa47d35..6c6c9dd 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -23,6 +23,7 @@ type Controllers struct { Terminal *controllers.TerminalController Settings *controllers.SettingsController Dependency *controllers.DependencyController + Agent *controllers.AgentController } func mustSubFS(fsys fs.FS, dir string) fs.FS { @@ -197,6 +198,31 @@ func Setup(c *Controllers) *gin.Engine { deps.POST("/reinstall-all", c.Dependency.ReinstallAll) deps.GET("/installed", c.Dependency.GetInstalled) } + + // Agent routes (Agent 管理) + 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) + } + } + + // 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) } } diff --git a/internal/services/agent_service.go b/internal/services/agent_service.go new file mode 100644 index 0000000..90abb9e --- /dev/null +++ b/internal/services/agent_service.go @@ -0,0 +1,378 @@ +package services + +import ( + "baihu/internal/database" + "baihu/internal/logger" + "baihu/internal/models" + "baihu/internal/utils" + "crypto/rand" + "encoding/hex" + "fmt" + "os" + "path/filepath" + "strings" + "time" +) + +// AgentService Agent 服务 +type AgentService struct{} + +// NewAgentService 创建 Agent 服务 +func NewAgentService() *AgentService { + return &AgentService{} +} + +// generateToken 生成随机 Token +func generateToken() string { + bytes := make([]byte, 32) + rand.Read(bytes) + 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 + } + + // 创建新的待审核 Agent + 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, + } + + if err := database.DB.Create(agent).Error; err != nil { + return nil, err + } + + logger.Infof("[Agent] 新 Agent 注册: %s (%s)", req.Name, ip) + return agent, 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 不存在"} + } + + if agent.Status != "pending" { + return nil, &ServiceError{Message: "Agent 状态不是待审核"} + } + + token := generateToken() + 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.Token = token + agent.Status = "online" + agent.LastSeen = &now + + 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 +} + +// Update 更新 Agent +func (s *AgentService) Update(id uint, name, description string, enabled bool) error { + return database.DB.Model(&models.Agent{}).Where("id = ?", id).Updates(map[string]interface{}{ + "name": name, + "description": description, + "enabled": enabled, + }).Error +} + +// Delete 删除 Agent +func (s *AgentService) Delete(id uint) error { + // 检查是否有关联任务 + var count int64 + database.DB.Model(&models.Task{}).Where("agent_id = ?", id).Count(&count) + if count > 0 { + return &ServiceError{Message: "该 Agent 下还有关联任务，无法删除"} + } + + return database.DB.Delete(&models.Agent{}, id).Error +} + +// GetByID 根据 ID 获取 Agent +func (s *AgentService) GetByID(id uint) *models.Agent { + var agent models.Agent + if err := database.DB.First(&agent, id).Error; err != nil { + return nil + } + return &agent +} + +// GetByToken 根据 Token 获取 Agent +func (s *AgentService) GetByToken(token string) *models.Agent { + var agent models.Agent + if err := database.DB.Where("token = ?", token).First(&agent).Error; err != nil { + return nil + } + return &agent +} + +// List 获取已审核的 Agent 列表 +func (s *AgentService) List() []models.Agent { + var agents []models.Agent + database.DB.Where("status != ?", "pending").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 +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 +} + +// Heartbeat Agent 心跳 +func (s *AgentService) Heartbeat(token, ip, version, buildTime, hostname, osType, arch string) (*models.Agent, error) { + agent := s.GetByToken(token) + if agent == nil { + return nil, &ServiceError{Message: "无效的 Token"} + } + + if !agent.Enabled { + return nil, &ServiceError{Message: "Agent 已禁用"} + } + + now := models.LocalTime(time.Now()) + updates := map[string]interface{}{ + "status": "online", + "last_seen": now, + "ip": ip, + } + if version != "" { + updates["version"] = version + } + if buildTime != "" { + updates["build_time"] = buildTime + } + if hostname != "" { + updates["hostname"] = hostname + } + if osType != "" { + updates["os"] = osType + } + if arch != "" { + updates["arch"] = arch + } + + database.DB.Model(&models.Agent{}).Where("id = ?", agent.ID).Updates(updates) + + agent.Status = "online" + agent.LastSeen = &now + agent.IP = ip + agent.Version = version + agent.BuildTime = buildTime + agent.Hostname = hostname + agent.OS = osType + agent.Arch = arch + + 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 + database.DB.Where("agent_id = ? AND enabled = ?", agentID, true).Find(&tasks) + + result := make([]models.AgentTask, len(tasks)) + for i, task := range tasks { + result[i] = models.AgentTask{ + ID: task.ID, + Name: task.Name, + Command: task.Command, + Schedule: task.Schedule, + Timeout: task.Timeout, + WorkDir: task.WorkDir, + Envs: task.Envs, + Enabled: task.Enabled, + } + } + + return result +} + +// ReportResult Agent 上报执行结果 +func (s *AgentService) ReportResult(result *models.AgentTaskResult) error { + // 压缩输出 + compressed, err := utils.CompressToBase64(result.Output) + if err != nil { + logger.Errorf("[Agent] 压缩日志失败: %v", err) + compressed = "" + } + + taskLog := &models.TaskLog{ + TaskID: result.TaskID, + AgentID: &result.AgentID, + Command: result.Command, + Output: compressed, + Status: result.Status, + Duration: result.Duration, + ExitCode: result.ExitCode, + } + + if err := database.DB.Create(taskLog).Error; err != nil { + return err + } + + // 更新任务的 last_run + database.DB.Model(&models.Task{}).Where("id = ?", result.TaskID).Update("last_run", time.Now()) + + // 更新统计 + sendStatsService := NewSendStatsService() + sendStatsService.IncrementStats(result.TaskID, result.Status) + + logger.Infof("[Agent] 收到任务结果 #%d (agent=%d, status=%s)", result.TaskID, result.AgentID, result.Status) + return nil +} + +// UpdateOfflineAgents 更新离线 Agent 状态（超过 2 分钟无心跳） +func (s *AgentService) UpdateOfflineAgents() { + cutoff := time.Now().Add(-2 * time.Minute) + database.DB.Model(&models.Agent{}). + Where("status = ? AND last_seen < ?", "online", cutoff). + Update("status", "offline") +} + +// GetLatestVersion 获取最新 Agent 版本 +func (s *AgentService) GetLatestVersion() string { + // 优先从 /opt/agent 读取（容器内） + versionFile := "/opt/agent/version.txt" + data, err := os.ReadFile(versionFile) + if err != nil { + // 回退到 data/agent（本地开发） + data, err = os.ReadFile("data/agent/version.txt") + if err != nil { + return "" + } + } + return strings.TrimSpace(string(data)) +} + +// GetAvailablePlatforms 获取可用的平台列表 +func (s *AgentService) GetAvailablePlatforms() []map[string]string { + platforms := []map[string]string{} + + // 优先从 /opt/agent 读取（容器内） + agentDir := "/opt/agent" + files, err := os.ReadDir(agentDir) + if err != nil { + // 回退到 data/agent（本地开发） + agentDir = "data/agent" + files, err = os.ReadDir(agentDir) + if err != nil { + return platforms + } + } + + for _, f := range files { + name := f.Name() + if strings.HasPrefix(name, "baihu-agent-") { + // baihu-agent-linux-amd64, baihu-agent-windows-amd64.exe + parts := strings.Split(strings.TrimSuffix(name, ".exe"), "-") + if len(parts) >= 4 { + platforms = append(platforms, map[string]string{ + "os": parts[2], + "arch": parts[3], + "filename": name, + }) + } + } + } + + return platforms +} + +// GetAgentBinary 获取 Agent 二进制文件 +func (s *AgentService) GetAgentBinary(osType, arch string) ([]byte, string, error) { + filename := fmt.Sprintf("baihu-agent-%s-%s", osType, arch) + if osType == "windows" { + filename += ".exe" + } + + // 优先从 /opt/agent 读取（容器内） + filePath := filepath.Join("/opt/agent", filename) + data, err := os.ReadFile(filePath) + if err != nil { + // 回退到 data/agent（本地开发） + filePath = filepath.Join("data/agent", filename) + data, err = os.ReadFile(filePath) + if err != nil { + return nil, "", &ServiceError{Message: "未找到对应平台的 Agent 程序"} + } + } + + return data, filename, nil +} + +// SetForceUpdate 设置强制更新标志 +func (s *AgentService) SetForceUpdate(id uint) error { + return database.DB.Model(&models.Agent{}).Where("id = ?", id).Update("force_update", true).Error +} + +// ClearForceUpdate 清除强制更新标志 +func (s *AgentService) ClearForceUpdate(id uint) error { + return database.DB.Model(&models.Agent{}).Where("id = ?", id).Update("force_update", false).Error +} + +// ServiceError 服务错误 +type ServiceError struct { + Message string +} + +func (e *ServiceError) Error() string { + return e.Message +} diff --git a/internal/services/cron_service.go b/internal/services/cron_service.go index 812334a..ce2024d 100644 --- a/internal/services/cron_service.go +++ b/internal/services/cron_service.go @@ -11,6 +11,9 @@ import ( "github.com/robfig/cron/v3" ) +// 东八区时区 +var cstZone = time.FixedZone("CST", 8*3600) + // CronService manages scheduled tasks using robfig/cron type CronService struct { cron *cron.Cron @@ -22,8 +25,8 @@ type CronService struct { // NewCronService creates a new cron service func NewCronService(taskService *TaskService, executorService *ExecutorService) *CronService { - // 使用秒级精度的 cron parser，支持 6 位表达式（秒 分 时 日 月 周） - c := cron.New(cron.WithSeconds()) + // 使用秒级精度的 cron parser，支持 6 位表达式（秒 分 时 日 月 周），使用东八区时区 + c := cron.New(cron.WithSeconds(), cron.WithLocation(cstZone)) return &CronService{ cron: c, @@ -52,7 +55,8 @@ func (cs *CronService) loadTasks() { tasks := cs.taskService.GetTasks() count := 0 for _, task := range tasks { - if task.Enabled { + // 只调度本地任务（agent_id 为空） + if task.Enabled && task.AgentID == nil { err := cs.addTask(&task, false) if err != nil { return