tasks.go
1 // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // Requirement: Any integration or derivative work must explicitly attribute 16 // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its 17 // documentation or user interface, as detailed in the NOTICE file. 18 19 package agent 20 21 import ( 22 "context" 23 "encoding/json" 24 "fmt" 25 "os" 26 "path" 27 "path/filepath" 28 "strings" 29 "sync" 30 "time" 31 32 "github.com/Tencent/AI-Infra-Guard/pkg/vulstruct" 33 iputil "github.com/projectdiscovery/utils/ip" 34 35 "github.com/Tencent/AI-Infra-Guard/common/utils" 36 37 "github.com/Tencent/AI-Infra-Guard/common/utils/models" 38 39 "github.com/Tencent/AI-Infra-Guard/common/runner" 40 "github.com/Tencent/AI-Infra-Guard/internal/gologger" 41 "github.com/Tencent/AI-Infra-Guard/internal/options" 42 "github.com/google/uuid" 43 ) 44 45 // ResultCallback 任务结果回调函数类型 46 type ResultCallback func(result map[string]interface{}) 47 48 // ActionLogCallback 插件日志回调函数类型 49 type ActionLogCallback func(actionId, tool, planStepId, actionLog string) 50 51 // ToolUsedCallback 插件工作状态回调函数类型 52 type ToolUsedCallback func(planStepId, statusId, description string, tools []Tool) 53 54 // NewPlanStepCallback 新建执行步骤回调函数类型 55 type NewPlanStepCallback func(stepId, title string) 56 57 // StatusUpdateCallback 更新步骤状态回调函数类型 58 type StatusUpdateCallback func(planStepId, statusId, agentStatus, brief, description string) 59 60 // PlanUpdateCallback 更新任务计划回调函数类型 61 type PlanUpdateCallback func(tasks []SubTask) 62 63 type ErrorCallback func(error string) 64 65 // TaskCallbacks 任务回调函数集合 66 type TaskCallbacks struct { 67 ResultCallback ResultCallback // 任务结果回调 68 ToolUseLogCallback ActionLogCallback // 插件日志回调 69 ToolUsedCallback ToolUsedCallback // 插件状态回调 70 NewPlanStepCallback NewPlanStepCallback // 新建执行步骤回调 71 StepStatusUpdateCallback StatusUpdateCallback // 更新步骤状态回调 72 PlanUpdateCallback PlanUpdateCallback // 更新任务计划回调 73 ErrorCallback ErrorCallback // 错误回调 74 } 75 76 type TaskInterface interface { 77 GetName() string 78 Execute(ctx context.Context, request TaskRequest, callbacks TaskCallbacks) error 79 } 80 81 // ScanRequest 扫描请求结构 82 type ScanRequest struct { 83 Target []string `json:"-"` 84 Headers map[string]string `json:"headers"` 85 Timeout int `json:"timeout,omitempty"` 86 Model struct { 87 Model string `json:"model"` 88 Token string `json:"token"` 89 BaseUrl string `json:"base_url"` 90 } `json:"model,omitempty"` 91 } 92 93 type AIInfraScanAgent struct { 94 Server string 95 } 96 97 func (t *AIInfraScanAgent) GetName() string { 98 return TaskTypeAIInfraScan 99 } 100 func (t *AIInfraScanAgent) Execute(ctx context.Context, request TaskRequest, callbacks TaskCallbacks) error { 101 var reqScan ScanRequest 102 if len(request.Params) > 0 { 103 if err := json.Unmarshal(request.Params, &reqScan); err != nil { 104 return err 105 } 106 } 107 108 language := request.Language 109 if language == "" { 110 language = "zh" 111 } 112 113 // 初始化语言文本 114 texts := initTexts(language) 115 116 // 处理目标和附件 117 targets, err := t.prepareTargets(request, reqScan, texts) 118 if err != nil { 119 return err 120 } 121 reqScan.Target = targets 122 123 if reqScan.Timeout == 0 { 124 reqScan.Timeout = 30 125 } 126 127 // 判断是否使用AI模式 128 useAI := reqScan.Model.Token != "" || reqScan.Model.BaseUrl != "" 129 var model *models.OpenAI 130 if useAI { 131 if reqScan.Model.BaseUrl == "" && reqScan.Model.Model == "" { 132 return fmt.Errorf("model parameters are required") 133 } 134 model = &models.OpenAI{ 135 BaseUrl: reqScan.Model.BaseUrl, 136 Model: reqScan.Model.Model, 137 Key: reqScan.Model.Token, 138 } 139 } 140 141 return t.executeScan(ctx, request, reqScan, texts, callbacks, model) 142 } 143 144 // scanTexts 包含所有语言相关的文本 145 type scanTexts struct { 146 initEnv, execScan, genReport, aigWorking, aigCompleted, initConfig, portDetection, targetConfig, scanComplete, reportGen string 147 createTempDir, downloadFile, readFile, portScan, foundPort, portCount, targetCount, foundVuln, noVuln, execError, config, scanResult string 148 initDescTemplate, portDetectDescTemplate, portCompleteTemplate, targetCountTemplate, webAppTemplate, vulnFoundTemplate, noVulnTemplate, errorTemplate string 149 scanningDesc, execScanDesc, scanCompleteDesc, taskCompleteDesc, reportGenDesc, reportGenToolDesc, scanResultTemplate string 150 downloadFileLog, execScanTool, scanTool, generateReportTool, nmapTool, aiScannerTool, reportGeneratorTool, scanOperation, targetSystem, generateReport string 151 } 152 153 // initTexts 初始化语言文本 154 func initTexts(language string) scanTexts { 155 var texts scanTexts 156 if language == "en" { 157 texts.initEnv = "Initialize scan environment" 158 texts.execScan = "Execute AI infrastructure scan" 159 texts.genReport = "Generate scan report" 160 texts.aigWorking = "A.I.G is working" 161 texts.aigCompleted = "A.I.G completed work" 162 texts.initConfig = "Initialization configuration completed" 163 texts.portDetection = "Auto-detecting ports" 164 texts.targetConfig = "Target configuration completed" 165 texts.scanComplete = "Scan completed" 166 texts.reportGen = "Generating scan report" 167 texts.createTempDir = "Failed to create temporary directory" 168 texts.downloadFile = "Failed to download file" 169 texts.readFile = "Failed to read file" 170 texts.portScan = "Port scan" 171 texts.foundPort = "Found port" 172 texts.portCount = "Port count" 173 texts.targetCount = "Target count" 174 texts.foundVuln = "Found vulnerabilities" 175 texts.noVuln = "Scan completed, no vulnerabilities found" 176 texts.execError = "Execution error" 177 texts.config = "Configuration" 178 texts.scanResult = "Scan result" 179 texts.initDescTemplate = "Starting to initialize AI infrastructure scan environment" 180 texts.portDetectDescTemplate = "Auto-detecting IP: %s" 181 texts.portCompleteTemplate = "%s port detection completed" 182 texts.targetCountTemplate = "Target count: %d" 183 texts.webAppTemplate = "WEB Application: %s " 184 texts.vulnFoundTemplate = "URL:%s %sFound vulnerabilities:%d\n" 185 texts.noVulnTemplate = "URL:%s %sScan completed, no vulnerabilities found\n" 186 texts.errorTemplate = "Execution error: host:%s %s\n" 187 texts.scanningDesc = "Scanning..." 188 texts.execScanDesc = "Executing AI infrastructure scan" 189 texts.scanCompleteDesc = "Scan completed" 190 texts.taskCompleteDesc = "AI infrastructure scan task completed" 191 texts.reportGenDesc = "Generating scan report" 192 texts.reportGenToolDesc = "Generating scan report" 193 texts.scanResultTemplate = "Scan results: %d items" 194 texts.downloadFileLog = "Starting to download file: %s" 195 texts.execScanTool = "Execute scan" 196 texts.scanTool = "Scan" 197 texts.generateReportTool = "Generate report" 198 texts.nmapTool = "nmap" 199 texts.aiScannerTool = "ai_scanner" 200 texts.reportGeneratorTool = "report_generator" 201 texts.scanOperation = "Scan" 202 texts.targetSystem = "Target system" 203 texts.generateReport = "Generate report" 204 } else { 205 texts.initEnv = "准备扫描环境" 206 texts.execScan = "执行深度扫描" 207 texts.genReport = "智能分析与报告生成" 208 texts.aigWorking = "Agent正在工作" 209 texts.aigCompleted = "Agent完成工作" 210 texts.initConfig = "初始化配置完成" 211 texts.portDetection = "正在自动识别端口" 212 texts.targetConfig = "目标配置完成" 213 texts.scanComplete = "扫描完成" 214 texts.reportGen = "生成扫描报告" 215 texts.createTempDir = "创建临时目录失败" 216 texts.downloadFile = "下载文件失败" 217 texts.readFile = "读取文件失败" 218 texts.portScan = "端口扫描" 219 texts.foundPort = "发现端口" 220 texts.portCount = "端口数量" 221 texts.targetCount = "目标数量" 222 texts.foundVuln = "发现漏洞" 223 texts.noVuln = "扫描完成,未发现漏洞" 224 texts.execError = "执行错误" 225 texts.config = "配置" 226 texts.scanResult = "Agent发现" 227 texts.initDescTemplate = "开始初始化AI基础设施扫描环境" 228 texts.portDetectDescTemplate = "正在自动识别IP: %s" 229 texts.portCompleteTemplate = "%s 端口探测完成" 230 texts.targetCountTemplate = "目标数量: %d" 231 texts.webAppTemplate = "WEB应用: %s " 232 texts.vulnFoundTemplate = "URL:%s %s发现漏洞:%d\n" 233 texts.noVulnTemplate = "URL:%s %s扫描完成,未发现漏洞\n" 234 texts.errorTemplate = "执行错误: host:%s %s\n" 235 texts.scanningDesc = "开始对目标系统进行多维度扫描" 236 texts.execScanDesc = "正在执行AI基础设施扫描" 237 texts.scanCompleteDesc = "完成扫描" 238 texts.taskCompleteDesc = "AI基础设施扫描任务完成" 239 texts.reportGenDesc = "我需要提供更有价值的洞察..." 240 texts.reportGenToolDesc = "正在生成扫描报告" 241 texts.scanResultTemplate = "扫描结果: %d 条" 242 texts.downloadFileLog = "开始下载文件: %s" 243 texts.execScanTool = "执行扫描" 244 texts.scanTool = "扫描" 245 texts.generateReportTool = "生成报告" 246 texts.nmapTool = "nmap" 247 texts.aiScannerTool = "ai_scanner" 248 texts.reportGeneratorTool = "report_generator" 249 texts.scanOperation = "扫描" 250 texts.targetSystem = "目标系统" 251 texts.generateReport = "生成结构化扫描报告" 252 } 253 return texts 254 } 255 256 // prepareTargets 处理目标和附件 257 func (t *AIInfraScanAgent) prepareTargets(request TaskRequest, reqScan ScanRequest, texts scanTexts) ([]string, error) { 258 targets := strings.Split(strings.TrimSpace(request.Content), "\n") 259 260 if len(request.Attachments) == 0 { 261 return targets, nil 262 } 263 264 tempDir := "temp_uploads" 265 if err := os.MkdirAll(tempDir, 0755); err != nil { 266 gologger.Errorf("%s: %v", texts.createTempDir, err) 267 return nil, err 268 } 269 270 for _, file := range request.Attachments { 271 gologger.Infof(texts.downloadFileLog, file) 272 fileName := filepath.Join(tempDir, fmt.Sprintf("tmp-%d%s", time.Now().UnixMicro(), filepath.Ext(file))) 273 // Verify the path is within tempDir to prevent path traversal 274 absTempDir, _ := filepath.Abs(tempDir) 275 absFileName, _ := filepath.Abs(fileName) 276 if !strings.HasPrefix(absFileName, absTempDir+string(os.PathSeparator)) { 277 gologger.WithError(fmt.Errorf("非法路径: %s", file)).Errorln(texts.downloadFile) 278 return nil, fmt.Errorf("非法文件路径") 279 } 280 if err := utils.DownloadFile(t.Server, request.SessionId, file, fileName); err != nil { 281 gologger.WithError(err).Errorln(texts.downloadFile) 282 return nil, err 283 } 284 lines, err := os.ReadFile(fileName) 285 if err != nil { 286 gologger.WithError(err).Errorln(texts.readFile) 287 return nil, err 288 } 289 targets = append(targets, strings.Split(string(lines), "\n")...) 290 } 291 292 return targets, nil 293 } 294 295 // scanPortsAndPrepareTargets 扫描端口并准备最终目标列表 296 func (t *AIInfraScanAgent) scanPortsAndPrepareTargets(targets []string, step01 string, texts scanTexts, callbacks TaskCallbacks) ([]string, error) { 297 finalTargets := []string{} 298 var hosts []string 299 300 for _, target := range targets { 301 if iputil.IsIP(target) { 302 hosts = append(hosts, target) 303 } 304 finalTargets = append(finalTargets, target) 305 } 306 307 for _, host := range hosts { 308 statusNmap := uuid.NewString() 309 toolId := uuid.NewString() 310 callbacks.StepStatusUpdateCallback(step01, statusNmap, AgentStatusRunning, texts.portDetection, fmt.Sprintf(texts.portDetectDescTemplate, host)) 311 callbacks.ToolUsedCallback(step01, statusNmap, texts.nmapTool, []Tool{ 312 CreateTool(toolId, texts.nmapTool, SubTaskStatusDoing, texts.portScan, texts.nmapTool, "-T4 -p 11434,1337,7000-9000,18789", ""), 313 }) 314 315 portScanResult, err := utils.NmapScan(host, "11434,1337,7000-9000,18789") 316 if err != nil { 317 return nil, err 318 } 319 320 success := 0 321 for _, port := range portScanResult.Hosts { 322 address := port.Address.Addr 323 for _, ported := range port.Ports.PortList { 324 if ported.State.State == "open" { 325 finalTargets = append(finalTargets, fmt.Sprintf("%s:%d", address, ported.PortID)) 326 success += 1 327 callbacks.ToolUseLogCallback(toolId, texts.nmapTool, step01, fmt.Sprintf("%s: %s:%d\n", texts.foundPort, address, ported.PortID)) 328 } 329 } 330 } 331 332 callbacks.ToolUsedCallback(step01, statusNmap, texts.nmapTool, []Tool{ 333 CreateTool(toolId, texts.nmapTool, SubTaskStatusDone, texts.portScan, texts.nmapTool, "-T4", fmt.Sprintf("%s: %d", texts.portCount, success)), 334 }) 335 callbacks.StepStatusUpdateCallback(step01, statusNmap, AgentStatusCompleted, fmt.Sprintf(texts.portCompleteTemplate, host), "") 336 } 337 338 return finalTargets, nil 339 } 340 341 // executeScan 执行扫描任务的统一入口 342 func (t *AIInfraScanAgent) executeScan(ctx context.Context, request TaskRequest, reqScan ScanRequest, texts scanTexts, callbacks TaskCallbacks, model *models.OpenAI) error { 343 // 创建任务计划 344 taskTitles := []string{texts.initEnv, texts.execScan, texts.genReport} 345 var tasks []SubTask 346 for _, title := range taskTitles { 347 tasks = append(tasks, CreateSubTask(SubTaskStatusTodo, title, 0, uuid.NewString())) 348 } 349 callbacks.PlanUpdateCallback(tasks) 350 351 // 步骤1: 初始化环境 352 step01 := tasks[0].StepId 353 callbacks.NewPlanStepCallback(step01, texts.initEnv) 354 statusId01 := uuid.NewString() 355 callbacks.StepStatusUpdateCallback(step01, statusId01, AgentStatusRunning, "Thinking", "") 356 357 // 配置选项 358 opts := &options.Options{ 359 TimeOut: reqScan.Timeout, 360 RateLimit: 200, 361 FPTemplates: t.Server, 362 AdvTemplates: t.Server, 363 WebServer: false, 364 Target: reqScan.Target, 365 LoadRemote: true, 366 } 367 368 headers := make([]string, 0) 369 for k, v := range reqScan.Headers { 370 headers = append(headers, k+":"+v) 371 } 372 opts.Headers = headers 373 374 // 扫描端口并准备目标 375 targets, err := t.scanPortsAndPrepareTargets(reqScan.Target, step01, texts, callbacks) 376 if err != nil { 377 return err 378 } 379 opts.Target = targets 380 381 // AI模式下的初始化反馈 382 if model != nil { 383 config := "" 384 var configMu sync.Mutex 385 386 // 临时回调收集配置信息 387 tempCallback := func(data interface{}) { 388 if v, ok := data.(runner.Step01); ok { 389 configMu.Lock() 390 config += v.Text + "\n" 391 configMu.Unlock() 392 } 393 } 394 opts.SetCallback(tempCallback) 395 396 // 创建runner获取配置 397 r, err := runner.New(opts) 398 if err != nil { 399 return fmt.Errorf("new runner failed: %v", err) 400 } 401 r.Close() 402 403 // AI分析初始化配置 404 prompt := fmt.Sprintf(`你在执行AI基础设施扫描任务，你正在完成todo:%s 405 然后将以下文本转换为进度任务中的todo,加入你自己的思考，而不是简单罗列: 406 target count:%s 407 %s 408 ## 返回格式(只返回json格式) 409 `+"```json\n"+` 410 [ 411 {"title":"思考","desc":"在开始扫描前，需要确保所有必要的工具和数据库都已就绪。这就像医生手术前检查器械一样重要。"}, 412 {"title":"执行1","desc":"✓ 目标锁定成功 - 识别到1个待扫描目标 ✓ 指纹库加载完成 - 已装载36种识别模式 ✓ 漏洞数据库就绪 - 涵盖394个已知漏洞特征"}, 413 {"title":"Agent反思","desc":"环境配置符合预期，所有组件状态良好。我现在已经具备了执行任务所需的全部能力。"}, 414 ] 415 `+"\n```\n", texts.initEnv, fmt.Sprintf(texts.targetCountTemplate, len(targets)), config) 416 417 if request.Language == "en" { 418 prompt += "## Return in English" 419 } 420 421 response, err := model.ChatResponse(ctx, prompt) 422 gologger.Infof("AI分析初始化配置结果: %s", response) 423 if err == nil { 424 type Item struct { 425 Title string `json:"title"` 426 Desc string `json:"desc"` 427 } 428 var items []Item 429 data := models.GetJsonString(response) 430 err = json.Unmarshal([]byte(data), &items) 431 if err != nil { 432 callbacks.StepStatusUpdateCallback(step01, statusId01, AgentStatusCompleted, texts.initConfig, "") 433 callbacks.StepStatusUpdateCallback(step01, uuid.NewString(), AgentStatusCompleted, texts.targetConfig, fmt.Sprintf(texts.targetCountTemplate, len(targets))) 434 } 435 if len(items) > 0 { 436 callbacks.StepStatusUpdateCallback(step01, statusId01, AgentStatusCompleted, items[0].Title, items[0].Desc) 437 if len(items) >= 2 { 438 for _, item := range items[1 : len(items)-1] { 439 s1 := uuid.NewString() 440 //callbacks.StepStatusUpdateCallback(step01, s1, AgentStatusRunning, "思考中", "AI思考中") 441 //time.Sleep(time.Millisecond * 600) 442 callbacks.StepStatusUpdateCallback(step01, s1, AgentStatusCompleted, item.Title, item.Desc) 443 } 444 } 445 } 446 } 447 } else { 448 callbacks.StepStatusUpdateCallback(step01, statusId01, AgentStatusCompleted, texts.initConfig, "") 449 callbacks.StepStatusUpdateCallback(step01, uuid.NewString(), AgentStatusCompleted, texts.targetConfig, fmt.Sprintf(texts.targetCountTemplate, len(targets))) 450 } 451 452 // 更新任务计划 453 tasks[0].Status = SubTaskStatusDone 454 tasks[1].Status = SubTaskStatusDoing 455 tasks[1].StartedAt = time.Now().Unix() 456 callbacks.PlanUpdateCallback(tasks) 457 458 // 步骤2: 执行扫描 459 step02 := tasks[1].StepId 460 callbacks.NewPlanStepCallback(step02, texts.execScan) 461 statusId02 := uuid.NewString() 462 callbacks.StepStatusUpdateCallback(step02, statusId02, AgentStatusCompleted, texts.aigWorking, texts.scanningDesc) 463 464 toolId02 := uuid.NewString() 465 callbacks.ToolUsedCallback(step02, statusId02, texts.execScanTool, 466 []Tool{CreateTool(toolId02, texts.aiScannerTool, ToolStatusDoing, texts.execScanDesc, texts.scanOperation, texts.targetSystem, "")}) 467 468 // 收集扫描结果 469 scanResults := make([]runner.CallbackScanResult, 0) 470 mu := sync.Mutex{} 471 472 // analysisWg waits for all concurrent AI-analysis / screenshot goroutines to finish 473 // before we compute the final score and report. 474 var analysisWg sync.WaitGroup 475 476 // analysisSem limits concurrent AI + screenshot goroutines to avoid overwhelming 477 // the model API or the screenshot service when there are many targets. 478 const maxConcurrentAnalysis = 5 479 analysisSem := make(chan struct{}, maxConcurrentAnalysis) 480 481 processFunc := func(data interface{}) { 482 switch v := data.(type) { 483 case runner.CallbackScanResult: 484 var log string 485 var appFinger string 486 if v.Fingerprint != "" { 487 appFinger = fmt.Sprintf(texts.webAppTemplate, v.Fingerprint) 488 } 489 if len(v.Vulnerabilities) > 0 { 490 log = fmt.Sprintf(texts.vulnFoundTemplate, v.TargetURL, appFinger, len(v.Vulnerabilities)) 491 } else { 492 log = fmt.Sprintf(texts.noVulnTemplate, v.TargetURL, appFinger) 493 } 494 495 callbacks.ToolUseLogCallback(toolId02, texts.aiScannerTool, step02, log) 496 497 // Run AI analysis and screenshot concurrently, outside the result mutex, 498 // so that multiple targets are processed in parallel instead of sequentially. 499 analysisWg.Add(1) 500 go func(result runner.CallbackScanResult, logMsg string) { 501 defer analysisWg.Done() 502 503 // Acquire semaphore slot 504 analysisSem <- struct{}{} 505 defer func() { <-analysisSem }() 506 507 // AI模式下的额外分析 508 if model != nil { 509 status := uuid.NewString() 510 callbacks.StepStatusUpdateCallback(step02, status, AgentStatusRunning, texts.scanResult, "AI analysis") 511 512 prompt := fmt.Sprintf("这是AI基础设施扫描的扫描结果，请你根据以下文本进行总结和归纳，你最后要补充一句(后面将调用未授权检测工具继续扫描,不需要一模一样的文字，大致意思是这样就可以):'我将进行截图分析,继续探索网页上可能的漏洞点'，扫描结果如下:\n%s\n", logMsg) 513 if request.Language == "en" { 514 prompt += "## 返回使用全英文" 515 } 516 response, _ := model.ChatResponse(context.Background(), prompt) 517 callbacks.StepStatusUpdateCallback(step02, status, AgentStatusCompleted, texts.scanResult, response) 518 } 519 520 // 截图和AI分析 521 newUuid := uuid.NewString() 522 func() { 523 callbacks.StepStatusUpdateCallback(step02, newUuid, AgentStatusRunning, "A.I.G is Thinking", "") 524 defer callbacks.StepStatusUpdateCallback(step02, newUuid, AgentStatusCompleted, "A.I.G Finished", "") 525 526 var screenshotData []byte 527 var vulInfo *vulstruct.Info 528 var summary string 529 var err error 530 531 if model != nil { 532 screenshotData, vulInfo, summary, err = runner.Analysis(result.TargetURL, result.Resp, request.Language, model) 533 if err != nil { 534 gologger.WithError(err).Errorf("AI分析失败: %v", err) 535 return 536 } 537 result.Reason = summary 538 } else { 539 screenshotData, err = runner.ScreenShot(result.TargetURL) 540 if err != nil { 541 gologger.WithError(err).Errorf("截图失败: %v", err) 542 return 543 } 544 } 545 546 if len(screenshotData) > 0 { 547 tmpPath := path.Join(os.TempDir(), fmt.Sprintf("%d.jpg", time.Now().UnixMicro())) 548 if err := os.WriteFile(tmpPath, screenshotData, 0644); err != nil { 549 gologger.WithError(err).Errorf("write file failed: %v", err) 550 return 551 } 552 info, err := utils.UploadFile(t.Server, tmpPath) 553 if err != nil { 554 gologger.WithError(err).Errorf("upload file failed: %v", err) 555 return 556 } 557 result.ScreenShot = "/api/v1/images/" + info.Data.FileUrl 558 559 if model != nil && vulInfo != nil && (vulInfo.Severity == "high" || vulInfo.Severity == "medium") { 560 result.Vulnerabilities = append(result.Vulnerabilities, *vulInfo) 561 } 562 } 563 564 // AI模式生成摘要 565 if model != nil { 566 vData, _ := json.Marshal(result.Vulnerabilities) 567 summaryPrompt := "根据以下我提供的漏洞信息，请总结一下发现x个漏洞，会导致xx业务风险，建议xx修复，几句简短的话概括，若未提供漏洞信息，就说目前暂时无漏洞发现。漏洞信息如下:\n" + string(vData) 568 if request.Language == "en" { 569 summaryPrompt += "## 返回使用全英文" 570 } 571 summary2, _ := model.ChatResponse(context.Background(), summaryPrompt) 572 result.Summary = summary2 573 } 574 }() 575 576 mu.Lock() 577 scanResults = append(scanResults, result) 578 mu.Unlock() 579 }(v, log) 580 581 case runner.CallbackErrorInfo: 582 callbacks.ToolUseLogCallback(toolId02, texts.aiScannerTool, step02, fmt.Sprintf(texts.errorTemplate, v.Target, v.Error)) 583 case runner.CallbackProcessInfo, runner.CallbackReportInfo, runner.Step01: 584 // 忽略这些类型 585 default: 586 gologger.Errorf("processFunc unknown type: %T\n", v) 587 } 588 } 589 590 opts.SetCallback(processFunc) 591 r, err := runner.New(opts) 592 if err != nil { 593 return fmt.Errorf("new runner failed: %v", err) 594 } 595 defer r.Close() 596 597 // 执行扫描 598 r.RunEnumeration() 599 600 // Wait for all concurrent AI-analysis / screenshot goroutines to finish 601 // before computing the final score and generating the report. 602 analysisWg.Wait() 603 604 // 计算安全评分 605 advies := make([]vulstruct.Info, 0) 606 for _, item := range scanResults { 607 advies = append(advies, item.Vulnerabilities...) 608 } 609 score := r.CalcSecScore(advies) 610 611 callbacks.StepStatusUpdateCallback(step02, statusId02, AgentStatusCompleted, texts.aigCompleted, texts.scanCompleteDesc) 612 callbacks.ToolUsedCallback(step02, statusId02, texts.execScanTool, 613 []Tool{CreateTool(toolId02, texts.aiScannerTool, ToolStatusDone, texts.scanComplete, texts.scanOperation, texts.targetSystem, fmt.Sprintf(texts.scanResultTemplate, len(scanResults)))}) 614 callbacks.StepStatusUpdateCallback(step02, uuid.NewString(), AgentStatusCompleted, texts.aigCompleted, texts.taskCompleteDesc) 615 616 // 更新任务计划 617 tasks[1].Status = SubTaskStatusDone 618 tasks[2].Status = SubTaskStatusDoing 619 tasks[2].StartedAt = time.Now().Unix() 620 callbacks.PlanUpdateCallback(tasks) 621 622 // 步骤3: 生成报告 623 step03 := tasks[2].StepId 624 callbacks.NewPlanStepCallback(step03, texts.genReport) 625 statustool := uuid.NewString() 626 callbacks.StepStatusUpdateCallback(step03, statustool, AgentStatusCompleted, texts.aigWorking, texts.reportGenDesc) 627 628 toolId03 := uuid.NewString() 629 callbacks.ToolUsedCallback(step03, statustool, texts.generateReportTool, 630 []Tool{CreateTool(toolId03, texts.reportGeneratorTool, ToolStatusDone, texts.reportGenToolDesc, texts.generateReport, "", fmt.Sprintf("%d", len(scanResults)))}) 631 632 // 发送最终结果 633 result := map[string]interface{}{ 634 "total": len(advies), 635 "score": score.SecScore, 636 "results": scanResults, 637 } 638 tasks[2].Status = SubTaskStatusDone 639 callbacks.PlanUpdateCallback(tasks) 640 callbacks.ResultCallback(result) 641 642 return nil 643 }