import requests
import csv
import json
from datetime import datetime
import time

def query_dns_batch(apikey, resources, output_file, batch_size=100):
    """
    批量查询DNS数据并保存为CSV文件
    :param apikey: API密钥
    :param resources: IP地址或域名列表
    :param output_file: 输出的CSV文件路径
    :param batch_size: 每次查询的资源数量（最大支持100个）
    """
    url = "https://api.threatbook.cn/v3/scene/dns"
    headers = [
        "地址", "一级分类", "二级分类", "是否恶意", "可信度",
        "严重级别", "情报类型", "标签"
    ]

    results = {}
    
    # 按批次处理
    for i in range(0, len(resources), batch_size):
        batch_resources = resources[i:i + batch_size]
        params = {
            "apikey": apikey,
            "resource": ",".join(batch_resources),  # 批量查询多个资源
            "lang": "zh"  # 返回中文结果
        }

        try:
            response = requests.get(url, params=params)
            data = response.json()

            if data.get("response_code") == 0:
                results.update(data.get("data", {}))
            else:
                print(f"查询失败，错误消息: {data.get('verbose_msg')}")

            # 避免频率限制
            time.sleep(1)
        except Exception as e:
            print(f"请求失败，批次: {batch_resources}, 错误: {e}")
    
    # 写入CSV文件
    with open(output_file, 'w', newline='', encoding='utf-8-sig') as f:
        writer = csv.writer(f)
        writer.writerow(headers)
        indcator_types = ["ips","domains"]
        for indicator_type in indcator_types:
            for resource, record in results.get(indicator_type).items():
                # 提取网站分类
                first_categories = ""
                second_categories = ""
                if indicator_type == "domains":
                    first_categories = json.dumps(record.get("categories", {}).get("first_cats", {}), ensure_ascii=False)  # 分类
                    second_categories = record.get("categories", {}).get("second_cats", {})
                # 提取字段
                is_malicious = record.get("is_malicious", "")  # 是否恶意
                confidence_level = record.get("confidence_level", "")  # 可信度
                severity = record.get("severity", "")  # 严重程度
                judgments = ",".join(record.get("judgments", []))  # 威胁类型   

                
                # 提取标签信息
                tags_classes = record.get("tags_classes", [])
                tags_info = "; ".join(
                    [", ".join(tag.get("tags", [])) if isinstance(tag.get("tags", []), list) else str(tag.get("tags", ""))
                     for tag in tags_classes] 
                )
                
            

                # 写入行数据
                row = [
                    resource,
                    first_categories,
                    second_categories,
                    is_malicious,
                    confidence_level,
                    severity,
                    judgments,
                    tags_info
                ]
                writer.writerow(row)
                print(f"已处理资源: {resource}")

if __name__ == "__main__":
    # 示例输入
    input_resources = ["8.8.8.8", "fget-career.com"]  # 待查询的资源列表
    api_key = "******"  # 替换为实际的API Key
    output_csv = f"./compromise_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    
    # 调用查询函数
    query_dns_batch(api_key, input_resources, output_csv)
    print(f"查询完成，结果已保存到文件: {output_csv}")