"""
漏洞API测试演示
测试所有漏洞相关API方法
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from sarm_sdk.apis.vulnerability import Vulnerability
from data_generator import DataGenerator
import json


def test_vulnerability_api():
    """测试漏洞API的所有方法"""
    print("=== 开始测试漏洞API ===")
    
    # 初始化API客户端和数据生成器
    vuln_api = Vulnerability()
    generator = DataGenerator()
    
    try:
        # 1. 测试创建单个漏洞
        print("\n1. 测试创建单个漏洞...")
        vuln_data = generator.generate_vulnerability("high")
        create_result = vuln_api.create(vuln_data)
        print(f"创建结果: {json.dumps(create_result, indent=2, ensure_ascii=False)}")
        
        if create_result.get('code') == 200:
            created_vuln_id = create_result.get('data', {}).get('vuln_unique_id')
            print(f"成功创建漏洞，ID: {created_vuln_id}")
        
        # 2. 测试批量创建漏洞
        print("\n2. 测试批量创建漏洞...")
        batch_vulns = []
        severities = ["critical", "high", "medium", "low"]
        for i in range(4):
            vuln = generator.generate_vulnerability(severities[i % len(severities)])
            batch_vulns.append(vuln)
        
        batch_result = vuln_api.create_batch(batch_vulns)
        print(f"批量创建结果: {json.dumps(batch_result, indent=2, ensure_ascii=False)}")
        
        # 3. 测试查询漏洞列表
        print("\n3. 测试查询漏洞列表...")
        list_result = vuln_api.get_list()
        print(f"查询结果: {json.dumps(list_result, indent=2, ensure_ascii=False)}")
        
        # 4. 测试更新漏洞
        print("\n4. 测试更新漏洞...")
        if batch_vulns:
            update_vuln = batch_vulns[0].copy()
            update_vuln["title"] = "更新后的漏洞标题"
            update_vuln["description"] = "这是更新后的漏洞描述"
            update_vuln["status"] = "fixed"
            update_result = vuln_api.update(update_vuln)
            print(f"更新结果: {json.dumps(update_result, indent=2, ensure_ascii=False)}")
        
        # 5. 测试批量更新漏洞
        print("\n5. 测试批量更新漏洞...")
        if batch_vulns and len(batch_vulns) > 1:
            update_vulns = []
            for vuln in batch_vulns[1:]:
                updated_vuln = vuln.copy()
                updated_vuln["status"] = "under_review"
                update_vulns.append(updated_vuln)
            
            batch_update_result = vuln_api.update_batch(update_vulns)
            print(f"批量更新结果: {json.dumps(batch_update_result, indent=2, ensure_ascii=False)}")
        
        # 6. 测试批量删除漏洞
        print("\n6. 测试批量删除漏洞...")
        if batch_vulns and len(batch_vulns) > 2:
            delete_vuln_ids = [
                vuln["vuln_unique_id"] 
                for vuln in batch_vulns[2:]
            ]
            delete_result = vuln_api.delete_batch(delete_vuln_ids)
            print(f"删除结果: {json.dumps(delete_result, indent=2, ensure_ascii=False)}")
        
        print("\n=== 漏洞API测试完成 ===")
        
    except Exception as e:
        print(f"测试过程中出现错误: {str(e)}")
        print(f"错误类型: {type(e).__name__}")


if __name__ == "__main__":
    test_vulnerability_api()