#!/usr/bin/env python3
"""Extract Indicators of Compromise (IOCs) from text files.

Supports extraction of IP addresses, domains, URLs, email addresses,
file hashes (MD5, SHA-1, SHA-256), mutexes, registry keys, and file paths.
Includes defanging, deduplication, and multiple output formats.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Optional


# --- IOC Regex Patterns ---

PATTERNS = {
    "ipv4": re.compile(
        r"\b(?:(?:25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)\.){3}"
        r"(?:25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)\b"
    ),
    "ipv6": re.compile(
        r"\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}\b"
    ),
    "domain": re.compile(
        r"\b(?:[a-zA-Z0-9](?:[a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?\.)+"
        r"(?:com|net|org|info|biz|xyz|top|online|site|club|ru|cn|tk|ml|ga|cf|gq|"
        r"io|co|me|tv|cc|pw|in|de|uk|fr|it|nl|br|au|ca|es|jp|kr|za|ua|su|onion)\b"
    ),
    "url": re.compile(
        r"https?://[^\s<>\"')\]}>]{3,200}"
    ),
    "email": re.compile(
        r"\b[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}\b"
    ),
    "md5": re.compile(
        r"\b[a-fA-F0-9]{32}\b"
    ),
    "sha1": re.compile(
        r"\b[a-fA-F0-9]{40}\b"
    ),
    "sha256": re.compile(
        r"\b[a-fA-F0-9]{64}\b"
    ),
    "registry_key": re.compile(
        r"\b(?:HKEY_LOCAL_MACHINE|HKEY_CURRENT_USER|HKEY_CLASSES_ROOT|"
        r"HKEY_USERS|HKEY_CURRENT_CONFIG|HKLM|HKCU|HKCR|HKU|HKCC)"
        r"\\[^\s\"'<>]{3,200}"
    ),
    "file_path_windows": re.compile(
        r"\b[A-Z]:\\(?:[^\s\\/:*?\"<>|]+\\)*[^\s\\/:*?\"<>|]+\b"
    ),
    "file_path_unix": re.compile(
        r"(?:^|\s)/(?:tmp|var|etc|usr|opt|home|root|bin|sbin)/[^\s\"'<>]{2,200}"
    ),
    "mutex": re.compile(
        r"\b(?:Global\\|Local\\)[^\s\"'<>]{3,100}"
    ),
    "user_agent": re.compile(
        r"(?:Mozilla/[45]\.0\s*\([^)]+\)[^\r\n]{10,200})"
    ),
    "bitcoin_address": re.compile(
        r"\b[13][a-km-zA-HJ-NP-Z1-9]{25,34}\b"
    ),
    "cve": re.compile(
        r"\bCVE-\d{4}-\d{4,7}\b"
    ),
}

# Known false positive patterns to filter
FALSE_POSITIVE_HASHES = {
    "d41d8cd98f00b204e9800998ecf8427e",  # MD5 of empty string
    "da39a3ee5e6b4b0d3255bfef95601890afd80709",  # SHA1 of empty string
    "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",  # SHA256 empty
}


def defang_ioc(ioc: str, ioc_type: str) -> str:
    """Defang an IOC for safe sharing."""
    if ioc_type in ("ipv4", "ipv6"):
        return ioc.replace(".", "[.]")
    elif ioc_type == "domain":
        return ioc.replace(".", "[.]")
    elif ioc_type == "url":
        result = ioc.replace("http://", "hxxp://")
        result = result.replace("https://", "hxxps://")
        result = result.replace(".", "[.]")
        return result
    elif ioc_type == "email":
        return ioc.replace("@", "[@]").replace(".", "[.]")
    return ioc


def refang_ioc(ioc: str) -> str:
    """Refang a defanged IOC back to its original form."""
    result = ioc.replace("[.]", ".")
    result = result.replace("[@]", "@")
    result = result.replace("hxxp://", "http://")
    result = result.replace("hxxps://", "https://")
    return result


def is_false_positive_hash(value: str) -> bool:
    """Check if a hash is a known false positive."""
    return value.lower() in FALSE_POSITIVE_HASHES


def classify_hash(value: str) -> Optional[str]:
    """Classify a hex string as MD5, SHA-1, or SHA-256."""
    length = len(value)
    if length == 32:
        return "md5"
    elif length == 40:
        return "sha1"
    elif length == 64:
        return "sha256"
    return None


def extract_iocs(text: str, ioc_types: Optional[list] = None) -> dict:
    """Extract IOCs from text, returning deduplicated results by type."""
    results = defaultdict(set)
    types_to_extract = ioc_types or list(PATTERNS.keys())

    # First refang any already-defanged IOCs in the text
    text_clean = refang_ioc(text)

    for ioc_type in types_to_extract:
        if ioc_type not in PATTERNS:
            continue
        pattern = PATTERNS[ioc_type]
        for match in pattern.finditer(text_clean):
            value = match.group(0).strip()

            # Filter false positive hashes
            if ioc_type in ("md5", "sha1", "sha256"):
                if is_false_positive_hash(value):
                    continue
                # Verify it's actually hex and not just a long word
                try:
                    int(value, 16)
                except ValueError:
                    continue

            # Clean URL trailing characters
            if ioc_type == "url":
                value = value.rstrip(".,;:!?)]}>\"'")

            results[ioc_type].add(value)

    # Remove domains that are part of extracted URLs
    if "url" in results and "domain" in results:
        url_domains = set()
        for url in results["url"]:
            domain_match = re.search(r"https?://([^/:\s]+)", url)
            if domain_match:
                url_domains.add(domain_match.group(1).lower())
        # Keep domains not already captured by URLs
        # (don't remove - they're still valid IOCs)

    # Convert sets to sorted lists
    return {k: sorted(v) for k, v in results.items() if v}


def format_output(iocs: dict, fmt: str, defang: bool = False,
                  source: str = "") -> str:
    """Format IOC results in the specified format."""
    if fmt == "json":
        output = {
            "metadata": {
                "timestamp": datetime.utcnow().isoformat(),
                "source": source,
                "total_iocs": sum(len(v) for v in iocs.values()),
            },
            "indicators": {},
        }
        for ioc_type, values in iocs.items():
            output["indicators"][ioc_type] = []
            for v in values:
                entry = {"value": v}
                if defang:
                    entry["defanged"] = defang_ioc(v, ioc_type)
                output["indicators"][ioc_type].append(entry)
        return json.dumps(output, indent=2)

    elif fmt == "csv":
        lines = ["type,value,defanged"]
        for ioc_type, values in sorted(iocs.items()):
            for v in values:
                d = defang_ioc(v, ioc_type) if defang else ""
                # Escape commas in values
                v_escaped = f'"{v}"' if "," in v else v
                d_escaped = f'"{d}"' if "," in d else d
                lines.append(f"{ioc_type},{v_escaped},{d_escaped}")
        return "\n".join(lines)

    elif fmt == "text":
        lines = []
        for ioc_type, values in sorted(iocs.items()):
            lines.append(f"\n=== {ioc_type.upper()} ({len(values)}) ===")
            for v in values:
                if defang:
                    lines.append(f"  {defang_ioc(v, ioc_type)}")
                else:
                    lines.append(f"  {v}")
        return "\n".join(lines)

    elif fmt == "stix":
        # Simplified STIX 2.1 bundle output
        objects = []
        for ioc_type, values in iocs.items():
            for v in values:
                stix_type = _map_to_stix_type(ioc_type)
                if stix_type:
                    obj = {
                        "type": "indicator",
                        "spec_version": "2.1",
                        "pattern_type": "stix",
                        "pattern": f"[{stix_type} = '{v}']",
                        "valid_from": datetime.utcnow().isoformat() + "Z",
                    }
                    objects.append(obj)
        bundle = {"type": "bundle", "objects": objects}
        return json.dumps(bundle, indent=2)

    return ""


def _map_to_stix_type(ioc_type: str) -> Optional[str]:
    """Map internal IOC type to STIX 2.1 observable type."""
    mapping = {
        "ipv4": "ipv4-addr:value",
        "ipv6": "ipv6-addr:value",
        "domain": "domain-name:value",
        "url": "url:value",
        "email": "email-addr:value",
        "md5": "file:hashes.MD5",
        "sha1": "file:hashes.'SHA-1'",
        "sha256": "file:hashes.'SHA-256'",
    }
    return mapping.get(ioc_type)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract IOCs from text files (IPs, domains, URLs, hashes, etc.)"
    )
    parser.add_argument(
        "--input", "-i",
        required=True,
        help="Input file to extract IOCs from (use - for stdin)",
    )
    parser.add_argument(
        "--output", "-o",
        default=None,
        help="Output file path (default: stdout)",
    )
    parser.add_argument(
        "--format", "-f",
        choices=["json", "csv", "text", "stix"],
        default="json",
        help="Output format (default: json)",
    )
    parser.add_argument(
        "--defang",
        action="store_true",
        help="Defang IOCs in output for safe sharing",
    )
    parser.add_argument(
        "--types", "-t",
        nargs="+",
        choices=list(PATTERNS.keys()),
        default=None,
        help="IOC types to extract (default: all)",
    )
    parser.add_argument(
        "--source", "-s",
        default="",
        help="Source label for metadata",
    )

    args = parser.parse_args()

    # Read input
    try:
        if args.input == "-":
            text = sys.stdin.read()
        else:
            if not os.path.isfile(args.input):
                print(f"[!] Error: File not found: {args.input}", file=sys.stderr)
                sys.exit(1)
            with open(args.input, "r", errors="replace") as f:
                text = f.read()
    except Exception as e:
        print(f"[!] Error reading input: {e}", file=sys.stderr)
        sys.exit(1)

    # Extract IOCs
    iocs = extract_iocs(text, ioc_types=args.types)

    if not iocs:
        print("[-] No IOCs found in input", file=sys.stderr)
        sys.exit(0)

    # Format output
    output = format_output(
        iocs, fmt=args.format, defang=args.defang, source=args.source
    )

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

    # Print summary to stderr
    total = sum(len(v) for v in iocs.values())
    print(f"\n[+] Total IOCs extracted: {total}", file=sys.stderr)
    for ioc_type, values in sorted(iocs.items()):
        print(f"    {ioc_type}: {len(values)}", file=sys.stderr)


if __name__ == "__main__":
    main()