#!/usr/bin/env python3
"""Agent for performing malware IOC extraction from files, reports, and samples."""

import json
import argparse
import re
import hashlib
from datetime import datetime, timezone
from pathlib import Path


IOC_PATTERNS = {
    "ipv4": re.compile(r"\b(?:(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\.){3}(?:25[0-5]|2[0-4]\d|1\d\d|[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])?\.)+(?:com|net|org|io|ru|cn|xyz|top|info|biz|cc|tk|ml|ga|cf|gq|pw)\b"),
    "url": re.compile(r"https?://[^\s<>\"'\)]+"),
    "email": re.compile(r"\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b"),
    "md5": re.compile(r"\b[a-f0-9]{32}\b"),
    "sha1": re.compile(r"\b[a-f0-9]{40}\b"),
    "sha256": re.compile(r"\b[a-f0-9]{64}\b"),
    "cve": re.compile(r"CVE-\d{4}-\d{4,7}", re.I),
    "registry_key": re.compile(r"(?:HKLM|HKCU|HKCR|HKU|HKCC)\\[^\s\"']+"),
    "file_path_windows": re.compile(r"[A-Z]:\\(?:[^\s\\\"]+\\)*[^\s\\\"]+\.\w{1,5}"),
    "file_path_unix": re.compile(r"/(?:tmp|var|etc|usr|home|opt|bin|sbin)/[^\s\"']+"),
    "mutex": re.compile(r"(?:Global|Local)\\[^\s\"']+"),
    "bitcoin_addr": re.compile(r"\b[13][a-km-zA-HJ-NP-Z1-9]{25,34}\b"),
    "mitre_technique": re.compile(r"T\d{4}(?:\.\d{3})?"),
    "user_agent": re.compile(r"Mozilla/5\.0[^\n\"]{20,200}"),
}

DEFANGED_PATTERNS = {
    "ip_defanged": (re.compile(r"\b\d+\[\.\]\d+\[\.\]\d+\[\.\]\d+\b"), lambda m: m.group().replace("[.]", ".")),
    "url_defanged": (re.compile(r"hxxps?://[^\s]+"), lambda m: m.group().replace("hxxp", "http")),
    "domain_defanged": (re.compile(r"\b\S+\[\.\]\S+\b"), lambda m: m.group().replace("[.]", ".")),
}


def extract_iocs_from_text(text):
    """Extract all IOC types from raw text."""
    refanged = text
    for name, (pattern, fixer) in DEFANGED_PATTERNS.items():
        refanged = pattern.sub(fixer, refanged)
    extracted = {}
    for ioc_type, pattern in IOC_PATTERNS.items():
        matches = list(set(pattern.findall(refanged)))
        if matches:
            extracted[ioc_type] = sorted(matches)[:200]
    private_ip = re.compile(r"^(?:10\.|172\.(?:1[6-9]|2\d|3[01])\.|192\.168\.|127\.)")
    if "ipv4" in extracted:
        extracted["ipv4"] = [ip for ip in extracted["ipv4"] if not private_ip.match(ip)]
    return extracted


def extract_from_file(file_path):
    """Extract IOCs from a file (text, PDF text, or report)."""
    content = Path(file_path).read_text(encoding="utf-8", errors="replace")
    iocs = extract_iocs_from_text(content)
    total = sum(len(v) for v in iocs.values())
    return {
        "source": file_path, "total_iocs": total,
        "by_type": {k: len(v) for k, v in iocs.items()},
        "indicators": iocs,
    }


def hash_file(file_path):
    """Calculate file hashes for malware sample identification."""
    data = Path(file_path).read_bytes()
    return {
        "file": file_path,
        "size_bytes": len(data),
        "md5": hashlib.md5(data).hexdigest(),
        "sha1": hashlib.sha1(data).hexdigest(),
        "sha256": hashlib.sha256(data).hexdigest(),
    }


def extract_strings(file_path, min_length=6):
    """Extract printable strings from binary file."""
    data = Path(file_path).read_bytes()
    ascii_strings = re.findall(rb"[\x20-\x7e]{%d,}" % min_length, data)
    wide_strings = re.findall(rb"(?:[\x20-\x7e]\x00){%d,}" % min_length, data)
    all_strings = [s.decode("ascii", errors="replace") for s in ascii_strings]
    all_strings += [s.decode("utf-16-le", errors="replace") for s in wide_strings]
    iocs = extract_iocs_from_text("\n".join(all_strings))
    suspicious = []
    suspicious_kw = ["http", "socket", "connect", "download", "upload", "exec", "cmd.exe",
                     "powershell", "reg add", "CreateRemoteThread", "VirtualAlloc", "WriteProcessMemory",
                     "LoadLibrary", "GetProcAddress", "WinExec", "ShellExecute"]
    for s in all_strings:
        if any(kw.lower() in s.lower() for kw in suspicious_kw):
            suspicious.append(s[:200])
    return {
        "file": file_path, "total_strings": len(all_strings),
        "suspicious_strings": suspicious[:30],
        "extracted_iocs": {k: len(v) for k, v in iocs.items()},
        "ioc_details": iocs,
    }


def generate_ioc_report(file_path, output=None):
    """Generate comprehensive IOC extraction report."""
    hashes = hash_file(file_path)
    strings = extract_strings(file_path)
    report = {
        "generated": datetime.now(timezone.utc).isoformat(),
        "file_info": hashes,
        "strings_analysis": {
            "total": strings["total_strings"],
            "suspicious": strings["suspicious_strings"],
        },
        "extracted_iocs": strings["ioc_details"],
        "ioc_summary": strings["extracted_iocs"],
    }
    if output:
        with open(output, "w") as f:
            json.dump(report, f, indent=2)
    return report


def main():
    parser = argparse.ArgumentParser(description="Malware IOC Extraction Agent")
    sub = parser.add_subparsers(dest="command")
    t = sub.add_parser("text", help="Extract IOCs from text/report file")
    t.add_argument("--file", required=True)
    h = sub.add_parser("hash", help="Calculate file hashes")
    h.add_argument("--file", required=True)
    s = sub.add_parser("strings", help="Extract strings and IOCs from binary")
    s.add_argument("--file", required=True)
    s.add_argument("--min-length", type=int, default=6)
    r = sub.add_parser("report", help="Generate full IOC report")
    r.add_argument("--file", required=True)
    r.add_argument("--output", help="Output JSON file")
    args = parser.parse_args()
    if args.command == "text":
        result = extract_from_file(args.file)
    elif args.command == "hash":
        result = hash_file(args.file)
    elif args.command == "strings":
        result = extract_strings(args.file, args.min_length)
    elif args.command == "report":
        result = generate_ioc_report(args.file, args.output)
    else:
        parser.print_help()
        return
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()