name: "virustotal_search"
command: "python3"
args:
  - "-c"
  - |
    import sys
    import json
    import requests
    import os
    import time

    # ==================== VirusTotal 配置 ====================
    # 请在此处配置您的 VirusTotal API 密钥
    # 您也可以在环境变量中设置：VT_API_KEY
    # enable 默认为 false，需开启才能调用该MCP
    VT_API_KEY = ""  # 请填写您的 VirusTotal API 密钥
    # =======================================================

    # VirusTotal API 基础 URL
    BASE_URL = "https://www.virustotal.com/api/v3"

    def parse_args():
        """解析命令行参数"""
        # 尝试从第一个参数读取 JSON 配置
        if len(sys.argv) > 1:
            try:
                arg1 = str(sys.argv[1])
                config = json.loads(arg1)
                if isinstance(config, dict):
                    return config
            except (json.JSONDecodeError, TypeError, ValueError):
                pass

        # 传统位置参数方式
        config = {}
        if len(sys.argv) > 1:
            config['domain'] = str(sys.argv[1])
        if len(sys.argv) > 2:
            try:
                config['limit'] = int(sys.argv[2])
            except (ValueError, TypeError):
                pass
        if len(sys.argv) > 3:
            config['include_ips'] = sys.argv[3].lower() in ('true', '1', 'yes')
        return config

    def query_virustotal_subdomains(domain, api_key, limit=100, include_ips=False):
        """
        查询 VirusTotal 的子域名信息
        
        Args:
            domain: 要查询的域名
            api_key: VirusTotal API 密钥
            limit: 返回结果数量限制
            include_ips: 是否包含 IP 地址信息
        
        Returns:
            dict: 包含查询结果的字典
        """
        # 构建 API 请求 URL
        url = f"{BASE_URL}/domains/{domain}/subdomains"
        
        headers = {
            "x-apikey": api_key,
            "accept": "application/json"
        }
        
        params = {
            "limit": min(limit, 40)  # API 限制最大 40
        }
        
        all_results = []
        next_url = None
        
        try:
            # 处理分页
            while True:
                if next_url:
                    response = requests.get(next_url, headers=headers, timeout=30)
                else:
                    response = requests.get(url, headers=headers, params=params, timeout=30)
                
                response.raise_for_status()
                data = response.json()
                
                # 提取子域名数据
                if 'data' in data and data['data']:
                    for item in data['data']:
                        if 'id' in item:
                            subdomain_info = {
                                'subdomain': item['id'],
                                'type': item.get('type', 'domain'),
                            }
                            
                            # 如果 include_ips 为 True，尝试获取解析 IP
                            if include_ips and 'attributes' in item:
                                attributes = item.get('attributes', {})
                                # 这里简化处理，实际可能需要额外的 API 调用
                                subdomain_info['last_dns_records'] = attributes.get('last_dns_records', [])
                            
                            all_results.append(subdomain_info)
                    
                    # 检查是否有下一页
                    if 'links' in data and 'next' in data['links'] and len(all_results) < limit:
                        next_url = data['links']['next']
                        # 避免请求过快
                        time.sleep(0.5)
                    else:
                        break
                else:
                    break
                
                # 如果已达到限制，停止获取
                if len(all_results) >= limit:
                    break
            
            # 处理返回结果
            if all_results:
                return {
                    "status": "success",
                    "domain": domain,
                    "total_found": len(all_results),
                    "results": all_results[:limit],
                    "message": f"成功获取 {len(all_results[:limit])} 个子域名"
                }
            else:
                return {
                    "status": "success",
                    "domain": domain,
                    "total_found": 0,
                    "results": [],
                    "message": f"未找到 {domain} 的子域名"
                }
                
        except requests.exceptions.RequestException as e:
            error_msg = str(e)
            error_result = {
                "status": "error",
                "message": f"API 请求失败: {error_msg}",
                "suggestion": "请检查网络连接、API 密钥是否正确，或 VirusTotal API 服务是否可用"
            }
            
            # 处理特定 HTTP 状态码
            if hasattr(e, 'response') and e.response:
                status_code = e.response.status_code
                if status_code == 401:
                    error_result["message"] = "API 密钥无效或未授权"
                    error_result["suggestion"] = "请检查 VirusTotal API 密钥是否正确，或在 https://www.virustotal.com/ 获取有效密钥"
                elif status_code == 429:
                    error_result["message"] = "API 请求频率超限"
                    error_result["suggestion"] = "请稍后再试，VirusTotal API 有严格的速率限制（免费版每分钟4次）"
                elif status_code == 404:
                    error_result["message"] = f"域名 '{domain}' 不存在或未找到"
            
            return error_result

    try:
        config = parse_args()
        
        if not isinstance(config, dict):
            error_result = {
                "status": "error",
                "message": f"参数解析错误: 期望字典类型，但得到 {type(config).__name__}",
                "type": "TypeError"
            }
            print(json.dumps(error_result, ensure_ascii=False, indent=2))
            sys.exit(1)
        
        # 获取 API 密钥（从配置或环境变量）
        api_key = os.getenv('VT_API_KEY', VT_API_KEY).strip()
        
        if not api_key:
            error_result = {
                "status": "error",
                "message": "缺少 VirusTotal API 密钥",
                "required_config": ["VT_API_KEY"],
                "note": "请在 YAML 文件的 VT_API_KEY 配置项中填写您的 VirusTotal API 密钥，或在环境变量 VT_API_KEY 中设置。API 密钥可在 https://www.virustotal.com/ 注册获取"
            }
            print(json.dumps(error_result, ensure_ascii=False, indent=2))
            sys.exit(1)
        
        # 获取必需参数
        domain = config.get('domain', '').strip()
        if not domain:
            error_result = {
                "status": "error",
                "message": "缺少必需参数: domain（要查询的域名）",
                "required_params": ["domain"],
                "examples": [
                    "example.com",
                    "google.com",
                    "baidu.com"
                ]
            }
            print(json.dumps(error_result, ensure_ascii=False, indent=2))
            sys.exit(1)
        
        # 获取可选参数
        limit = config.get('limit', 100)
        try:
            limit = int(limit)
            if limit < 1:
                limit = 100
            elif limit > 1000:
                limit = 1000  # 限制最大 1000
        except (ValueError, TypeError):
            limit = 100
        
        include_ips = config.get('include_ips', False)
        if isinstance(include_ips, str):
            include_ips = include_ips.lower() in ('true', '1', 'yes')
        
        # 执行查询
        result = query_virustotal_subdomains(domain, api_key, limit, include_ips)
        
        # 输出结果
        print(json.dumps(result, ensure_ascii=False, indent=2))
        
    except Exception as e:
        error_result = {
            "status": "error",
            "message": f"执行出错: {str(e)}",
            "type": type(e).__name__
        }
        print(json.dumps(error_result, ensure_ascii=False, indent=2))
        sys.exit(1)

enabled: false

short_description: "VirusTotal 子域名查询工具，通过 VirusTotal API 被动收集域名子域名"

description: |
  VirusTotal 子域名查询工具，利用 VirusTotal 聚合的历史 DNS 数据来发现目标域名的子域名。
  
  **主要功能：**
  - 被动子域名收集：从 VirusTotal 历史 DNS 数据中检索子域名
  - 分页查询：支持大量子域名的获取
  - IP 关联：可选包含 DNS 解析记录
  - 去重处理：自动去重返回结果
  
  **使用场景：**
  - 安全测试前期信息收集
  - 企业网络资产发现
  - 攻击面分析
  - 威胁情报收集
  - 渗透测试信息收集
  
  **数据来源：**
  VirusTotal 聚合了来自多个来源的 DNS 数据，包括：
  - 历史 DNS 解析记录
  - 被动 DNS 数据库
  - 证书透明度日志
  - 安全扫描数据
  
  **注意事项：**
  - **API 密钥必需**：需要在 VirusTotal 注册账号并获取 API 密钥
  - **速率限制**：免费版 API 每分钟限制 4 次请求
  - **数据时效性**：数据基于历史扫描记录，可能不是实时的
  - **使用授权**：仅允许对您拥有合法授权的目标进行查询
  - **配额限制**：免费版每月有查询配额限制

parameters:
  - name: "domain"
    type: "string"
    description: |
      要查询的目标域名（必需）。
      
      **格式要求：**
      - 仅输入主域名，不要包含协议头（http://）或路径
      - 支持二级域名查询
      
      **示例值：**
      - "example.com"
      - "google.com"
      - "baidu.com"
      - "github.com"
      
      **注意事项：**
      - 域名格式必须正确
      - 查询结果可能包含跨域子域名
    required: true
    position: 2
    format: "positional"
  
  - name: "limit"
    type: "int"
    description: |
      返回结果数量限制（可选）。
      
      **说明：**
      - 默认值：40
      - 最大值：1000（API 限制）
      - 建议值：100-500
      
      **注意事项：**
      - 设置过大的值可能导致请求超时
      - API 单次返回限制为 40 条，超过会自动分页
    required: false
    position: 3
    format: "positional"
    default: 40
  
  - name: "include_ips"
    type: "bool"
    description: |
      是否包含 IP 地址信息（可选）。
      
      **说明：**
      - true：在结果中包含 DNS 解析记录
      - false：仅返回子域名列表
      
      **注意事项：**
      - 包含 IP 信息会增加 API 调用次数
      - 可能包含历史解析 IP，不一定准确
    required: false
    position: 4
    format: "positional"
    default: false