Add files via upload

2026-05-21 23:26:50 +02:00 · 2026-02-20 16:16:48 +08:00
parent 2edd5ffe95
commit dd883677b8
9 changed files with 1551 additions and 2 deletions
@@ -318,6 +318,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
 	roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger)
 	roleHandler.SetSkillsManager(skillsManager) // 设置Skills管理器到RoleHandler
 	skillsHandler := handler.NewSkillsHandler(skillsManager, cfg, configPath, log.Logger)
+	fofaHandler := handler.NewFofaHandler(cfg, log.Logger)
 	if db != nil {
 		skillsHandler.SetDB(db) // 设置数据库连接以便获取调用统计
 	}
@@ -415,6 +416,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
 		vulnerabilityHandler,
 		roleHandler,
 		skillsHandler,
+		fofaHandler,
 		mcpServer,
 		authManager,
 		openAPIHandler,
@@ -478,6 +480,7 @@ func setupRoutes(
 	vulnerabilityHandler *handler.VulnerabilityHandler,
 	roleHandler *handler.RoleHandler,
 	skillsHandler *handler.SkillsHandler,
+	fofaHandler *handler.FofaHandler,
 	mcpServer *mcp.Server,
 	authManager *security.AuthManager,
 	openAPIHandler *handler.OpenAPIHandler,
@@ -506,6 +509,9 @@ func setupRoutes(
 		protected.GET("/agent-loop/tasks", agentHandler.ListAgentTasks)
 		protected.GET("/agent-loop/tasks/completed", agentHandler.ListCompletedTasks)

+		// 信息收集 - FOFA 查询（后端代理）
+		protected.POST("/fofa/search", fofaHandler.Search)
+
 		// 批量任务管理
 		protected.POST("/batch-tasks", agentHandler.CreateBatchQueue)
 		protected.GET("/batch-tasks", agentHandler.ListBatchQueues)
@@ -18,6 +18,7 @@ type Config struct {
 	Log         LogConfig             `yaml:"log"`
 	MCP         MCPConfig             `yaml:"mcp"`
 	OpenAI      OpenAIConfig          `yaml:"openai"`
+	FOFA        FofaConfig            `yaml:"fofa,omitempty" json:"fofa,omitempty"`
 	Agent       AgentConfig           `yaml:"agent"`
 	Security    SecurityConfig        `yaml:"security"`
 	Database    DatabaseConfig        `yaml:"database"`
@@ -52,6 +53,13 @@ type OpenAIConfig struct {
 	MaxTotalTokens int    `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"`
 }

+type FofaConfig struct {
+	// Email 为 FOFA 账号邮箱；APIKey 为 FOFA API Key（建议使用只读权限的 Key）
+	Email   string `yaml:"email,omitempty" json:"email,omitempty"`
+	APIKey  string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
+	BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://fofa.info/api/v1/search/all
+}
+
 type SecurityConfig struct {
 	Tools               []ToolConfig `yaml:"tools,omitempty"`                 // 向后兼容：支持在主配置文件中定义工具
 	ToolsDir            string       `yaml:"tools_dir,omitempty"`             // 工具配置文件目录（新方式）
@@ -145,6 +145,7 @@ func (h *ConfigHandler) SetAppUpdater(updater AppUpdater) {
 // GetConfigResponse 获取配置响应
 type GetConfigResponse struct {
 	OpenAI    config.OpenAIConfig    `json:"openai"`
+	FOFA      config.FofaConfig      `json:"fofa"`
 	MCP       config.MCPConfig       `json:"mcp"`
 	Tools     []ToolConfigInfo       `json:"tools"`
 	Agent     config.AgentConfig     `json:"agent"`
@@ -216,6 +217,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {

 	c.JSON(http.StatusOK, GetConfigResponse{
 		OpenAI:    h.config.OpenAI,
+		FOFA:      h.config.FOFA,
 		MCP:       h.config.MCP,
 		Tools:     tools,
 		Agent:     h.config.Agent,
@@ -472,6 +474,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
 // UpdateConfigRequest 更新配置请求
 type UpdateConfigRequest struct {
 	OpenAI    *config.OpenAIConfig    `json:"openai,omitempty"`
+	FOFA      *config.FofaConfig      `json:"fofa,omitempty"`
 	MCP       *config.MCPConfig       `json:"mcp,omitempty"`
 	Tools     []ToolEnableStatus      `json:"tools,omitempty"`
 	Agent     *config.AgentConfig     `json:"agent,omitempty"`
@@ -506,6 +509,12 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
 		)
 	}

+	// 更新FOFA配置
+	if req.FOFA != nil {
+		h.config.FOFA = *req.FOFA
+		h.logger.Info("更新FOFA配置", zap.String("email", h.config.FOFA.Email))
+	}
+
 	// 更新MCP配置
 	if req.MCP != nil {
 		h.config.MCP = *req.MCP
@@ -0,0 +1,223 @@
+package handler
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+	"os"
+	"strings"
+	"time"
+
+	"cyberstrike-ai/internal/config"
+
+	"github.com/gin-gonic/gin"
+	"go.uber.org/zap"
+)
+
+type FofaHandler struct {
+	cfg    *config.Config
+	logger *zap.Logger
+	client *http.Client
+}
+
+func NewFofaHandler(cfg *config.Config, logger *zap.Logger) *FofaHandler {
+	return &FofaHandler{
+		cfg:    cfg,
+		logger: logger,
+		client: &http.Client{Timeout: 30 * time.Second},
+	}
+}
+
+type fofaSearchRequest struct {
+	Query  string `json:"query" binding:"required"`
+	Size   int    `json:"size,omitempty"`
+	Page   int    `json:"page,omitempty"`
+	Fields string `json:"fields,omitempty"`
+	Full   bool   `json:"full,omitempty"`
+}
+
+type fofaAPIResponse struct {
+	Error   bool            `json:"error"`
+	ErrMsg  string          `json:"errmsg"`
+	Size    int             `json:"size"`
+	Page    int             `json:"page"`
+	Total   int             `json:"total"`
+	Mode    string          `json:"mode"`
+	Query   string          `json:"query"`
+	Results [][]interface{} `json:"results"`
+}
+
+type fofaSearchResponse struct {
+	Query        string                   `json:"query"`
+	Size         int                      `json:"size"`
+	Page         int                      `json:"page"`
+	Total        int                      `json:"total"`
+	Fields       []string                 `json:"fields"`
+	ResultsCount int                      `json:"results_count"`
+	Results      []map[string]interface{} `json:"results"`
+}
+
+func (h *FofaHandler) resolveCredentials() (email, apiKey string) {
+	// 优先环境变量（便于容器部署），其次配置文件
+	email = strings.TrimSpace(os.Getenv("FOFA_EMAIL"))
+	apiKey = strings.TrimSpace(os.Getenv("FOFA_API_KEY"))
+	if email != "" && apiKey != "" {
+		return email, apiKey
+	}
+	if h.cfg != nil {
+		if email == "" {
+			email = strings.TrimSpace(h.cfg.FOFA.Email)
+		}
+		if apiKey == "" {
+			apiKey = strings.TrimSpace(h.cfg.FOFA.APIKey)
+		}
+	}
+	return email, apiKey
+}
+
+func (h *FofaHandler) resolveBaseURL() string {
+	if h.cfg != nil {
+		if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" {
+			return v
+		}
+	}
+	return "https://fofa.info/api/v1/search/all"
+}
+
+// Search FOFA 查询（后端代理，避免前端暴露 key）
+func (h *FofaHandler) Search(c *gin.Context) {
+	var req fofaSearchRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
+		return
+	}
+
+	req.Query = strings.TrimSpace(req.Query)
+	if req.Query == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "query 不能为空"})
+		return
+	}
+	if req.Size <= 0 {
+		req.Size = 100
+	}
+	if req.Page <= 0 {
+		req.Page = 1
+	}
+	// FOFA 接口 size 上限和账户权限相关，这里只做一个合理的保护
+	if req.Size > 10000 {
+		req.Size = 10000
+	}
+	if req.Fields == "" {
+		req.Fields = "host,ip,port,domain,title,protocol,country,province,city,server"
+	}
+
+	email, apiKey := h.resolveCredentials()
+	if email == "" || apiKey == "" {
+		c.JSON(http.StatusBadRequest, gin.H{
+			"error":   "FOFA 未配置：请在系统设置中填写 FOFA Email/API Key，或设置环境变量 FOFA_EMAIL/FOFA_API_KEY",
+			"need":    []string{"fofa.email", "fofa.api_key"},
+			"env_key": []string{"FOFA_EMAIL", "FOFA_API_KEY"},
+		})
+		return
+	}
+
+	baseURL := h.resolveBaseURL()
+	qb64 := base64.StdEncoding.EncodeToString([]byte(req.Query))
+
+	u, err := url.Parse(baseURL)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "FOFA base_url 无效: " + err.Error()})
+		return
+	}
+
+	params := u.Query()
+	params.Set("email", email)
+	params.Set("key", apiKey)
+	params.Set("qbase64", qb64)
+	params.Set("size", fmt.Sprintf("%d", req.Size))
+	params.Set("page", fmt.Sprintf("%d", req.Page))
+	params.Set("fields", strings.TrimSpace(req.Fields))
+	if req.Full {
+		params.Set("full", "true")
+	} else {
+		// 明确传 false，便于排查
+		params.Set("full", "false")
+	}
+	u.RawQuery = params.Encode()
+
+	httpReq, err := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, u.String(), nil)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()})
+		return
+	}
+
+	resp, err := h.client.Do(httpReq)
+	if err != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "请求 FOFA 失败: " + err.Error()})
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("FOFA 返回非 2xx: %d", resp.StatusCode)})
+		return
+	}
+
+	var apiResp fofaAPIResponse
+	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "解析 FOFA 响应失败: " + err.Error()})
+		return
+	}
+	if apiResp.Error {
+		msg := strings.TrimSpace(apiResp.ErrMsg)
+		if msg == "" {
+			msg = "FOFA 返回错误"
+		}
+		c.JSON(http.StatusBadGateway, gin.H{"error": msg})
+		return
+	}
+
+	fields := splitAndCleanCSV(req.Fields)
+	results := make([]map[string]interface{}, 0, len(apiResp.Results))
+	for _, row := range apiResp.Results {
+		item := make(map[string]interface{}, len(fields))
+		for i, f := range fields {
+			if i < len(row) {
+				item[f] = row[i]
+			} else {
+				item[f] = nil
+			}
+		}
+		results = append(results, item)
+	}
+
+	c.JSON(http.StatusOK, fofaSearchResponse{
+		Query:        req.Query,
+		Size:         apiResp.Size,
+		Page:         apiResp.Page,
+		Total:        apiResp.Total,
+		Fields:       fields,
+		ResultsCount: len(results),
+		Results:      results,
+	})
+}
+
+func splitAndCleanCSV(s string) []string {
+	parts := strings.Split(s, ",")
+	out := make([]string, 0, len(parts))
+	seen := make(map[string]struct{}, len(parts))
+	for _, p := range parts {
+		v := strings.TrimSpace(p)
+		if v == "" {
+			continue
+		}
+		if _, ok := seen[v]; ok {
+			continue
+		}
+		seen[v] = struct{}{}
+		out = append(out, v)
+	}
+	return out
+}