import requests
import csv
import time
from datetime import datetime

def query_ip_analysis(apikey, ip_list, output_file):
    """
    批量查询IP分析数据并保存为CSV文件
    :param apikey: API密钥
    :param ip_list: IP地址列表
    :param output_file: 输出的CSV文件路径
    """
    url = "https://api.threatbook.cn/v3/ip/query"
    headers = [
        "IP地址", "威胁类型", "标签信息",
        "ASN编号", "ASN名称及注册信息", "ASN风险等级", "更新时间", "国家", "国家代码",
        "省", "城市", "经度", "纬度"
    ]

    results = {}
    
    # 循环处理每个IP
    for ip in ip_list:
        params = {
            "apikey": apikey,
            "resource": ip,
            "lang": "zh" #返回中文结果
        }
        try:
            print(f"正在查询IP地址: {ip}")
            response = requests.get(url, params=params)
            data = response.json()
            
            if data.get("response_code") == 0:
                results[ip] = data.get("data", {}).get(ip, {})
            else:
                print(f"查询失败，IP: {ip}, 错误消息: {data.get('verbose_msg')}")
            
            # 避免频率限制
            time.sleep(1)
        except Exception as e:
            print(f"请求失败，IP: {ip}, 错误: {e}")
    
    # 写入CSV文件
    with open(output_file, 'w', newline='', encoding='utf-8-sig') as f:
        writer = csv.writer(f)
        writer.writerow(headers)
        
        for ip, record in results.items():
            # 提取核心字段
            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] # 标签信息
            )

            # 提取ASN信息
            asn = record.get("asn", {})
            asn_number = asn.get("number", "")  # ASN编号
            asn_info = asn.get("info", "")  # ASN名称及注册信息
            asn_rank = asn.get("rank", "")  # ASN风险等级

            # 提取位置和时间信息
            update_time = record.get("update_time", "")  # 更新时间
            location = record.get("basic", {}).get("location", {})  # 地理位置
            country = location.get("country", "")  # 国家
            country_code = location.get("country_code", "")  # 国家代码
            province = location.get("province", "")  # 省
            city = location.get("city", "")  # 城市
            lng = location.get("lng", "")  # 经度
            lat = location.get("lat", "")  # 纬度

            # 写入行数据
            row = [
                ip,
                judgments,
                tags_info,
                asn_number,
                asn_info,
                asn_rank,
                update_time,
                country,
                country_code,
                province,
                city,
                lng,
                lat
            ]
            writer.writerow(row)
            print(f"已处理IP: {ip}")

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