feat: 加固 Agent CLI 执行边界并更新模型目录
收紧 SSH 命令护栏、权限校验与远程执行会话,并同步主流大模型供应商列表。 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
1 parent
67660ae1d6
commit
2fd1624bec
18 files changed
+480
-49
No files matched your search
@@ -523,12 +523,13 @@ func (ac *agentConn) buildTargetInfo() string {
|
||||
|
||||
// confirmFn 高危命令确认回调:推 confirm 事件并阻塞等前端回传 // Logan.Li
|
||||
func (ac *agentConn) confirmFn(toolName string, args map[string]interface{}, risk tools.RiskLevel) (bool, error) {
|
||||
safeArgs := guardrail.SanitizeMap(args)
|
||||
confirmID := fmt.Sprintf("cf-%d", time.Now().UnixNano())
|
||||
cmd, _ := args["command"].(string)
|
||||
cmd, _ := safeArgs["command"].(string)
|
||||
if cmd == "" && (toolName == "kubernetes_action" || toolName == "k8s_action") {
|
||||
cmd = fmt.Sprintf("%v %v", args["action"], args["name"])
|
||||
cmd = fmt.Sprintf("%v %v", safeArgs["action"], safeArgs["name"])
|
||||
}
|
||||
ac.write(wsEvent{Type: "confirm", Phase: core.PhaseApproval, Summary: "等待用户确认", ConfirmID: confirmID, ToolName: toolName, Command: cmd, Risk: risk.Display(), ToolArgs: args})
|
||||
ac.write(wsEvent{Type: "confirm", Phase: core.PhaseApproval, Summary: "等待用户确认", ConfirmID: confirmID, ToolName: toolName, Command: cmd, Risk: risk.Display(), ToolArgs: safeArgs})
|
||||
|
||||
ch := make(chan bool, 1)
|
||||
ac.cfMu.Lock()
|
||||
@@ -537,7 +538,7 @@ func (ac *agentConn) confirmFn(toolName string, args map[string]interface{}, ris
|
||||
|
||||
select {
|
||||
case approved := <-ch:
|
||||
ac.recordConfirmSpan(toolName, args, risk, approved)
|
||||
ac.recordConfirmSpan(toolName, safeArgs, risk, approved)
|
||||
return approved, nil
|
||||
case <-ac.taskCtx.Done():
|
||||
ac.delConfirm(confirmID)
|
||||
|
||||
@@ -84,8 +84,7 @@ func (h *AgentCLIHandler) DeleteSession(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, model.Response{Code: 500, Message: "删除消息失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
_ = h.db.Where("run_id = ?", sess.RunID).Delete(&model.AIRunSpan{}).Error
|
||||
_ = h.db.Where("run_id = ?", sess.RunID).Delete(&model.OpsCommand{}).Error
|
||||
// 会话历史可删除,但远程执行 Span 与命令审计必须独立保留。 // Logan.Li
|
||||
if err := h.db.Delete(&model.AIRun{}, id).Error; err != nil {
|
||||
c.JSON(http.StatusOK, model.Response{Code: 500, Message: "删除会话失败: " + err.Error()})
|
||||
return
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"ezops/model"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestDeleteAgentCLISessionRetainsExecutionAudit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("打开测试数据库失败: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.AIRun{}, &model.AIRunMessage{}, &model.AIRunSpan{}, &model.OpsCommand{}); err != nil {
|
||||
t.Fatalf("创建测试表失败: %v", err)
|
||||
}
|
||||
run := model.AIRun{RunID: "run-delete-audit", UserID: 7, Feature: "agent_cli", StartedAt: time.Now()}
|
||||
if err := db.Create(&run).Error; err != nil {
|
||||
t.Fatalf("创建会话失败: %v", err)
|
||||
}
|
||||
db.Create(&model.AIRunMessage{RunID: run.RunID, Role: "tool", Type: "tool_result", Content: "ok"})
|
||||
db.Create(&model.AIRunSpan{RunID: run.RunID, SpanID: "span-delete-audit", Kind: "tool", Name: "ssh_execute", StartedAt: time.Now()})
|
||||
db.Create(&model.OpsCommand{RunID: run.RunID, CommandID: "cmd-delete-audit", UserID: 7, Command: "uptime", ExecutedAt: time.Now()})
|
||||
|
||||
router := gin.New()
|
||||
router.DELETE("/sessions/:id", func(c *gin.Context) {
|
||||
c.Set("userID", uint(7))
|
||||
NewAgentCLIHandler(db).DeleteSession(c)
|
||||
})
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodDelete, "/sessions/"+strconv.FormatUint(uint64(run.ID), 10), nil))
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("删除响应状态=%d, body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
|
||||
assertCount := func(name string, target interface{}, want int64) {
|
||||
t.Helper()
|
||||
var count int64
|
||||
if err := db.Model(target).Where("run_id = ?", run.RunID).Count(&count).Error; err != nil {
|
||||
t.Fatalf("查询 %s 失败: %v", name, err)
|
||||
}
|
||||
if count != want {
|
||||
t.Fatalf("%s 数量=%d, want %d", name, count, want)
|
||||
}
|
||||
}
|
||||
assertCount("消息", &model.AIRunMessage{}, 0)
|
||||
assertCount("Span", &model.AIRunSpan{}, 1)
|
||||
assertCount("命令审计", &model.OpsCommand{}, 1)
|
||||
}
|
||||
@@ -141,11 +141,13 @@ func resolveEndpoint(provider, endpoint string) string {
|
||||
return "https://api.moonshot.cn/v1"
|
||||
case "kimi": // Kimi 编码套餐 OpenAI 兼容端点(注意:该端点有客户端白名单,通用调用可能被拒) // Logan.Li
|
||||
return "https://api.kimi.com/coding/v1"
|
||||
case "mimo": // 小米 MiMo OpenAI 兼容端点 // Logan.Li
|
||||
return "https://api.mimo.mi.com/v1"
|
||||
case "mimo": // 小米 MiMo 按量 OpenAI 兼容端点 // Logan.Li
|
||||
return "https://api.xiaomimimo.com/v1"
|
||||
case "mimo-plan": // 小米 MiMo Token Plan 包月专用端点(OpenAI 兼容,套餐密钥) // Logan.Li
|
||||
return "https://token-plan-cn.xiaomimimo.com/v1"
|
||||
case "anthropic": // Anthropic 官方 Messages 协议(anthropic.go 客户端拼 /v1/messages) // Logan.Li
|
||||
return "https://api.anthropic.com"
|
||||
case "glm-coding": // 智谱 GLM 编码套餐 Anthropic 协议端点(模型 GLM-5.2) // Logan.Li
|
||||
case "glm-coding": // 智谱 GLM 编码套餐 Anthropic 协议端点(模型 GLM-5.3) // Logan.Li
|
||||
return "https://open.bigmodel.cn/api/anthropic"
|
||||
case "minimax": // MiniMax 按量 OpenAI 兼容端点 // Logan.Li
|
||||
return "https://api.minimaxi.com/v1"
|
||||
|
||||
@@ -146,6 +146,29 @@ func AdminAuth() gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// PermissionAuth 校验当前用户是否拥有指定权限码,管理员直接放行。 // Logan.Li
|
||||
func PermissionAuth(db *gorm.DB, permissionCode string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if GetIsAdmin(c) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
var count int64
|
||||
err := db.Table("iam_user_roles ur").
|
||||
Joins("JOIN iam_roles r ON r.id = ur.role_id AND r.status = 1").
|
||||
Joins("JOIN iam_role_permissions rp ON rp.role_id = r.id").
|
||||
Joins("JOIN iam_permissions p ON p.id = rp.permission_id AND p.status = 1").
|
||||
Where("ur.user_id = ? AND p.code = ?", GetUserID(c), permissionCode).
|
||||
Count(&count).Error
|
||||
if err != nil || count == 0 {
|
||||
c.JSON(http.StatusForbidden, model.Response{Code: 403, Message: "无此功能权限"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// GetUserID 获取当前用户ID
|
||||
func GetUserID(c *gin.Context) uint {
|
||||
userID, exists := c.Get("userID")
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestPermissionAuth(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("打开测试数据库失败: %v", err)
|
||||
}
|
||||
statements := []string{
|
||||
"CREATE TABLE iam_user_roles (user_id INTEGER, role_id INTEGER)",
|
||||
"CREATE TABLE iam_roles (id INTEGER PRIMARY KEY, status INTEGER)",
|
||||
"CREATE TABLE iam_role_permissions (role_id INTEGER, permission_id INTEGER)",
|
||||
"CREATE TABLE iam_permissions (id INTEGER PRIMARY KEY, code TEXT, status INTEGER)",
|
||||
"INSERT INTO iam_user_roles(user_id, role_id) VALUES (7, 1)",
|
||||
"INSERT INTO iam_roles(id, status) VALUES (1, 1)",
|
||||
"INSERT INTO iam_role_permissions(role_id, permission_id) VALUES (1, 10)",
|
||||
"INSERT INTO iam_permissions(id, code, status) VALUES (10, 'aiops:agent-cli', 1)",
|
||||
}
|
||||
for _, statement := range statements {
|
||||
if err := db.Exec(statement).Error; err != nil {
|
||||
t.Fatalf("准备测试数据失败: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userID uint
|
||||
isAdmin bool
|
||||
wantStatus int
|
||||
}{
|
||||
{name: "granted", userID: 7, wantStatus: http.StatusNoContent},
|
||||
{name: "denied", userID: 8, wantStatus: http.StatusForbidden},
|
||||
{name: "admin", userID: 9, isAdmin: true, wantStatus: http.StatusNoContent},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set("userID", tc.userID)
|
||||
c.Set("isAdmin", tc.isAdmin)
|
||||
})
|
||||
router.GET("/agent", PermissionAuth(db, "aiops:agent-cli"), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/agent", nil))
|
||||
if recorder.Code != tc.wantStatus {
|
||||
t.Fatalf("状态码=%d, want %d, body=%s", recorder.Code, tc.wantStatus, recorder.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -6,7 +6,7 @@ import "time"
|
||||
// 支持 OpenAI 兼容接口(OpenAI / DeepSeek / 智谱GLM / 通义千问 / Moonshot / Kimi / 小米MiMo / Anthropic(需代理) / Ollama / 自定义)
|
||||
type LLMConfig struct {
|
||||
ID uint `gorm:"primarykey" json:"id"`
|
||||
Provider string `gorm:"size:50;not null;default:openai" json:"provider"` // deepseek/glm/glm-coding/kimi/moonshot/minimax/minimax-coding/mimo/siliconflow/openai/anthropic/ollama/vllm/custom
|
||||
Provider string `gorm:"size:50;not null;default:openai" json:"provider"` // deepseek/glm/glm-coding/kimi/moonshot/minimax/minimax-coding/mimo/mimo-plan/siliconflow/openai/anthropic/ollama/vllm/custom
|
||||
Model string `gorm:"size:100;not null;default:gpt-4o-mini" json:"model"` // 模型名,如 gpt-4o-mini / deepseek-v4-flash
|
||||
APIKey string `gorm:"size:500" json:"api_key"` // API Key(GET 时会被脱敏)
|
||||
Endpoint string `gorm:"size:500" json:"endpoint"` // 基础地址,如 https://api.openai.com/v1
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"time"
|
||||
|
||||
"ezops/pkg/agentcli/evolver"
|
||||
"ezops/pkg/agentcli/guardrail"
|
||||
"ezops/pkg/agentcli/llm"
|
||||
"ezops/pkg/agentcli/prompt"
|
||||
"ezops/pkg/agentcli/safety"
|
||||
@@ -293,9 +294,10 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
Content: "错误: 参数解析失败: " + err.Error()})
|
||||
continue
|
||||
}
|
||||
safeArgs := guardrail.SanitizeMap(args)
|
||||
|
||||
if msg := a.suppressBroadKubernetesGet(tc.Function.Name, args, history); msg != "" {
|
||||
a.emit(Event{Type: EventToolCall, Phase: PhaseToolRunning, Round: round, ToolName: tc.Function.Name, ToolArgs: args, Risk: tools.RiskSafe.Display(), Summary: "工具执行中"})
|
||||
a.emit(Event{Type: EventToolCall, Phase: PhaseToolRunning, Round: round, ToolName: tc.Function.Name, ToolArgs: safeArgs, Risk: tools.RiskSafe.Display(), Summary: "工具执行中"})
|
||||
a.emit(Event{Type: EventToolResult, Phase: PhaseObserving, Round: round, ToolName: tc.Function.Name,
|
||||
Output: msg, Success: true, Duration: "0s", OutputSize: len(msg)})
|
||||
obs = append(obs, llm.Message{Role: "tool", ToolCallID: tc.ID, Name: tc.Function.Name, Content: msg})
|
||||
@@ -303,12 +305,12 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
}
|
||||
|
||||
risk := a.toolRisk(tool, args)
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseSafety, Round: round, ToolName: tc.Function.Name, ToolArgs: args, Risk: risk.Display(), Summary: "安全检查"})
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseSafety, Round: round, ToolName: tc.Function.Name, ToolArgs: safeArgs, Risk: risk.Display(), Summary: "安全检查"})
|
||||
|
||||
// 安全检查(高危触发 confirmFn,由 controller 异步确认)
|
||||
approved, err := a.safetyCtl.Check(tool, args)
|
||||
if err != nil {
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseError, Round: round, ToolName: tc.Function.Name, ToolArgs: args, Risk: risk.Display(), Message: err.Error()})
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseError, Round: round, ToolName: tc.Function.Name, ToolArgs: safeArgs, Risk: risk.Display(), Message: guardrail.SanitizeText(err.Error())})
|
||||
obs = append(obs, llm.Message{Role: "tool", ToolCallID: tc.ID, Name: tc.Function.Name,
|
||||
Content: "安全检查失败: " + err.Error()})
|
||||
continue
|
||||
@@ -316,19 +318,19 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
if !approved {
|
||||
obs = append(obs, llm.Message{Role: "tool", ToolCallID: tc.ID, Name: tc.Function.Name,
|
||||
Content: "操作已被用户拒绝执行"})
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseBlocked, Round: round, ToolName: tc.Function.Name, ToolArgs: args, Risk: risk.Display(), Summary: "用户拒绝执行"})
|
||||
a.emit(Event{Type: EventStatus, Phase: PhaseBlocked, Round: round, ToolName: tc.Function.Name, ToolArgs: safeArgs, Risk: risk.Display(), Summary: "用户拒绝执行"})
|
||||
a.emit(Event{Type: EventToolResult, Round: round, ToolName: tc.Function.Name,
|
||||
Phase: PhaseObserving, Output: "(用户拒绝执行)", Success: false, OutputSize: len("(用户拒绝执行)")})
|
||||
continue
|
||||
}
|
||||
|
||||
// 执行
|
||||
a.emit(Event{Type: EventToolCall, Phase: PhaseToolRunning, Round: round, ToolName: tc.Function.Name, ToolArgs: args, Risk: risk.Display(), Summary: "工具执行中"})
|
||||
a.emit(Event{Type: EventToolCall, Phase: PhaseToolRunning, Round: round, ToolName: tc.Function.Name, ToolArgs: safeArgs, Risk: risk.Display(), Summary: "工具执行中"})
|
||||
if cached, ok := a.cachedToolResult(tc.Function.Name, args); ok {
|
||||
a.emit(Event{Type: EventToolResult, Phase: PhaseObserving, Round: round, ToolName: tc.Function.Name,
|
||||
Output: cached, Success: true, Duration: "0s", OutputSize: len(cached)})
|
||||
records = append(records, evolver.ToolCallRecord{
|
||||
ToolName: tc.Function.Name, Args: args, Output: cached,
|
||||
ToolName: tc.Function.Name, Args: safeArgs, Output: cached,
|
||||
Duration: 0, Success: true, RiskLevel: int(risk),
|
||||
})
|
||||
obs = append(obs, llm.Message{Role: "tool", ToolCallID: tc.ID, Name: tc.Function.Name, Content: cached})
|
||||
@@ -349,6 +351,9 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
errMsg = execErr.Error()
|
||||
success = false
|
||||
}
|
||||
// 工具结果必须先脱敏,再进入自动分析、事件、缓存、Evolver 和模型上下文。 // Logan.Li
|
||||
output = guardrail.SanitizeText(output)
|
||||
errMsg = guardrail.SanitizeText(errMsg)
|
||||
|
||||
// 超长命令输出自动分析,避免 token 爆炸 // Logan.Li
|
||||
if (tc.Function.Name == "ssh_execute" || tc.Function.Name == "pod_execute") && success && len(output) > a.autoAnalyzeThreshold() {
|
||||
@@ -365,10 +370,10 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
|
||||
a.emit(Event{Type: EventToolResult, Phase: PhaseObserving, Round: round, ToolName: tc.Function.Name,
|
||||
Output: output, Success: success, Duration: dur.Truncate(time.Millisecond).String(), OutputSize: len(output)})
|
||||
a.storeToolResult(tc.Function.Name, args, output, success)
|
||||
a.storeToolResult(tc.Function.Name, safeArgs, output, success)
|
||||
|
||||
records = append(records, evolver.ToolCallRecord{
|
||||
ToolName: tc.Function.Name, Args: args, Output: output,
|
||||
ToolName: tc.Function.Name, Args: safeArgs, Output: output,
|
||||
Duration: dur, Success: success, RiskLevel: int(risk),
|
||||
})
|
||||
|
||||
@@ -376,7 +381,7 @@ func (a *Agent) processToolCalls(ctx context.Context, calls []llm.ToolCall, hist
|
||||
var content string
|
||||
switch {
|
||||
case execErr != nil:
|
||||
content = "执行错误: " + execErr.Error()
|
||||
content = "执行错误: " + errMsg
|
||||
failInc++
|
||||
case !success:
|
||||
content = fmt.Sprintf("执行失败: %s\n输出: %s", errMsg, output)
|
||||
|
||||
@@ -15,7 +15,7 @@ var sensitivePatterns = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)((?:mysql|postgres|postgresql|redis|mongodb)://[^:\s/@]+:)[^@\s]+(@)`),
|
||||
regexp.MustCompile(`(?is)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----`),
|
||||
regexp.MustCompile(`(?i)((?:Cookie|Set-Cookie):\s*)[^\r\n]+`),
|
||||
regexp.MustCompile(`(?i)(AKIA[0-9A-Z]{16})`),
|
||||
regexp.MustCompile(`(?i)AKIA[0-9A-Z]{16}`),
|
||||
}
|
||||
|
||||
var sensitiveKeys = map[string]bool{
|
||||
@@ -65,6 +65,18 @@ func SanitizeJSON(raw string) string {
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// SanitizeMap 返回脱敏后的参数副本,避免事件、审计和记忆泄露工具参数中的凭证。 // Logan.Li
|
||||
func SanitizeMap(input map[string]interface{}) map[string]interface{} {
|
||||
if input == nil {
|
||||
return nil
|
||||
}
|
||||
sanitized, ok := sanitizeValue(input).(map[string]interface{})
|
||||
if !ok {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
return sanitized
|
||||
}
|
||||
|
||||
func sanitizeValue(v interface{}) interface{} {
|
||||
switch val := v.(type) {
|
||||
case map[string]interface{}:
|
||||
|
||||
@@ -6,9 +6,9 @@ import (
|
||||
)
|
||||
|
||||
func TestSanitizeTextMasksCommonSecrets(t *testing.T) {
|
||||
input := "Authorization: Bearer abcdefghijklmnopqrstuvwxyz password=secret123 mysql://root:pass@db/prod"
|
||||
input := "Authorization: Bearer abcdefghijklmnopqrstuvwxyz password=secret123 mysql://root:pass@db/prod AKIA1234567890ABCDEF"
|
||||
got := SanitizeText(input)
|
||||
for _, leak := range []string{"abcdefghijklmnopqrstuvwxyz", "secret123", "root:pass@db"} {
|
||||
for _, leak := range []string{"abcdefghijklmnopqrstuvwxyz", "secret123", "root:pass@db", "AKIA1234567890ABCDEF"} {
|
||||
if strings.Contains(got, leak) {
|
||||
t.Fatalf("敏感信息未脱敏: %s in %s", leak, got)
|
||||
}
|
||||
@@ -29,3 +29,22 @@ func TestSanitizeJSONMasksSensitiveKeys(t *testing.T) {
|
||||
t.Fatalf("非敏感字段不应被改写: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeMapReturnsMaskedCopy(t *testing.T) {
|
||||
input := map[string]interface{}{
|
||||
"command": "curl -H 'Authorization: Bearer abcdefghijklmnop' https://example.test",
|
||||
"nested": map[string]interface{}{"password": "plain-secret"},
|
||||
}
|
||||
|
||||
got := SanitizeMap(input)
|
||||
if got["command"] == input["command"] {
|
||||
t.Fatal("命令中的令牌未脱敏")
|
||||
}
|
||||
nested := got["nested"].(map[string]interface{})
|
||||
if nested["password"] != "[REDACTED]" {
|
||||
t.Fatalf("密码字段未脱敏: %#v", nested["password"])
|
||||
}
|
||||
if input["nested"].(map[string]interface{})["password"] != "plain-secret" {
|
||||
t.Fatal("脱敏不得修改原始参数")
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,8 @@ type SSHExecuteTool struct {
|
||||
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 {
|
||||
@@ -82,12 +84,12 @@ func (t *SSHExecuteTool) Execute(ctx context.Context, args map[string]interface{
|
||||
return &Result{Success: false, Error: "参数 command 不能为空"}, nil
|
||||
}
|
||||
|
||||
timeout := t.timeout
|
||||
timeout := boundedSSHTimeout(t.timeout)
|
||||
if sec := IntParam(args, "timeout"); sec > 0 {
|
||||
timeout = time.Duration(sec) * time.Second
|
||||
timeout = boundedSSHTimeout(time.Duration(sec) * time.Second)
|
||||
}
|
||||
|
||||
res, err := t.client.Exec(command, timeout)
|
||||
res, err := t.client.ExecContext(ctx, command, timeout)
|
||||
if err != nil {
|
||||
return &Result{Success: false, Error: fmt.Sprintf("SSH 连接失败: %v", err)}, nil
|
||||
}
|
||||
@@ -103,4 +105,15 @@ func (t *SSHExecuteTool) Execute(ctx context.Context, args map[string]interface{
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestBoundedSSHTimeout(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in time.Duration
|
||||
want time.Duration
|
||||
}{
|
||||
{name: "non-positive", in: 0, want: time.Second},
|
||||
{name: "normal", in: 60 * time.Second, want: 60 * time.Second},
|
||||
{name: "too-large", in: 24 * time.Hour, want: 300 * time.Second},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := boundedSSHTimeout(tc.in); got != tc.want {
|
||||
t.Fatalf("boundedSSHTimeout(%v)=%v, want %v", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -29,13 +29,21 @@ func NewPersistentClient(cfg SSHConfig) (*PersistentClient, error) {
|
||||
|
||||
// Exec 在长连接上执行命令;连接断开时自动重连一次 // Logan.Li
|
||||
func (p *PersistentClient) Exec(command string, timeout time.Duration) (*ExecResult, error) {
|
||||
return p.ExecContext(context.Background(), command, timeout)
|
||||
}
|
||||
|
||||
// ExecContext 在长连接上执行命令,并在调用方取消时关闭当前远端 session。 // Logan.Li
|
||||
func (p *PersistentClient) ExecContext(ctx context.Context, command string, timeout time.Duration) (*ExecResult, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
result, err := p.execLocked(command, timeout)
|
||||
result, err := p.execLockedContext(ctx, command, timeout)
|
||||
if err == nil {
|
||||
return result, nil
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return result, ctx.Err()
|
||||
}
|
||||
// 连接级失败:重连后重试一次 // Logan.Li
|
||||
if p.client != nil {
|
||||
_ = p.client.Close()
|
||||
@@ -46,11 +54,16 @@ func (p *PersistentClient) Exec(command string, timeout time.Duration) (*ExecRes
|
||||
return result, err
|
||||
}
|
||||
p.client = client
|
||||
return p.execLocked(command, timeout)
|
||||
return p.execLockedContext(ctx, command, timeout)
|
||||
}
|
||||
|
||||
// execLocked 假定已持锁,在当前 client 上执行 // Logan.Li
|
||||
func (p *PersistentClient) execLocked(command string, timeout time.Duration) (*ExecResult, error) {
|
||||
return p.execLockedContext(context.Background(), command, timeout)
|
||||
}
|
||||
|
||||
// execLockedContext 假定已持锁,并将调用方上下文传递给当前 SSH session。 // Logan.Li
|
||||
func (p *PersistentClient) execLockedContext(ctx context.Context, command string, timeout time.Duration) (*ExecResult, error) {
|
||||
if p.client == nil {
|
||||
client, err := connect(p.cfg)
|
||||
if err != nil {
|
||||
@@ -58,7 +71,7 @@ func (p *PersistentClient) execLocked(command string, timeout time.Duration) (*E
|
||||
}
|
||||
p.client = client
|
||||
}
|
||||
return execWithClient(p.client, p.cfg, command, timeout)
|
||||
return execWithClientContext(ctx, p.client, p.cfg, command, timeout)
|
||||
}
|
||||
|
||||
// Close 关闭长连接 // Logan.Li
|
||||
@@ -75,6 +88,11 @@ func (p *PersistentClient) Close() error {
|
||||
|
||||
// execWithClient 在已有 SSH client 上执行命令(供 Exec / PersistentClient 共用) // Logan.Li
|
||||
func execWithClient(client *ssh.Client, cfg SSHConfig, command string, timeout time.Duration) (*ExecResult, error) {
|
||||
return execWithClientContext(context.Background(), client, cfg, command, timeout)
|
||||
}
|
||||
|
||||
// execWithClientContext 在已有 SSH client 上执行命令,并响应调用方取消。 // Logan.Li
|
||||
func execWithClientContext(parent context.Context, client *ssh.Client, cfg SSHConfig, command string, timeout time.Duration) (*ExecResult, error) {
|
||||
start := time.Now()
|
||||
result := &ExecResult{Host: cfg.Host, Port: cfg.Port}
|
||||
|
||||
@@ -90,7 +108,7 @@ func execWithClient(client *ssh.Client, cfg SSHConfig, command string, timeout t
|
||||
timeout = 60 * time.Second
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
ctx, cancel := context.WithTimeout(parent, timeout)
|
||||
defer cancel()
|
||||
|
||||
type execOut struct {
|
||||
@@ -117,10 +135,16 @@ func execWithClient(client *ssh.Client, cfg SSHConfig, command string, timeout t
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = session.Signal(ssh.SIGKILL)
|
||||
_ = session.Close()
|
||||
result.Duration = time.Since(start).Milliseconds()
|
||||
result.Error = fmt.Sprintf("执行超时(%v)", timeout)
|
||||
if parent.Err() != nil {
|
||||
result.Error = "执行已取消"
|
||||
} else {
|
||||
result.Error = fmt.Sprintf("执行超时(%v)", timeout)
|
||||
}
|
||||
result.ExitCode = -1
|
||||
return result, fmt.Errorf("exec timeout after %v", timeout)
|
||||
return result, ctx.Err()
|
||||
case out := <-ch:
|
||||
result.Duration = time.Since(start).Milliseconds()
|
||||
result.Output = string(out.output)
|
||||
@@ -141,4 +165,4 @@ func execWithClient(client *ssh.Client, cfg SSHConfig, command string, timeout t
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
@@ -952,8 +952,8 @@ func Register(r *gin.Engine, db *gorm.DB, log *zap.Logger, cfg *config.Config, c
|
||||
// 指标浏览器 — 监控指标元数据 + 时序查询 // Logan.Li
|
||||
// Agent CLI — AI 运维 Agent(WebSocket 推流 + 会话历史回放) // Logan.Li
|
||||
agentCLIHandler := controller.NewAgentCLIHandler(db)
|
||||
r.GET("/ws/agent-cli", middleware.JWTAuth(&cfg.JWT, db), agentCLIHandler.Run)
|
||||
agentCLIGroup := authorized.Group("/aiops/agent-cli")
|
||||
r.GET("/ws/agent-cli", middleware.JWTAuth(&cfg.JWT, db), middleware.PermissionAuth(db, "aiops:agent-cli"), agentCLIHandler.Run)
|
||||
agentCLIGroup := authorized.Group("/aiops/agent-cli", middleware.PermissionAuth(db, "aiops:agent-cli"))
|
||||
{
|
||||
agentCLIGroup.GET("/sessions", agentCLIHandler.ListSessions)
|
||||
agentCLIGroup.GET("/sessions/:id", agentCLIHandler.GetSession)
|
||||
|
||||
Reference in new issue
Block a user