Logan.LiandCursor
2fd1624bec
收紧 SSH 命令护栏、权限校验与远程执行会话，并同步主流大模型供应商列表。 Co-authored-by: Cursor <cursoragent@cursor.com>
120 lines
3.7 KiB
Go
120 lines
3.7 KiB
Go
// ssh_execute.go - 远程命令执行工具 // Logan.Li
|
||
// 复用 aiops 既有 pkg/remoteexec（Go 原生 SSH），凭据由 controller 从 model.Server 构建后注入。
|
||
// 任务级复用 PersistentClient 长连接，避免每条命令重建 SSH 握手。
|
||
// 实现 DynamicRiskTool：依据具体命令动态评估风险，决定是否触发用户确认。
|
||
package tools
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
|
||
"ezops/pkg/remoteexec"
|
||
)
|
||
|
||
// SSHExecuteTool 远程命令执行工具 // Logan.Li
|
||
type SSHExecuteTool struct {
|
||
client *remoteexec.PersistentClient // 任务级 SSH 长连接
|
||
timeout time.Duration // 默认命令超时
|
||
}
|
||
|
||
const maxSSHCommandTimeout = 300 * time.Second
|
||
|
||
// NewSSHExecuteTool 创建远程执行工具（内部建立 SSH 长连接） // Logan.Li
|
||
func NewSSHExecuteTool(cfg remoteexec.SSHConfig, timeout time.Duration) (*SSHExecuteTool, error) {
|
||
if timeout <= 0 {
|
||
timeout = 60 * time.Second
|
||
}
|
||
client, err := remoteexec.NewPersistentClient(cfg)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("SSH 长连接建立失败: %w", err)
|
||
}
|
||
return &SSHExecuteTool{client: client, timeout: timeout}, nil
|
||
}
|
||
|
||
// Close 释放 SSH 长连接 // Logan.Li
|
||
func (t *SSHExecuteTool) Close() error {
|
||
if t.client == nil {
|
||
return nil
|
||
}
|
||
return t.client.Close()
|
||
}
|
||
|
||
// Name 工具名 // Logan.Li
|
||
func (t *SSHExecuteTool) Name() string { return "ssh_execute" }
|
||
|
||
// Description 给 LLM 的工具说明 // Logan.Li
|
||
func (t *SSHExecuteTool) Description() string {
|
||
return "在远程目标服务器上执行 shell 命令并返回合并的 stdout+stderr 输出与退出码。" +
|
||
"用于排障（查看日志/进程/磁盘/网络）、安装与配置软件、重启服务等。" +
|
||
"高危命令（rm -rf /、mkfs、reboot 等）需用户确认。"
|
||
}
|
||
|
||
// Parameters JSON Schema 参数定义 // Logan.Li
|
||
func (t *SSHExecuteTool) Parameters() map[string]interface{} {
|
||
return map[string]interface{}{
|
||
"type": "object",
|
||
"properties": map[string]interface{}{
|
||
"command": map[string]interface{}{
|
||
"type": "string",
|
||
"description": "要执行的 shell 命令，例如 df -h、systemctl status nginx、tail -n 100 /var/log/messages",
|
||
},
|
||
"timeout": map[string]interface{}{
|
||
"type": "integer",
|
||
"description": "命令超时秒数（默认 60）",
|
||
},
|
||
},
|
||
"required": []string{"command"},
|
||
}
|
||
}
|
||
|
||
// RiskLevel 静态风险等级（高危） // Logan.Li
|
||
func (t *SSHExecuteTool) RiskLevel() RiskLevel { return RiskHigh }
|
||
|
||
// RiskLevelForArgs 依据命令动态评估风险 // Logan.Li
|
||
func (t *SSHExecuteTool) RiskLevelForArgs(args map[string]interface{}) RiskLevel {
|
||
return AssessBashRisk(StringParam(args, "command"))
|
||
}
|
||
|
||
// Execute 执行远程命令 // Logan.Li
|
||
func (t *SSHExecuteTool) Execute(ctx context.Context, args map[string]interface{}) (*Result, error) {
|
||
command := strings.TrimSpace(StringParam(args, "command"))
|
||
if command == "" {
|
||
return &Result{Success: false, Error: "参数 command 不能为空"}, nil
|
||
}
|
||
|
||
timeout := boundedSSHTimeout(t.timeout)
|
||
if sec := IntParam(args, "timeout"); sec > 0 {
|
||
timeout = boundedSSHTimeout(time.Duration(sec) * time.Second)
|
||
}
|
||
|
||
res, err := t.client.ExecContext(ctx, command, timeout)
|
||
if err != nil {
|
||
return &Result{Success: false, Error: fmt.Sprintf("SSH 连接失败: %v", err)}, nil
|
||
}
|
||
|
||
output := res.Output
|
||
if !strings.HasSuffix(output, "\n") && output != "" {
|
||
output += "\n"
|
||
}
|
||
output += fmt.Sprintf("[exit_code=%d]", res.ExitCode)
|
||
|
||
return &Result{
|
||
Success: res.Success,
|
||
Output: output,
|
||
Error: res.Error,
|
||
}, nil
|
||
}
|
||
|
||
// boundedSSHTimeout 将模型提供的超时限制在平台允许范围内。 // Logan.Li
|
||
func boundedSSHTimeout(timeout time.Duration) time.Duration {
|
||
if timeout <= 0 {
|
||
return time.Second
|
||
}
|
||
if timeout > maxSSHCommandTimeout {
|
||
return maxSSHCommandTimeout
|
||
}
|
||
return timeout
|
||
}
|