#!/usr/bin/env python3
"""
CVE Template Field Enhancement Script
Adds missing impact and remediation fields from ProjectDiscovery API
"""

import re
import sys
import requests
from pathlib import Path

API_URL = "https://api.projectdiscovery.io/v2/vulnerability"

def check_missing_fields(content):
    """Check if impact or remediation fields are missing."""
    has_impact = bool(re.search(r'^\s{2}impact:', content, re.MULTILINE))
    has_remediation = bool(re.search(r'^\s{2}remediation:', content, re.MULTILINE))
    return (not has_impact, not has_remediation)

def extract_cve_id(content):
    """Extract CVE ID from template."""
    match = re.search(r'^id:\s*(CVE-\d{4}-\d+)', content, re.MULTILINE)
    return match.group(1) if match else None

def fetch_from_api(cve_id):
    """Fetch impact and remediation from ProjectDiscovery API."""
    try:
        response = requests.get(f"{API_URL}/{cve_id}", timeout=10)
        if response.status_code == 200:
            data = response.json().get('data', {})
            return data.get('impact'), data.get('remediation')
    except:
        pass
    return None, None

def add_fields(content, impact, remediation, needs_impact, needs_remediation):
    """Add impact and remediation fields after description."""
    lines = content.split('\n')
    result = []
    i = 0

    while i < len(lines):
        line = lines[i]
        result.append(line)

        # Found description field
        if re.match(r'^\s{2}description:', line):
            i += 1
            # Copy all description content (lines starting with 4+ spaces or |)
            while i < len(lines):
                if lines[i].strip() == '' or lines[i].startswith('    ') or lines[i].strip() == '|':
                    result.append(lines[i])
                    i += 1
                else:
                    # End of description block - add new fields here
                    if needs_impact and impact:
                        result.append('  impact: |')
                        result.append(f'    {impact}')
                    if needs_remediation and remediation:
                        result.append('  remediation: |')
                        result.append(f'    {remediation}')
                    break
            continue

        i += 1

    return '\n'.join(result)

def enhance_file(file_path):
    """Enhance a single CVE template file."""
    try:
        content = file_path.read_text()

        # Check what's missing
        needs_impact, needs_remediation = check_missing_fields(content)
        if not needs_impact and not needs_remediation:
            print(f"✓ {file_path.name} - already complete")
            return False

        # Get CVE ID
        cve_id = extract_cve_id(content)
        if not cve_id:
            print(f"✗ {file_path.name} - no CVE ID found")
            return False

        # Fetch from API
        impact, remediation = fetch_from_api(cve_id)
        if not impact and not remediation:
            print(f"⚠ {cve_id} - not found in API")
            return False

        # Add fields
        enhanced = add_fields(content, impact, remediation, needs_impact, needs_remediation)
        file_path.write_text(enhanced)

        print(f"✓ {cve_id} - enhanced")
        return True

    except Exception as e:
        print(f"✗ {file_path.name} - error: {e}")
        return False

def main():
    if len(sys.argv) < 2:
        print("Usage: python enhance-cve-fields.py <file1> [file2] ...")
        sys.exit(1)

    enhanced = 0
    for file_arg in sys.argv[1:]:
        file_path = Path(file_arg)
        if file_path.exists():
            if enhance_file(file_path):
                enhanced += 1

    print(f"\nEnhanced {enhanced} file(s)")
    sys.exit(0)

if __name__ == '__main__':
    main()
