feat: add agent exec
This commit is contained in:
engigu
@@ -10,6 +10,9 @@ baihu
|
||||
|
||||
# Data & Logs
|
||||
data/
|
||||
!data/agent/
|
||||
data/agent/*
|
||||
!data/agent/version.txt
|
||||
envs/
|
||||
.kiro/
|
||||
# logs/
|
||||
|
||||
+32
-1
@@ -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" \
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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=
|
||||
+976
@@ -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 <file> 配置文件路径 (默认: config.ini)
|
||||
-l, --log <file> 日志文件路径 (默认: 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())
|
||||
}
|
||||
}
|
||||
@@ -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 <token>
|
||||
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 下次心跳时将自动更新")
|
||||
}
|
||||
@@ -15,5 +15,6 @@ func Migrate() error {
|
||||
&models.LoginLog{},
|
||||
&models.SendStats{},
|
||||
&models.Dependency{},
|
||||
&models.Agent{},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -39,6 +39,7 @@ func RegisterControllers() *Controllers {
|
||||
Terminal: controllers.NewTerminalController(),
|
||||
Settings: controllers.NewSettingsController(userService, loginLogService, executorService),
|
||||
Dependency: controllers.NewDependencyController(),
|
||||
Agent: controllers.NewAgentController(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user