#!/usr/bin/env python3
"""Validate and enrich extracted IOCs.

Filters out private/reserved IPs, validates domain formats, checks hash lengths,
and optionally enriches IOCs via VirusTotal, AbuseIPDB, or OTX APIs.
"""

from __future__ import annotations

import argparse
import ipaddress
import json
import os
import re
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Optional
from urllib.error import URLError
from urllib.request import Request, urlopen


# --- Validation Functions ---

RESERVED_DOMAINS = {
    "localhost", "example.com", "example.org", "example.net",
    "test.com", "invalid", "localhost.localdomain",
}

COMMON_FP_DOMAINS = {
    "microsoft.com", "google.com", "windows.com", "w3.org",
    "schema.org", "xml.org", "mozilla.org", "apache.org",
    "github.com", "googleapis.com", "gstatic.com",
}


def is_private_ip(ip_str: str) -> bool:
    """Check if an IP address is private, reserved, or loopback."""
    try:
        addr = ipaddress.ip_address(ip_str)
        return (
            addr.is_private
            or addr.is_reserved
            or addr.is_loopback
            or addr.is_multicast
            or addr.is_link_local
            or addr.is_unspecified
        )
    except ValueError:
        return True  # Invalid IP treated as private


def is_valid_domain(domain: str) -> bool:
    """Validate a domain name format."""
    if not domain or len(domain) > 253:
        return False
    if domain.lower() in RESERVED_DOMAINS:
        return False
    # Must have at least one dot
    if "." not in domain:
        return False
    # Each label: 1-63 chars, alphanumeric or hyphen, no leading/trailing hyphen
    labels = domain.split(".")
    for label in labels:
        if not label or len(label) > 63:
            return False
        if not re.match(r"^[a-zA-Z0-9]([a-zA-Z0-9\-]*[a-zA-Z0-9])?$", label):
            return False
    # TLD must be alphabetic
    if not labels[-1].isalpha():
        return False
    return True


def is_valid_hash(value: str, hash_type: str) -> bool:
    """Validate a hash value."""
    expected_lengths = {"md5": 32, "sha1": 40, "sha256": 64}
    expected = expected_lengths.get(hash_type)
    if not expected:
        return False
    if len(value) != expected:
        return False
    try:
        int(value, 16)
        return True
    except ValueError:
        return False


def is_valid_url(url: str) -> bool:
    """Basic URL validation."""
    if not url.startswith(("http://", "https://")):
        return False
    if len(url) < 10 or len(url) > 2048:
        return False
    # Must have a domain component
    domain_match = re.search(r"https?://([^/:\s]+)", url)
    if not domain_match:
        return False
    return True


def validate_ioc(value: str, ioc_type: str, filter_private: bool = True,
                  filter_common: bool = False) -> dict:
    """Validate a single IOC and return validation result."""
    result = {
        "value": value,
        "type": ioc_type,
        "valid": False,
        "reason": None,
    }

    if ioc_type in ("ipv4", "ipv6"):
        try:
            ipaddress.ip_address(value)
            if filter_private and is_private_ip(value):
                result["reason"] = "private_or_reserved"
            else:
                result["valid"] = True
        except ValueError:
            result["reason"] = "invalid_format"

    elif ioc_type == "domain":
        if not is_valid_domain(value):
            result["reason"] = "invalid_format"
        elif filter_common and value.lower() in COMMON_FP_DOMAINS:
            result["reason"] = "common_false_positive"
        else:
            result["valid"] = True

    elif ioc_type == "url":
        if is_valid_url(value):
            result["valid"] = True
        else:
            result["reason"] = "invalid_format"

    elif ioc_type in ("md5", "sha1", "sha256"):
        if is_valid_hash(value, ioc_type):
            result["valid"] = True
        else:
            result["reason"] = "invalid_hash"

    elif ioc_type == "email":
        if re.match(r"^[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}$", value):
            result["valid"] = True
        else:
            result["reason"] = "invalid_format"

    else:
        # Accept other types as-is (registry keys, mutexes, file paths)
        result["valid"] = True

    return result


# --- Enrichment Functions ---

def enrich_virustotal(value: str, ioc_type: str, api_key: str) -> Optional[dict]:
    """Query VirusTotal for IOC enrichment."""
    base_url = "https://www.virustotal.com/api/v3"

    type_endpoints = {
        "ipv4": f"/ip_addresses/{value}",
        "domain": f"/domains/{value}",
        "url": f"/urls",  # Requires URL ID
        "md5": f"/files/{value}",
        "sha1": f"/files/{value}",
        "sha256": f"/files/{value}",
    }

    endpoint = type_endpoints.get(ioc_type)
    if not endpoint:
        return None

    try:
        url = base_url + endpoint
        req = Request(url)
        req.add_header("x-apikey", api_key)
        req.add_header("Accept", "application/json")
        response = urlopen(req, timeout=15)
        data = json.loads(response.read())

        enrichment = {"source": "virustotal"}
        attrs = data.get("data", {}).get("attributes", {})

        if ioc_type in ("md5", "sha1", "sha256"):
            stats = attrs.get("last_analysis_stats", {})
            enrichment["malicious"] = stats.get("malicious", 0)
            enrichment["undetected"] = stats.get("undetected", 0)
            enrichment["detection_ratio"] = (
                f"{stats.get('malicious', 0)}/{sum(stats.values())}"
                if stats else "unknown"
            )
        elif ioc_type in ("ipv4", "domain"):
            stats = attrs.get("last_analysis_stats", {})
            enrichment["malicious"] = stats.get("malicious", 0)
            enrichment["reputation"] = attrs.get("reputation", "unknown")

        return enrichment
    except (URLError, json.JSONDecodeError, KeyError) as e:
        return {"source": "virustotal", "error": str(e)}


def enrich_abuseipdb(ip: str, api_key: str) -> Optional[dict]:
    """Query AbuseIPDB for IP reputation."""
    try:
        url = f"https://api.abuseipdb.com/api/v2/check?ipAddress={ip}&maxAgeInDays=90"
        req = Request(url)
        req.add_header("Key", api_key)
        req.add_header("Accept", "application/json")
        response = urlopen(req, timeout=15)
        data = json.loads(response.read())

        report = data.get("data", {})
        return {
            "source": "abuseipdb",
            "abuse_confidence": report.get("abuseConfidenceScore", 0),
            "total_reports": report.get("totalReports", 0),
            "country": report.get("countryCode", "unknown"),
            "isp": report.get("isp", "unknown"),
        }
    except (URLError, json.JSONDecodeError) as e:
        return {"source": "abuseipdb", "error": str(e)}


def validate_ioc_file(input_data: dict, filter_private: bool = True,
                       filter_common: bool = False,
                       vt_api_key: str = None,
                       abuse_api_key: str = None,
                       rate_limit: float = 0.25) -> dict:
    """Validate all IOCs in the input data structure."""
    indicators = input_data.get("indicators", input_data)
    validated = {
        "metadata": {
            "timestamp": datetime.utcnow().isoformat(),
            "source": input_data.get("metadata", {}).get("source", ""),
            "validation_settings": {
                "filter_private_ips": filter_private,
                "filter_common_domains": filter_common,
                "enrichment_enabled": bool(vt_api_key or abuse_api_key),
            },
        },
        "valid_indicators": {},
        "filtered_indicators": {},
        "statistics": {
            "total_input": 0,
            "total_valid": 0,
            "total_filtered": 0,
        },
    }

    for ioc_type, values in indicators.items():
        if isinstance(values, list):
            ioc_values = []
            for item in values:
                if isinstance(item, dict):
                    ioc_values.append(item.get("value", str(item)))
                else:
                    ioc_values.append(str(item))
        else:
            continue

        valid_list = []
        filtered_list = []

        for value in ioc_values:
            validated["statistics"]["total_input"] += 1
            result = validate_ioc(value, ioc_type, filter_private, filter_common)

            if result["valid"]:
                entry = {"value": value, "type": ioc_type}

                # Enrich if API keys provided
                if vt_api_key and ioc_type in ("ipv4", "domain", "md5", "sha1", "sha256"):
                    enrichment = enrich_virustotal(value, ioc_type, vt_api_key)
                    if enrichment:
                        entry["virustotal"] = enrichment
                    time.sleep(rate_limit)

                if abuse_api_key and ioc_type in ("ipv4",):
                    enrichment = enrich_abuseipdb(value, abuse_api_key)
                    if enrichment:
                        entry["abuseipdb"] = enrichment
                    time.sleep(rate_limit)

                valid_list.append(entry)
                validated["statistics"]["total_valid"] += 1
            else:
                filtered_list.append({
                    "value": value,
                    "reason": result["reason"],
                })
                validated["statistics"]["total_filtered"] += 1

        if valid_list:
            validated["valid_indicators"][ioc_type] = valid_list
        if filtered_list:
            validated["filtered_indicators"][ioc_type] = filtered_list

    return validated


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Validate and optionally enrich extracted IOCs"
    )
    parser.add_argument(
        "--input", "-i",
        required=True,
        help="Input JSON file with extracted IOCs",
    )
    parser.add_argument(
        "--output", "-o",
        default=None,
        help="Output JSON file path (default: stdout)",
    )
    parser.add_argument(
        "--no-filter-private",
        action="store_true",
        help="Do not filter private/reserved IP addresses",
    )
    parser.add_argument(
        "--filter-common",
        action="store_true",
        help="Filter commonly seen domains (microsoft.com, google.com, etc.)",
    )
    parser.add_argument(
        "--enrich",
        action="store_true",
        help="Enable API enrichment (requires API keys)",
    )
    parser.add_argument(
        "--vt-api-key",
        default=os.environ.get("VT_API_KEY"),
        help="VirusTotal API key (or set VT_API_KEY env var)",
    )
    parser.add_argument(
        "--abuse-api-key",
        default=os.environ.get("ABUSEIPDB_API_KEY"),
        help="AbuseIPDB API key (or set ABUSEIPDB_API_KEY env var)",
    )
    parser.add_argument(
        "--rate-limit",
        type=float,
        default=0.25,
        help="Seconds between API requests (default: 0.25)",
    )
    parser.add_argument(
        "--format", default="json", choices=["json", "text", "csv"],
        help="Output format (default: json)",
    )

    args = parser.parse_args()

    if not os.path.isfile(args.input):
        print(f"[!] Error: File not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    try:
        with open(args.input, "r") as f:
            input_data = json.load(f)
    except json.JSONDecodeError as e:
        print(f"[!] Error: Invalid JSON: {e}", file=sys.stderr)
        sys.exit(1)

    vt_key = args.vt_api_key if args.enrich else None
    abuse_key = args.abuse_api_key if args.enrich else None

    if args.enrich and not (vt_key or abuse_key):
        print("[!] Warning: --enrich specified but no API keys provided", file=sys.stderr)

    results = validate_ioc_file(
        input_data,
        filter_private=not args.no_filter_private,
        filter_common=args.filter_common,
        vt_api_key=vt_key,
        abuse_api_key=abuse_key,
        rate_limit=args.rate_limit,
    )

    output_json = json.dumps(results, indent=2)

    if args.output:
        with open(args.output, "w") as f:
            f.write(output_json)
        print(f"[+] Validated IOCs saved to: {args.output}", file=sys.stderr)
    else:
        print(output_json)

    stats = results["statistics"]
    print(f"\n[+] Validation complete:", file=sys.stderr)
    print(f"    Input:    {stats['total_input']}", file=sys.stderr)
    print(f"    Valid:    {stats['total_valid']}", file=sys.stderr)
    print(f"    Filtered: {stats['total_filtered']}", file=sys.stderr)


if __name__ == "__main__":
    main()