Add files via upload

2026-05-21 23:26:50 +02:00 · 2025-11-08 22:58:51 +08:00
parent ae1f2531c9
commit 8d8bcbfaac
5 changed files with 943 additions and 257 deletions
@@ -271,8 +271,11 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
 		MCPExecutionIDs: make([]string, 0),
 	}

-	maxIterations := 10
+	maxIterations := 30
 	for i := 0; i < maxIterations; i++ {
+		// 检查是否是最后一次迭代
+		isLastIteration := (i == maxIterations-1)
+		
 		// 获取可用工具
 		tools := a.getAvailableTools()

@@ -282,6 +285,12 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
 				"iteration": i + 1,
 				"total":     maxIterations,
 			})
+		} else if isLastIteration {
+			sendProgress("iteration", fmt.Sprintf("第 %d 轮迭代（最后一次）", i+1), map[string]interface{}{
+				"iteration": i + 1,
+				"total":     maxIterations,
+				"isLast":    true,
+			})
 		} else {
 			sendProgress("iteration", fmt.Sprintf("第 %d 轮迭代", i+1), map[string]interface{}{
 				"iteration": i + 1,
@@ -439,6 +448,29 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
 					}
 				}
 			}
+			
+			// 如果是最后一次迭代，执行完工具后要求AI进行总结
+			if isLastIteration {
+				sendProgress("progress", "最后一次迭代：正在生成总结和下一步计划...", nil)
+				// 添加用户消息，要求AI进行总结
+				messages = append(messages, ChatMessage{
+					Role:    "user",
+					Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试，请提供详细的下一步执行计划。请直接回复，不要调用工具。",
+				})
+				// 立即调用OpenAI获取总结
+				summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具，强制AI直接回复
+				if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
+					summaryChoice := summaryResponse.Choices[0]
+					if summaryChoice.Message.Content != "" {
+						result.Response = summaryChoice.Message.Content
+						sendProgress("progress", "总结生成完成", nil)
+						return result, nil
+					}
+				}
+				// 如果获取总结失败，跳出循环，让后续逻辑处理
+				break
+			}
+			
 			continue
 		}

@@ -455,6 +487,33 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
 			})
 		}

+		// 如果是最后一次迭代，无论finish_reason是什么，都要求AI进行总结
+		if isLastIteration {
+			sendProgress("progress", "最后一次迭代：正在生成总结和下一步计划...", nil)
+			// 添加用户消息，要求AI进行总结
+			messages = append(messages, ChatMessage{
+				Role:    "user",
+				Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试，请提供详细的下一步执行计划。请直接回复，不要调用工具。",
+			})
+			// 立即调用OpenAI获取总结
+			summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具，强制AI直接回复
+			if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
+				summaryChoice := summaryResponse.Choices[0]
+				if summaryChoice.Message.Content != "" {
+					result.Response = summaryChoice.Message.Content
+					sendProgress("progress", "总结生成完成", nil)
+					return result, nil
+				}
+			}
+			// 如果获取总结失败，使用当前回复作为结果
+			if choice.Message.Content != "" {
+				result.Response = choice.Message.Content
+				return result, nil
+			}
+			// 如果都没有内容，跳出循环，让后续逻辑处理
+			break
+		}
+		
 		// 如果完成，返回结果
 		if choice.FinishReason == "stop" {
 			sendProgress("progress", "正在生成最终回复...", nil)
@@ -463,7 +522,27 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
 		}
 	}

-	result.Response = "达到最大迭代次数"
+	// 如果循环结束仍未返回，说明达到了最大迭代次数
+	// 尝试最后一次调用AI获取总结
+	sendProgress("progress", "达到最大迭代次数，正在生成总结...", nil)
+	finalSummaryPrompt := ChatMessage{
+		Role:    "user",
+		Content: "已达到最大迭代次数（30轮）。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试，请提供详细的下一步执行计划。请直接回复，不要调用工具。",
+	}
+	messages = append(messages, finalSummaryPrompt)
+	
+	summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具，强制AI直接回复
+	if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
+		summaryChoice := summaryResponse.Choices[0]
+		if summaryChoice.Message.Content != "" {
+			result.Response = summaryChoice.Message.Content
+			sendProgress("progress", "总结生成完成", nil)
+			return result, nil
+		}
+	}
+	
+	// 如果无法生成总结，返回友好的提示
+	result.Response = "已达到最大迭代次数（30轮）。系统已执行了多轮测试，但由于达到迭代上限，无法继续自动执行。建议您查看已执行的工具结果，或提出新的测试请求以继续测试。"
 	return result, nil
 }

@@ -317,95 +317,21 @@ func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConf
 		return cmdArgs
 	}

-	// 向后兼容：如果没有定义参数，使用旧的硬编码逻辑
-	switch toolName {
-	case "nmap":
-		// nmap -sT -sV -sC target [ports]
-		// 使用 -sT (TCP连接扫描) 而不是 -sS (SYN扫描)，因为 -sS 需要root权限
-		e.logger.Debug("处理nmap参数",
-			zap.Any("args", args),
-		)
-		
-		// 尝试多种方式获取target参数
-		var target string
-		var ok bool
-		
-		// 方式1: 直接获取target
-		if target, ok = args["target"].(string); !ok || target == "" {
-			// 方式2: 尝试从tool字段获取（兼容某些格式）
-			if toolVal, exists := args["tool"]; exists {
-				if toolMap, ok := toolVal.(map[string]interface{}); ok {
-					if t, ok := toolMap["target"].(string); ok {
-						target = t
-					}
-				}
-			}
+	// 如果没有定义参数配置，使用固定参数和通用处理
+	// 添加固定参数
+	cmdArgs = append(cmdArgs, toolConfig.Args...)
+	
+	// 通用处理：将参数转换为命令行参数
+	for key, value := range args {
+		if key == "_tool_name" {
+			continue
 		}
-		
-		if target == "" {
-			e.logger.Warn("nmap缺少target参数",
-				zap.Any("args", args),
-			)
-			return cmdArgs // 返回空数组，让上层处理错误
-		}
-		
-		e.logger.Debug("提取到target",
-			zap.String("target", target),
-		)
-		
-		// 处理URL格式的目标（提取域名）
-		if strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://") {
-			// 提取域名部分
-			target = strings.TrimPrefix(target, "http://")
-			target = strings.TrimPrefix(target, "https://")
-			// 移除路径部分
-			if idx := strings.Index(target, "/"); idx != -1 {
-				target = target[:idx]
-			}
-		}
-		
-		// 添加扫描选项：-sT (TCP连接扫描，不需要root权限), -sV (版本检测), -sC (默认脚本)
-		cmdArgs = append(cmdArgs, "-sT", "-sV", "-sC")
-		
-		// 添加端口范围（如果指定）
-		if ports, ok := args["ports"].(string); ok && ports != "" {
-			cmdArgs = append(cmdArgs, "-p", ports)
-		}
-		
-		// 添加目标
-		cmdArgs = append(cmdArgs, target)
-		
-		e.logger.Debug("nmap命令参数构建完成",
-			zap.Strings("cmdArgs", cmdArgs),
-		)
-	case "sqlmap":
-		// sqlmap -u url
-		if url, ok := args["url"].(string); ok {
-			cmdArgs = append(cmdArgs, "-u", url, "--batch", "--level=3", "--risk=2")
-		}
-	case "nikto":
-		// nikto -h target
-		if target, ok := args["target"].(string); ok {
-			cmdArgs = append(cmdArgs, "-h", target)
-		}
-	case "dirb":
-		// dirb url
-		if url, ok := args["url"].(string); ok {
-			cmdArgs = append(cmdArgs, url)
-		}
-	default:
-		// 通用处理
-		cmdArgs = append(cmdArgs, toolConfig.Args...)
-		for key, value := range args {
-			if key == "_tool_name" {
-				continue
-			}
-			cmdArgs = append(cmdArgs, fmt.Sprintf("--%s", key))
-			if strValue, ok := value.(string); ok {
-				cmdArgs = append(cmdArgs, strValue)
-			} else {
-				cmdArgs = append(cmdArgs, fmt.Sprintf("%v", value))
-			}
+		// 使用 --key value 格式
+		cmdArgs = append(cmdArgs, fmt.Sprintf("--%s", key))
+		if strValue, ok := value.(string); ok {
+			cmdArgs = append(cmdArgs, strValue)
+		} else {
+			cmdArgs = append(cmdArgs, fmt.Sprintf("%v", value))
 		}
 	}

@@ -592,54 +518,12 @@ func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]in
 		return schema
 	}

-	// 向后兼容：如果没有定义参数，使用旧的硬编码逻辑
-	switch toolConfig.Name {
-	case "nmap":
-		schema["properties"] = map[string]interface{}{
-			"target": map[string]interface{}{
-				"type":        "string",
-				"description": "目标IP地址或域名",
-			},
-			"ports": map[string]interface{}{
-				"type":        "string",
-				"description": "端口范围，例如: 1-1000",
-			},
-		}
-		schema["required"] = []string{"target"}
-	case "sqlmap":
-		schema["properties"] = map[string]interface{}{
-			"url": map[string]interface{}{
-				"type":        "string",
-				"description": "目标URL",
-			},
-		}
-		schema["required"] = []string{"url"}
-	case "nikto", "dirb":
-		schema["properties"] = map[string]interface{}{
-			"target": map[string]interface{}{
-				"type":        "string",
-				"description": "目标URL",
-			},
-		}
-		schema["required"] = []string{"target"}
-	case "exec":
-		schema["properties"] = map[string]interface{}{
-			"command": map[string]interface{}{
-				"type":        "string",
-				"description": "要执行的系统命令",
-			},
-			"shell": map[string]interface{}{
-				"type":        "string",
-				"description": "使用的shell（可选，默认为sh）",
-			},
-			"workdir": map[string]interface{}{
-				"type":        "string",
-				"description": "工作目录（可选）",
-			},
-		}
-		schema["required"] = []string{"command"}
-	}
-
+	// 如果没有定义参数配置，返回空schema
+	// 这种情况下工具可能只使用固定参数（args字段）
+	// 或者需要通过YAML配置文件定义参数
+	e.logger.Warn("工具未定义参数配置，返回空schema",
+		zap.String("tool", toolConfig.Name),
+	)
 	return schema
 }