From 9bafc2ab980aae911b7c48a97dbbadb25ce7738b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:31:20 +0800 Subject: [PATCH] Add files via upload --- internal/handler/config.go | 53 ++- internal/handler/fofa.go | 837 ++++++++++++++++++++++++++++------ internal/handler/fofa_test.go | 113 +++++ 3 files changed, 847 insertions(+), 156 deletions(-) diff --git a/internal/handler/config.go b/internal/handler/config.go index 6a1fc776..133f7400 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -258,17 +258,20 @@ func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) err // GetConfigResponse 获取配置响应 type GetConfigResponse struct { - OpenAI config.OpenAIConfig `json:"openai"` - Vision config.VisionConfig `json:"vision"` - FOFA config.FofaConfig `json:"fofa"` - MCP config.MCPConfig `json:"mcp"` - Tools []ToolConfigInfo `json:"tools"` - Agent config.AgentConfig `json:"agent"` - Hitl config.HitlConfig `json:"hitl,omitempty"` - Knowledge config.KnowledgeConfig `json:"knowledge"` - Robots config.RobotsConfig `json:"robots,omitempty"` - MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"` - C2 config.C2Public `json:"c2"` + OpenAI config.OpenAIConfig `json:"openai"` + Vision config.VisionConfig `json:"vision"` + FOFA config.FofaConfig `json:"fofa"` + ZoomEye config.SpaceSearchConfig `json:"zoomeye"` + Quake config.SpaceSearchConfig `json:"quake"` + Shodan config.SpaceSearchConfig `json:"shodan"` + MCP config.MCPConfig `json:"mcp"` + Tools []ToolConfigInfo `json:"tools"` + Agent config.AgentConfig `json:"agent"` + Hitl config.HitlConfig `json:"hitl,omitempty"` + Knowledge config.KnowledgeConfig `json:"knowledge"` + Robots config.RobotsConfig `json:"robots,omitempty"` + MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"` + C2 config.C2Public `json:"c2"` } // ToolConfigInfo 工具配置信息 @@ -363,6 +366,9 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) { OpenAI: h.config.OpenAI, Vision: h.config.Vision, FOFA: h.config.FOFA, + ZoomEye: h.config.ZoomEye, + Quake: h.config.Quake, + Shodan: h.config.Shodan, MCP: h.config.MCP, Tools: tools, Agent: h.config.Agent, @@ -703,6 +709,9 @@ type UpdateConfigRequest struct { OpenAI *config.OpenAIConfig `json:"openai,omitempty"` Vision *config.VisionConfig `json:"vision,omitempty"` FOFA *config.FofaConfig `json:"fofa,omitempty"` + ZoomEye *config.SpaceSearchConfig `json:"zoomeye,omitempty"` + Quake *config.SpaceSearchConfig `json:"quake,omitempty"` + Shodan *config.SpaceSearchConfig `json:"shodan,omitempty"` MCP *config.MCPConfig `json:"mcp,omitempty"` Tools []ToolEnableStatus `json:"tools,omitempty"` Agent *AgentConfigUpdate `json:"agent,omitempty"` @@ -777,6 +786,18 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) { h.config.FOFA = *req.FOFA h.logger.Info("更新FOFA配置", zap.String("base_url", h.config.FOFA.BaseURL)) } + if req.ZoomEye != nil { + h.config.ZoomEye = *req.ZoomEye + h.logger.Info("更新ZoomEye配置", zap.String("base_url", h.config.ZoomEye.BaseURL)) + } + if req.Quake != nil { + h.config.Quake = *req.Quake + h.logger.Info("更新Quake配置", zap.String("base_url", h.config.Quake.BaseURL)) + } + if req.Shodan != nil { + h.config.Shodan = *req.Shodan + h.logger.Info("更新Shodan配置", zap.String("base_url", h.config.Shodan.BaseURL)) + } // 更新MCP配置 if req.MCP != nil { @@ -1601,6 +1622,9 @@ func (h *ConfigHandler) saveConfig() error { updateOpenAIConfig(root, h.config.OpenAI) updateVisionConfig(root, h.config.Vision) updateFOFAConfig(root, h.config.FOFA) + updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye) + updateSpaceSearchConfig(root, "quake", h.config.Quake) + updateSpaceSearchConfig(root, "shodan", h.config.Shodan) updateKnowledgeConfig(root, h.config.Knowledge) updateC2Config(root, h.config.C2) updateRobotsConfig(root, h.config.Robots) @@ -1792,6 +1816,13 @@ func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) { setStringInMap(fofaNode, "api_key", cfg.APIKey) } +func updateSpaceSearchConfig(doc *yaml.Node, key string, cfg config.SpaceSearchConfig) { + root := doc.Content[0] + node := ensureMap(root, key) + setStringInMap(node, "base_url", cfg.BaseURL) + setStringInMap(node, "api_key", cfg.APIKey) +} + func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) { root := doc.Content[0] knowledgeNode := ensureMap(root, "knowledge") diff --git a/internal/handler/fofa.go b/internal/handler/fofa.go index e3eae38b..bea0494b 100644 --- a/internal/handler/fofa.go +++ b/internal/handler/fofa.go @@ -43,15 +43,17 @@ func NewFofaHandler(cfg *config.Config, logger *zap.Logger) *FofaHandler { } 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"` + Provider string `json:"provider,omitempty"` + 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 fofaParseRequest struct { - Text string `json:"text" binding:"required"` + Provider string `json:"provider,omitempty"` + Text string `json:"text" binding:"required"` } type fofaParseResponse struct { @@ -72,33 +74,118 @@ type fofaAPIResponse struct { } 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"` + Provider string `json:"provider,omitempty"` + 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"` + ExpectedCount int `json:"expected_count,omitempty"` + Shortfall int `json:"shortfall,omitempty"` + Warning string `json:"warning,omitempty"` + Results []map[string]interface{} `json:"results"` } -func (h *FofaHandler) resolveAPIKey() string { +func normalizeSpaceSearchProvider(provider string) string { + switch strings.ToLower(strings.TrimSpace(provider)) { + case "", "fofa": + return "fofa" + case "zoomeye", "zoom-eye": + return "zoomeye" + case "quake": + return "quake" + case "shodan": + return "shodan" + default: + return "" + } +} + +func providerDisplayName(provider string) string { + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "ZoomEye" + case "quake": + return "Quake" + case "shodan": + return "Shodan" + default: + return "FOFA" + } +} + +func defaultFieldsForProvider(provider string) string { + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "ip,port,domain,hostname,title,service,app,country,city" + case "quake": + return "ip,port,domain,service.name,service.http.title,location.country_cn,location.province_cn,location.city_cn" + case "shodan": + return "ip_str,port,hostnames,domains,org,isp,location.country_name,location.city,product,transport" + default: + return "host,ip,port,domain,title,protocol,country,province,city,server" + } +} + +func (h *FofaHandler) resolveAPIKey(provider string) string { // 优先环境变量（便于容器部署），其次配置文件。 - if apiKey := strings.TrimSpace(os.Getenv("FOFA_API_KEY")); apiKey != "" { + provider = normalizeSpaceSearchProvider(provider) + envKey := map[string]string{ + "fofa": "FOFA_API_KEY", + "zoomeye": "ZOOMEYE_API_KEY", + "quake": "QUAKE_API_KEY", + "shodan": "SHODAN_API_KEY", + }[provider] + if apiKey := strings.TrimSpace(os.Getenv(envKey)); apiKey != "" { return apiKey } if h.cfg != nil { - return strings.TrimSpace(h.cfg.FOFA.APIKey) + switch provider { + case "zoomeye": + return strings.TrimSpace(h.cfg.ZoomEye.APIKey) + case "quake": + return strings.TrimSpace(h.cfg.Quake.APIKey) + case "shodan": + return strings.TrimSpace(h.cfg.Shodan.APIKey) + default: + return strings.TrimSpace(h.cfg.FOFA.APIKey) + } } return "" } -func (h *FofaHandler) resolveBaseURL() string { +func (h *FofaHandler) resolveBaseURL(provider string) string { if h.cfg != nil { - if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" { - return v + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + if v := strings.TrimSpace(h.cfg.ZoomEye.BaseURL); v != "" { + return v + } + case "quake": + if v := strings.TrimSpace(h.cfg.Quake.BaseURL); v != "" { + return v + } + case "shodan": + if v := strings.TrimSpace(h.cfg.Shodan.BaseURL); v != "" { + return v + } + default: + if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" { + return v + } } } - return "https://fofa.info/api/v1/search/all" + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "https://api.zoomeye.org/v2/search" + case "quake": + return "https://quake.360.cn/api/v3/search/quake_service" + case "shodan": + return "https://api.shodan.io" + default: + return "https://fofa.info/api/v1/search/all" + } } // ParseNaturalLanguage 将自然语言解析为 FOFA 查询语法（仅生成，不执行查询） @@ -109,6 +196,11 @@ func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) { return } req.Text = strings.TrimSpace(req.Text) + provider := normalizeSpaceSearchProvider(req.Provider) + if provider == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "provider 不支持，可选：fofa、zoomeye、quake、shodan"}) + return + } if req.Text == "" { c.JSON(http.StatusBadRequest, gin.H{"error": "text 不能为空"}) return @@ -130,129 +222,158 @@ func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) { return } - systemPrompt := strings.TrimSpace(` -你是“FOFA 查询语法生成器”。任务：把用户输入的自然语言搜索意图，转换成 FOFA 查询语法。 + engineName := providerDisplayName(provider) + syntaxNotes := map[string]string{ + "fofa": ` +FOFA 官方查询语法参考： +- 基本格式：field="value"，字符串值使用英文双引号；多个条件用 &&（与）、||（或）、!（非）连接。 +- 组合优先级：复杂表达式必须使用 () 明确优先级，例如：(app="Apache" || app="nginx") && country="CN"。 +- 常用字段：app、title、body、header、host、domain、ip、port、protocol、country、province、city、server、icp、cert、icon_hash、fid。 +- 字段示例： + - app="Apache" + - title="后台管理" + - body="Powered by" + - header="JSESSIONID" + - domain="example.com" + - host="https://example.com" + - ip="1.1.1.1" + - port="443" + - country="CN" + - city="Hangzhou" + - cert="example.com" + - icon_hash="-247388890" +- 组合示例： + - app="Apache" && country="CN" + - title="login" || title="登录" + - (app="Apache" || app="nginx") && port="443" + - domain="example.com" && !title="404" + - cert="example.com" && port="443" + - header="JSESSIONID" && country="CN" +- 生成注意： + - 用户说“排除/不要/非”时优先使用 !field="value"。 + - 用户说“标题包含/页面标题”映射为 title；说“正文包含/页面包含”映射为 body；说“响应头/cookie/header”映射为 header。 + - 端口在 FOFA 中通常写成 port="443"。 +`, + "zoomeye": ` +ZoomEye 查询语法参考： +- 基本格式：field="value" 或 field=value；字符串/短语建议使用英文双引号。 +- 逻辑连接：可使用 && / || / !，也可使用 AND / OR / NOT；复杂表达式使用 () 明确优先级。 +- 常用字段：app、service、title、domain、hostname、ip、port、country、city、org、isp、asn、cidr、ssl、ssl.cert.fingerprint、iconhash。 +- 字段示例： + - app="Apache" + - service="ssh" + - title="登录" + - domain="example.com" + - hostname="example.com" + - ip="1.1.1.1" + - cidr="1.1.1.0/24" + - port=443 + - country="CN" + - city="Beijing" + - org="Tencent" + - ssl="example.com" + - ssl.cert.fingerprint="F3C98F223D82CC41CF83D94671CCC6C69873FABF" + - iconhash="-247388890" +- 组合示例： + - app="nginx" && country="CN" + - service="http" && (title="login" || title="登录") + - domain="example.com" && !app="cloudflare" + - port=443 && country="US" + - app="Elasticsearch" && port=9200 +- 生成注意： + - 用户说“服务/协议是 SSH、HTTP、RDP”优先映射为 service。 + - 用户说“站点/网站标题”映射为 title；说“域名/主域”优先映射为 domain 或 hostname。 + - 端口可以不加引号，例如 port=443；如果用户原文已给出冒号风格表达式且接近 ZoomEye 语法，可原样保留。 +`, + "quake": ` +Quake 查询语法参考： +- 基本格式：field:"value" 或 field:value；字符串/中文/短语使用英文双引号。 +- 逻辑连接：使用 AND、OR、NOT；复杂表达式必须使用 () 明确优先级。 +- 常用字段：domain、ip、port、service.name、service.http.title、service.http.server、service.http.response.header、service.http.favicon.hash、country_cn、province_cn、city_cn、location.country_cn、location.province_cn、location.city_cn、asn、org。 +- 字段示例： + - domain:"example.com" + - ip:"1.1.1.1" + - port:443 + - service.name:"http" + - service.name:"ssh" + - service.http.title:"登录" + - service.http.server:"nginx" + - service.http.response.header:"JSESSIONID" + - service.http.favicon.hash:"-247388890" + - country_cn:"中国" + - province_cn:"浙江" + - city_cn:"杭州" +- 组合示例： + - service.name:"http" AND country_cn:"中国" + - (service.name:"http" OR service.name:"https") AND port:443 + - domain:"example.com" AND NOT service.http.title:"404" + - service.http.title:"login" AND port:443 + - service.name:"ssh" AND country_cn:"中国" +- 生成注意： + - 用户说“中国/浙江/杭州”等中文地理位置时，Quake 优先使用 country_cn/province_cn/city_cn 并保留中文值。 + - 用户说“标题”映射为 service.http.title；说“Server/服务端软件”映射为 service.http.server；说“favicon/hash/icon”映射为 service.http.favicon.hash。 + - Quake 不使用 && / || 作为首选输出；优先输出 AND / OR / NOT。 +`, + "shodan": ` +Shodan 官方查询语法参考： +- 默认裸关键词只搜索 banner 的 data 内容；精确条件使用 filter:value。 +- filter 与 value 中间不能有空格；值包含空格时用英文双引号，例如 org:"Amazon Web Services"。 +- 多个过滤器并列表示同时满足（AND）；Shodan 查询不要使用 &&、||，除非用户明确给出并要求保留。 +- 常用过滤器：product、port、country、city、org、asn、hostname、net、ssl、ssl.cert.subject.cn、http.title、has_screenshot、vuln。 +- 字段示例： + - product:nginx + - port:443 + - country:CN + - city:Shanghai + - org:"Amazon" + - asn:AS15169 + - hostname:example.com + - ssl.cert.subject.cn:example.com + - http.title:"Dashboard" + - has_screenshot:true + - vuln:CVE-2021-41773 +- 组合示例： + - product:nginx country:CN + - apache country:DE + - org:"Amazon" port:443 + - ssl.cert.subject.cn:example.com port:443 + - http.title:"login" country:CN + - ssl:true port:443 hostname:example.com +- 生成注意： + - 用户说“产品/组件/服务软件”优先映射为 product；说“组织/公司/云厂商”映射为 org；说“证书 CN/SAN/域名证书”优先映射为 ssl.cert.subject.cn。 + - 国家用两位国家代码；如果用户给出中文国家名且无法确定代码，把推断写入 explanation 或 warnings。 + - Shodan 没有通用 NOT 排除语法；遇到“排除/不要”时应在 warnings 说明可能需要人工调整，不要强行编造过滤器。 +`, + }[provider] + + systemPrompt := strings.TrimSpace(fmt.Sprintf(` +你是“%s 查询语法生成器”。任务：把用户输入的自然语言搜索意图，转换成 %s 查询语法。 输出要求（非常重要）： 1) 只输出 JSON（不要 markdown、不要代码块、不要额外解释文本） 2) JSON 结构必须是： { - "query": "string，FOFA查询语法（可直接粘贴到 FOFA 或本系统查询框）", + "query": "string，%s 查询语法（可直接粘贴到 %s 或本系统查询框）", "explanation": "string，可选，解释你如何映射字段/逻辑", "warnings": ["string"...] 可选，列出歧义/风险/需要人工确认的点 } -3) 如果用户输入本身已经是 FOFA 查询语法（或非常接近 FOFA 语法的表达式），应当“原样返回”为 query： +3) 如果用户输入本身已经是 %s 查询语法（或非常接近该语法的表达式），应当“原样返回”为 query： - 不要擅自改写字段名、操作符、括号结构 - 不要改写任何字符串值（尤其是地理位置类值），不要做缩写/同义词替换/翻译/音译 -查询语法要点（来自 FOFA 语法参考）： -- 逻辑连接符：&&（与）、||（或），必要时用 () 包住子表达式以确认优先级（括号优先级最高） -- 当同一层级同时出现 && 与 ||（混用）时，用 () 明确优先级（避免歧义） -- 比较/匹配： - - = 匹配；当字段="" 时，可查询“不存在该字段”或“值为空”的情况 - - == 完全匹配；当字段=="" 时，可查询“字段存在且值为空”的情况 - - != 不匹配；当字段!="" 时，可查询“值不为空”的情况 - - *= 模糊匹配；可使用 * 或 ? 进行搜索 -- 直接输入关键词（不带字段）会在标题、HTML内容、HTTP头、URL字段中搜索；但当意图明确时优先用字段表达（更可控、更准确） +当前搜索引擎语法速查： +%s -字段示例速查（来自用户提供的案例，可直接套用/拼接）： -- 高级搜索操作符示例： - - title="beijing" （= 匹配） - - title=="" （== 完全匹配，字段存在且值为空） - - title="" （= 匹配，可能表示字段不存在或值为空） - - title!="" （!= 不匹配，可用于值不为空） - - title*="*Home*" （*= 模糊匹配，用 * 或 ?） - - (app="Apache" || app="Nginx") && country="CN" （混用 && / || 时用括号） -- 基础类（General）： - - ip="1.1.1.1" - - ip="220.181.111.1/24" - - ip="2600:9000:202a:2600:18:4ab7:f600:93a1" - - port="6379" - - domain="qq.com" - - host=".fofa.info" - - os="centos" - - server="Microsoft-IIS/10" - - asn="19551" - - org="LLC Baxet" - - is_domain=true / is_domain=false - - is_ipv6=true / is_ipv6=false -- 标记类（Special Label）： - - app="Microsoft-Exchange" - - fid="sSXXGNUO2FefBTcCLIT/2Q==" - - product="NGINX" - - product="Roundcube-Webmail" && product.version="1.6.10" - - category="服务" - - type="service" / type="subdomain" - - cloud_name="Aliyundun" - - is_cloud=true / is_cloud=false - - is_fraud=true / is_fraud=false - - is_honeypot=true / is_honeypot=false -- 协议类（type=service）： - - protocol="quic" - - banner="users" - - banner_hash="7330105010150477363" - - banner_fid="zRpqmn0FXQRjZpH8MjMX55zpMy9SgsW8" - - base_protocol="udp" / base_protocol="tcp" -- 网站类（type=subdomain）： - - title="beijing" - - header="elastic" - - header_hash="1258854265" - - body="网络空间测绘" - - body_hash="-2090962452" - - js_name="js/jquery.js" - - js_md5="82ac3f14327a8b7ba49baa208d4eaa15" - - cname="customers.spektrix.com" - - cname_domain="siteforce.com" - - icon_hash="-247388890" - - status_code="402" - - icp="京ICP证030173号" - - sdk_hash="Are3qNnP2Eqn7q5kAoUO3l+w3mgVIytO" -- 地理位置（Location）： - - country="CN" 或 country="中国" - - region="Zhejiang" 或 region="浙江"（仅支持中国地区中文） - - city="Hangzhou" -- 证书类（Certificate）： - - cert="baidu" - - cert.subject="Oracle Corporation" - - cert.issuer="DigiCert" - - cert.subject.org="Oracle Corporation" - - cert.subject.cn="baidu.com" - - cert.issuer.org="cPanel, Inc." - - cert.issuer.cn="Synology Inc. CA" - - cert.domain="huawei.com" - - cert.is_equal=true / cert.is_equal=false - - cert.is_valid=true / cert.is_valid=false - - cert.is_match=true / cert.is_match=false - - cert.is_expired=true / cert.is_expired=false - - jarm="2ad2ad0002ad2ad22c2ad2ad2ad2ad2eac92ec34bcc0cf7520e97547f83e81" - - tls.version="TLS 1.3" - - tls.ja3s="15af977ce25de452b96affa2addb1036" - - cert.sn="356078156165546797850343536942784588840297" - - cert.not_after.after="2025-03-01" / cert.not_after.before="2025-03-01" - - cert.not_before.after="2025-03-01" / cert.not_before.before="2025-03-01" -- 时间类（Last update time）： - - after="2023-01-01" - - before="2023-12-01" - - after="2023-01-01" && before="2023-12-01" -- 独立IP语法（需配合 ip_filter / ip_exclude）： - - ip_filter(banner="SSH-2.0-OpenSSH_6.7p2") && ip_filter(icon_hash="-1057022626") - - ip_filter(banner="SSH-2.0-OpenSSH_6.7p2" && asn="3462") && ip_exclude(title="EdgeOS") - - port_size="6" / port_size_gt="6" / port_size_lt="12" - - ip_ports="80,161" - - ip_country="CN" - - ip_region="Zhejiang" - - ip_city="Hangzhou" - - ip_after="2021-03-18" - - ip_before="2019-09-09" - -生成约束与注意事项： -- 字符串值一律用英文双引号包裹，例如 title="登录"、country="CN" -- 字符串值保持字面一致：不要缩写（例如 city="beijing" 不要变成 city="BJ"），不要用别名（例如 Beijing/Peking），不要擅自翻译/音译/改写大小写 -- 地理位置字段（country/region/city）更倾向于“按用户给定值输出”；不确定合法取值时，不要猜测，把备选写进 warnings -- 不要捏造不存在的 FOFA 字段；不确定时把不确定点写进 warnings，并输出一个保守的 query -- 当用户描述里有“多个与/或条件”，优先加 () 明确优先级，例如：(app="Apache" || app="Nginx") && country="CN" -- 当用户缺少关键条件导致范围过大或歧义（如地点/协议/端口/服务类型未说明），允许 query 为空字符串，并在 warnings 里明确需要补充的信息 -`) +通用生成约束： +- 严格遵守“当前搜索引擎语法速查”里的字段名、操作符和示例风格；不同数据源语法不同，不要混用。 +- 字符串值保持用户原意：不要无依据缩写、翻译、音译、替换同义词或改写大小写。 +- 地理位置、组织名、产品名、域名、证书名、CVE 编号等实体值必须尽量保留原文；确需推断（如“中国”到 CN）时在 explanation 或 warnings 中说明。 +- 不要捏造字段。不确定字段是否支持时，选择更通用且确定的字段，或把不确定点写进 warnings。 +- 当用户描述里有多个与/或条件，必须使用该数据源支持的括号和逻辑操作符明确优先级。 +- 如果用户输入已经是当前数据源查询语法或非常接近，应原样返回；只在明显有语法错误且能确定修复方式时轻微修正，并在 explanation 说明。 +- 如果需求范围过大、关键目标缺失或语义矛盾，允许 query 为空字符串，并在 warnings 中明确需要补充的信息。 +- 只生成资产测绘/信息收集查询语法，不生成扫描、利用、爆破、绕过、命令执行或攻击步骤。 +`, engineName, engineName, engineName, engineName, engineName, syntaxNotes)) userPrompt := fmt.Sprintf("自然语言意图：%s", req.Text) @@ -294,19 +415,20 @@ func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) { } content := strings.TrimSpace(apiResponse.Choices[0].Message.Content) - // 兼容模型偶尔返回 ```json ... ``` 的情况 - content = strings.TrimPrefix(content, "```json") - content = strings.TrimPrefix(content, "```") - content = strings.TrimSuffix(content, "```") - content = strings.TrimSpace(content) + jsonContent, extractErr := extractInfoCollectJSONObject(content) + if extractErr != nil { + snippet := trimSnippet(content, 1200) + c.JSON(http.StatusBadGateway, gin.H{ + "error": "AI 返回内容无法解析为 JSON，请稍后重试或换个描述方式", + "snippet": snippet, + }) + return + } var parsed fofaParseResponse - if err := json.Unmarshal([]byte(content), &parsed); err != nil { + if err := json.Unmarshal([]byte(jsonContent), &parsed); err != nil { // 直接回传一部分原文，方便排查，但避免太大 - snippet := content - if len(snippet) > 1200 { - snippet = snippet[:1200] - } + snippet := trimSnippet(content, 1200) c.JSON(http.StatusBadGateway, gin.H{ "error": "AI 返回内容无法解析为 JSON，请稍后重试或换个描述方式", "snippet": snippet, @@ -317,13 +439,106 @@ func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) { if parsed.Query == "" { // query 允许为空（表示需求不明确），但前端需要明确提示 if len(parsed.Warnings) == 0 { - parsed.Warnings = []string{"需求信息不足，未能生成可用的 FOFA 查询语法，请补充关键条件（如国家/端口/产品/域名等）。"} + parsed.Warnings = []string{"需求信息不足，未能生成可用的 " + engineName + " 查询语法，请补充关键条件（如国家/端口/产品/域名等）。"} } } c.JSON(http.StatusOK, parsed) } +func extractInfoCollectJSONObject(content string) (string, error) { + content = strings.TrimSpace(content) + if content == "" { + return "", errors.New("empty content") + } + candidates := []string{content} + if fenced := extractFencedJSON(content); fenced != "" { + candidates = append([]string{fenced}, candidates...) + } + for _, candidate := range candidates { + candidate = strings.TrimSpace(candidate) + if candidate == "" { + continue + } + if json.Valid([]byte(candidate)) { + return candidate, nil + } + if obj := scanBalancedJSONObject(candidate); obj != "" && json.Valid([]byte(obj)) { + return obj, nil + } + } + return "", errors.New("json object not found") +} + +func extractFencedJSON(content string) string { + start := strings.Index(content, "```") + if start < 0 { + return "" + } + rest := content[start+3:] + if nl := strings.Index(rest, "\n"); nl >= 0 { + lang := strings.ToLower(strings.TrimSpace(rest[:nl])) + if lang == "" || lang == "json" || strings.HasPrefix(lang, "json ") { + rest = rest[nl+1:] + } + } + end := strings.Index(rest, "```") + if end < 0 { + return "" + } + return strings.TrimSpace(rest[:end]) +} + +func scanBalancedJSONObject(content string) string { + start := strings.Index(content, "{") + if start < 0 { + return "" + } + depth := 0 + inString := false + escaped := false + for i := start; i < len(content); i++ { + ch := content[i] + if inString { + if escaped { + escaped = false + continue + } + switch ch { + case '\\': + escaped = true + case '"': + inString = false + } + continue + } + switch ch { + case '"': + inString = true + case '{': + depth++ + case '}': + depth-- + if depth == 0 { + return strings.TrimSpace(content[start : i+1]) + } + } + } + return "" +} + +func trimSnippet(s string, maxRunes int) string { + s = strings.TrimSpace(s) + if maxRunes <= 0 { + return "" + } + runes := []rune(s) + if len(runes) <= maxRunes { + return s + } + return string(runes[:maxRunes]) +} + // Search FOFA 查询（后端代理，避免前端暴露 key） func (h *FofaHandler) Search(c *gin.Context) { var req fofaSearchRequest @@ -331,6 +546,11 @@ func (h *FofaHandler) Search(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) return } + provider := normalizeSpaceSearchProvider(req.Provider) + if provider == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "provider 不支持，可选：fofa、zoomeye、quake、shodan"}) + return + } req.Query = strings.TrimSpace(req.Query) if req.Query == "" { @@ -348,10 +568,15 @@ func (h *FofaHandler) Search(c *gin.Context) { req.Size = 10000 } if req.Fields == "" { - req.Fields = "host,ip,port,domain,title,protocol,country,province,city,server" + req.Fields = defaultFieldsForProvider(provider) } - apiKey := h.resolveAPIKey() + if provider != "fofa" { + h.searchExternalProvider(c, provider, req) + return + } + + apiKey := h.resolveAPIKey(provider) if apiKey == "" { c.JSON(http.StatusBadRequest, gin.H{ "error": "FOFA 未配置：请在系统设置的资产管理中填写 FOFA API Key，或设置环境变量 FOFA_API_KEY", @@ -361,7 +586,7 @@ func (h *FofaHandler) Search(c *gin.Context) { return } - baseURL := h.resolveBaseURL() + baseURL := h.resolveBaseURL(provider) qb64 := base64.StdEncoding.EncodeToString([]byte(req.Query)) u, err := url.Parse(baseURL) @@ -439,6 +664,7 @@ func (h *FofaHandler) Search(c *gin.Context) { } c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: provider, Query: req.Query, Size: apiResp.Size, Page: apiResp.Page, @@ -449,6 +675,272 @@ func (h *FofaHandler) Search(c *gin.Context) { }) } +func (h *FofaHandler) searchExternalProvider(c *gin.Context, provider string, req fofaSearchRequest) { + apiKey := h.resolveAPIKey(provider) + if apiKey == "" { + envKey := map[string]string{ + "zoomeye": "ZOOMEYE_API_KEY", + "quake": "QUAKE_API_KEY", + "shodan": "SHODAN_API_KEY", + }[provider] + c.JSON(http.StatusBadRequest, gin.H{ + "error": providerDisplayName(provider) + " 未配置：请在 config.yaml 中填写 api_key，或设置环境变量 " + envKey, + "need": []string{provider + ".api_key"}, + "env_key": []string{envKey}, + }) + return + } + + switch provider { + case "zoomeye": + h.searchZoomEye(c, req, apiKey) + case "quake": + h.searchQuake(c, req, apiKey) + case "shodan": + h.searchShodan(c, req, apiKey) + } +} + +func (h *FofaHandler) searchZoomEye(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := h.resolveBaseURL("zoomeye") + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "ZoomEye base_url 无效: " + err.Error()}) + return + } + body := map[string]interface{}{ + "qbase64": base64.StdEncoding.EncodeToString([]byte(req.Query)), + "page": req.Page, + "pagesize": req.Size, + } + if fields := strings.TrimSpace(req.Fields); fields != "" { + body["fields"] = fields + } + var apiResp struct { + Code int `json:"code"` + Message string `json:"message"` + Query string `json:"query"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"pagesize"` + Data []map[string]interface{} `json:"data"` + } + if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "API-KEY", body, &apiResp, "ZoomEye") { + return + } + if apiResp.Code != 60000 { + msg := strings.TrimSpace(apiResp.Message) + if msg == "" { + msg = "ZoomEye 返回错误" + } + c.JSON(http.StatusBadGateway, gin.H{"error": msg}) + return + } + fields := splitAndCleanCSV(req.Fields) + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "zoomeye", + Query: firstNonEmptySpaceSearchValue(apiResp.Query, req.Query), + Size: firstPositive(apiResp.PageSize, req.Size), + Page: firstPositive(apiResp.Page, req.Page), + Total: apiResp.Total, + Fields: fields, + ResultsCount: len(apiResp.Data), + Results: projectRows(apiResp.Data, fields), + }) +} + +func (h *FofaHandler) searchQuake(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := h.resolveBaseURL("quake") + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Quake base_url 无效: " + err.Error()}) + return + } + fields := splitAndCleanCSV(req.Fields) + body := map[string]interface{}{ + "query": req.Query, + "size": req.Size, + "start": (req.Page - 1) * req.Size, + "latest": req.Full, + } + if len(fields) > 0 { + body["include"] = fields + } + var apiResp struct { + Code int `json:"code"` + Message string `json:"message"` + TotalCount int `json:"total_count"` + Data []map[string]interface{} `json:"data"` + Meta struct { + Pagination struct { + Total int `json:"total"` + } `json:"pagination"` + } `json:"meta"` + } + if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "X-QuakeToken", body, &apiResp, "Quake") { + return + } + if apiResp.Code != 0 { + msg := strings.TrimSpace(apiResp.Message) + if msg == "" { + msg = "Quake 返回错误" + } + c.JSON(http.StatusBadGateway, gin.H{"error": msg}) + return + } + total := firstPositive(apiResp.TotalCount, apiResp.Meta.Pagination.Total) + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "quake", + Query: req.Query, + Size: req.Size, + Page: req.Page, + Total: total, + Fields: fields, + ResultsCount: len(apiResp.Data), + Results: projectRows(apiResp.Data, fields), + }) +} + +func (h *FofaHandler) searchShodan(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := strings.TrimRight(h.resolveBaseURL("shodan"), "/") + "/shodan/host/search" + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Shodan base_url 无效: " + err.Error()}) + return + } + + var apiResp struct { + Total int `json:"total"` + Matches []map[string]interface{} `json:"matches"` + Error string `json:"error"` + } + targetSize := req.Size + if targetSize <= 0 { + targetSize = 100 + } + if targetSize > 1000 { + targetSize = 1000 + } + page := req.Page + matches := make([]map[string]interface{}, 0, targetSize) + pagesNeeded := (targetSize + 99) / 100 + for i := 0; i < pagesNeeded; i++ { + pageURL := *u + params := pageURL.Query() + params.Set("key", apiKey) + params.Set("query", req.Query) + params.Set("page", fmt.Sprintf("%d", page+i)) + params.Set("minify", "false") + if fields := strings.TrimSpace(req.Fields); fields != "" { + params.Set("fields", fields) + } + pageURL.RawQuery = params.Encode() + apiResp.Matches = nil + apiResp.Error = "" + if !h.doJSONRequest(c, http.MethodGet, pageURL.String(), "", "", nil, &apiResp, "Shodan") { + return + } + if strings.TrimSpace(apiResp.Error) != "" { + c.JSON(http.StatusBadGateway, gin.H{"error": apiResp.Error}) + return + } + if len(apiResp.Matches) == 0 { + break + } + matches = append(matches, apiResp.Matches...) + if len(matches) >= targetSize { + matches = matches[:targetSize] + break + } + } + fields := splitAndCleanCSV(req.Fields) + expectedCount := shodanExpectedResultCount(apiResp.Total, page, targetSize) + shortfall := expectedCount - len(matches) + warning := "" + if shortfall > 0 { + warning = fmt.Sprintf("Shodan 统计总数为 %d，但本次分页实际只返回 %d/%d 条明细", apiResp.Total, len(matches), expectedCount) + } + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "shodan", + Query: req.Query, + Size: targetSize, + Page: page, + Total: apiResp.Total, + Fields: fields, + ResultsCount: len(matches), + ExpectedCount: expectedCount, + Shortfall: max(0, shortfall), + Warning: warning, + Results: projectRows(matches, fields), + }) +} + +func shodanExpectedResultCount(total, page, size int) int { + if total <= 0 || size <= 0 { + return 0 + } + if page <= 0 { + page = 1 + } + startOffset := (page - 1) * 100 + remaining := total - startOffset + if remaining <= 0 { + return 0 + } + if remaining < size { + return remaining + } + return size +} + +func (h *FofaHandler) doJSONRequest(c *gin.Context, method, endpoint, apiKey, headerName string, body interface{}, out interface{}, label string) bool { + var reqBody *strings.Reader + if body != nil { + b, err := json.Marshal(body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) + return false + } + reqBody = strings.NewReader(string(b)) + } else { + reqBody = strings.NewReader("") + } + httpReq, err := http.NewRequestWithContext(c.Request.Context(), method, endpoint, reqBody) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) + return false + } + httpReq.Header.Set("User-Agent", "CyberStrikeAI/1.7.4") + httpReq.Header.Set("Accept", "application/json") + if body != nil { + httpReq.Header.Set("Content-Type", "application/json") + } + if headerName != "" && apiKey != "" { + httpReq.Header.Set(headerName, apiKey) + } + resp, err := h.client.Do(httpReq) + if err != nil { + status, message, timeout := safeFofaRequestError(err) + h.logger.Warn("请求空间测绘搜索失败", + zap.String("provider", label), + zap.Bool("timeout", timeout), + zap.String("error_type", fmt.Sprintf("%T", err)), + ) + c.JSON(status, gin.H{"error": strings.Replace(message, "FOFA", label, 1)}) + return false + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("%s 返回非 2xx: %d", label, resp.StatusCode)}) + return false + } + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "解析 " + label + " 响应失败: " + err.Error()}) + return false + } + return true +} + func safeFofaRequestError(err error) (status int, message string, timeout bool) { var netErr net.Error timeout = errors.Is(err, context.DeadlineExceeded) || @@ -480,3 +972,58 @@ func splitAndCleanCSV(s string) []string { } return out } + +func projectRows(rows []map[string]interface{}, fields []string) []map[string]interface{} { + if len(fields) == 0 { + return rows + } + out := make([]map[string]interface{}, 0, len(rows)) + for _, row := range rows { + item := make(map[string]interface{}, len(fields)) + for _, field := range fields { + item[field] = valueByPath(row, field) + } + out = append(out, item) + } + return out +} + +func valueByPath(row map[string]interface{}, path string) interface{} { + if row == nil { + return nil + } + if v, ok := row[path]; ok { + return v + } + parts := strings.Split(path, ".") + var current interface{} = row + for _, part := range parts { + m, ok := current.(map[string]interface{}) + if !ok { + return nil + } + current, ok = m[part] + if !ok { + return nil + } + } + return current +} + +func firstPositive(values ...int) int { + for _, v := range values { + if v > 0 { + return v + } + } + return 0 +} + +func firstNonEmptySpaceSearchValue(values ...string) string { + for _, v := range values { + if strings.TrimSpace(v) != "" { + return v + } + } + return "" +} diff --git a/internal/handler/fofa_test.go b/internal/handler/fofa_test.go index a35fd0d8..200ed461 100644 --- a/internal/handler/fofa_test.go +++ b/internal/handler/fofa_test.go @@ -3,6 +3,7 @@ package handler import ( "context" "encoding/json" + "fmt" "net/http" "net/http/httptest" "net/url" @@ -84,3 +85,115 @@ func TestSafeFofaRequestErrorDoesNotExposeURLOrAPIKey(t *testing.T) { t.Fatalf("safe error exposed request URL or API key: %q", message) } } + +func TestShodanSearchReportsShortfallWhenTotalExceedsMatches(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("SHODAN_API_KEY", "") + + shodanServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/shodan/host/search" { + t.Fatalf("unexpected path: %s", r.URL.Path) + } + if got := r.URL.Query().Get("key"); got != "test-shodan-key" { + t.Fatalf("Shodan key = %q, want test-shodan-key", got) + } + page := r.URL.Query().Get("page") + count := 0 + switch page { + case "1": + count = 100 + case "2": + count = 3 + default: + count = 0 + } + matches := make([]map[string]interface{}, 0, count) + for i := 0; i < count; i++ { + matches = append(matches, map[string]interface{}{ + "ip_str": fmt.Sprintf("192.0.2.%d", i+1), + "port": 80, + }) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "total": 104, + "matches": matches, + }) + })) + defer shodanServer.Close() + + h := NewFofaHandler(&config.Config{ + Shodan: config.SpaceSearchConfig{ + BaseURL: shodanServer.URL, + APIKey: "test-shodan-key", + }, + }, zap.NewNop()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + body := `{"provider":"shodan","query":"product:nginx","fields":"ip_str,port","size":1000,"page":1}` + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + h.Search(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var response fofaSearchResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.Total != 104 || response.ResultsCount != 103 { + t.Fatalf("counts: total=%d results_count=%d, want 104/103", response.Total, response.ResultsCount) + } + if response.ExpectedCount != 104 || response.Shortfall != 1 { + t.Fatalf("shortfall: expected=%d shortfall=%d, want 104/1", response.ExpectedCount, response.Shortfall) + } + if response.Warning == "" { + t.Fatal("warning should explain shortfall") + } +} + +func TestExtractInfoCollectJSONObject(t *testing.T) { + t.Parallel() + cases := []struct { + name string + in string + want string + }{ + { + name: "plain json", + in: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`, + want: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`, + }, + { + name: "fenced json", + in: "```json\n{\"query\":\"product:nginx\"}\n```", + want: `{"query":"product:nginx"}`, + }, + { + name: "prefixed explanation", + in: "解析结果如下：\n{\"query\":\"ssl.cert.subject.cn:example.com\",\"explanation\":\"ok\"}\n请确认。", + want: `{"query":"ssl.cert.subject.cn:example.com","explanation":"ok"}`, + }, + { + name: "braces inside string", + in: "结果：{\"query\":\"title:\\\"{admin}\\\"\",\"warnings\":[\"check\"]}", + want: `{"query":"title:\"{admin}\"","warnings":["check"]}`, + }, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := extractInfoCollectJSONObject(tc.in) + if err != nil { + t.Fatalf("extractInfoCollectJSONObject() error = %v", err) + } + if got != tc.want { + t.Fatalf("extractInfoCollectJSONObject() = %q, want %q", got, tc.want) + } + }) + } +}