aiops/backend/internal/api/handlers/agent_install.go

package handlers

import (
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"

	"ezops/internal/models"

	"github.com/gin-gonic/gin"
	"golang.org/x/crypto/ssh"
	"gorm.io/gorm"
)

// AgentInstallHandler Agent 远程安装处理器 // Logan.Li
type AgentInstallHandler struct {
	db *gorm.DB
}

func NewAgentInstallHandler(db *gorm.DB) *AgentInstallHandler {
	return &AgentInstallHandler{db: db}
}

// sshDial 复用 webshell.go 的 SSH 连接逻辑
func sshDial(server models.Server) (*ssh.Client, error) {
	config := &ssh.ClientConfig{
		User:            server.Username,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}

	switch server.AuthType {
	case "key":
		signer, err := ssh.ParsePrivateKey([]byte(server.SSHKey))
		if err != nil {
			config.Auth = []ssh.AuthMethod{ssh.Password(server.Password)}
		} else {
			config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
		}
	default:
		config.Auth = []ssh.AuthMethod{ssh.Password(server.Password)}
	}

	addr := server.IP + ":" + strconv.Itoa(server.Port)
	return ssh.Dial("tcp", addr, config)
}

// runSSHCommand 在远程主机执行命令并返回输出
func runSSHCommand(client *ssh.Client, cmd string) (string, error) {
	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	return string(output), err
}

// InstallAgent 远程安装 Agent 到指定服务器
// POST /servers/:id/install-agent
func (h *AgentInstallHandler) InstallAgent(c *gin.Context) {
	id, _ := strconv.Atoi(c.Param("id"))

	var server models.Server
	if err := h.db.First(&server, id).Error; err != nil {
		c.JSON(http.StatusNotFound, models.Response{Code: 404, Message: "服务器不存在"})
		return
	}

	if server.Status != "running" {
		c.JSON(http.StatusBadRequest, models.Response{Code: 400, Message: "服务器未运行，无法安装 Agent"})
		return
	}

	// 1. SSH 连接目标主机
	client, err := sshDial(server)
	if err != nil {
		log.Printf("[AgentInstall] SSH连接失败 %s: %v", server.IP, err)
		c.JSON(http.StatusOK, models.Response{
			Code:    500,
			Message: "SSH 连接失败: " + formatSSHError(err),
		})
		return
	}
	defer client.Close()

	// 2. 检测系统架构
	archOut, _ := runSSHCommand(client, "uname -m")
	arch := "amd64"
	if len(archOut) > 0 && (contains(archOut, "aarch64") || contains(archOut, "arm64")) {
		arch = "arm64"
	}

	// 3. 构造安装命令（使用 Agent API 服务地址）
	// 从请求头获取后端可访问地址
	serverURL := fmt.Sprintf("http://%s", c.Request.Host)
	token := c.Query("token")

	installCmd := fmt.Sprintf(
		`curl -fsSL %s/agent/install.sh | sudo bash -s -- --server %s --token %s`,
		serverURL, serverURL, token,
	)

	// 4. 执行安装
	installOut, err := runSSHCommand(client, installCmd)
	if err != nil {
		c.JSON(http.StatusOK, models.Response{
			Code:    500,
			Message: "安装命令执行失败（可能需要 sudo 权限或 curl 未安装）",
			Data:    installOut,
		})
		return
	}

	// 5. 验证安装是否成功
	verifyOut, _ := runSSHCommand(client, "systemctl is-active aiops-agent 2>/dev/null || echo 'not_found'")
	installed := contains(verifyOut, "active") || contains(verifyOut, "running")

	if installed {
		// 更新服务器 Agent 状态
		h.db.Model(&server).Updates(map[string]interface{}{
			"agent_status":  "online",
			"agent_version": "installed",
			"last_check_at": time.Now(),
		})
	}

	c.JSON(http.StatusOK, models.Response{
		Code:    200,
		Message: func() string {
			if installed {
				return "Agent 安装成功，已自动启动并设置开机自启"
			}
			return "安装命令已执行，但 Agent 服务未检测到运行（请查看日志排查）"
		}(),
		Data: gin.H{
			"installed": installed,
			"arch":      arch,
			"output":    installOut,
			"verify":    verifyOut,
		},
	})
}

// installOne 安装单个服务器的 Agent（内部方法）
func (h *AgentInstallHandler) installOne(server models.Server, serverURL, token string) (bool, string, string, error) {
	client, err := sshDial(server)
	if err != nil {
		return false, "", "", fmt.Errorf("SSH连接失败: %w", err)
	}
	defer client.Close()

	if server.Status != "running" {
		return false, "", "", fmt.Errorf("服务器未运行")
	}

	installCmd := fmt.Sprintf(
		`curl -fsSL %s/agent/install.sh | sudo bash -s -- --server %s --token %s`,
		serverURL, serverURL, token,
	)

	installOut, err := runSSHCommand(client, installCmd)
	if err != nil {
		return false, installOut, "", err
	}

	verifyOut, _ := runSSHCommand(client, "systemctl is-active aiops-agent 2>/dev/null || echo 'not_found'")
	installed := contains(verifyOut, "active") || contains(verifyOut, "running")

	if installed {
		h.db.Model(&server).Updates(map[string]interface{}{
			"agent_status":  "online",
			"agent_version": "installed",
			"last_check_at": time.Now(),
		})
	}

	return installed, installOut, verifyOut, nil
}

// BatchInstallAgent 批量安装 Agent
// POST /servers/install-agent
func (h *AgentInstallHandler) BatchInstallAgent(c *gin.Context) {
	var req struct {
		ServerIDs []uint `json:"server_ids"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || len(req.ServerIDs) == 0 {
		c.JSON(http.StatusBadRequest, models.Response{Code: 400, Message: "请提供要安装的服务器 ID 列表"})
		return
	}

	serverURL := fmt.Sprintf("http://%s", c.Request.Host)
	token := c.Query("token")

	type InstallResult struct {
		ServerID  uint   `json:"server_id"`
		Hostname  string `json:"hostname"`
		IP        string `json:"ip"`
		Success   bool   `json:"success"`
		Message   string `json:"message"`
		Installed bool   `json:"installed"`
	}

	results := make([]InstallResult, 0, len(req.ServerIDs))
	successCount := 0

	for _, sid := range req.ServerIDs {
		var server models.Server
		if err := h.db.First(&server, sid).Error; err != nil {
			results = append(results, InstallResult{
				ServerID: sid,
				Success:  false,
				Message:  "服务器不存在",
			})
			continue
		}

		installed, _, _, err := h.installOne(server, serverURL, token)
		if err != nil {
			results = append(results, InstallResult{
				ServerID: server.ID,
				Hostname: server.Hostname,
				IP:       server.IP,
				Success:  false,
				Message:  err.Error(),
			})
		} else {
			successCount++
			msg := "安装成功"
			if !installed {
				msg = "安装命令已执行，Agent 服务未检测到运行"
			}
			results = append(results, InstallResult{
				ServerID:  server.ID,
				Hostname:  server.Hostname,
				IP:        server.IP,
				Success:   true,
				Installed: installed,
				Message:   msg,
			})
		}
	}

	c.JSON(http.StatusOK, models.Response{
		Code:    200,
		Message: fmt.Sprintf("完成 %d/%d 台服务器", successCount, len(req.ServerIDs)),
		Data: gin.H{
			"total":   len(req.ServerIDs),
			"success": successCount,
			"results": results,
		},
	})
}