CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-31 16:17:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-20 17:31:20 +08:00
committed by GitHub
parent 437bae0b15
commit 9bafc2ab98
3 changed files with 847 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/handler/config.go
+42 -11
View File
@@ -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")
internal/handler/fofa.go
+692 -145
View File
@@ -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 ""
}
internal/handler/fofa_test.go
+113
View File
@@ -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)
}
})
}
}