#!/usr/bin/env python3
"""Extract IOCs (IPs, domains, URLs, emails, hashes, BTC, regkeys, paths) from
malware-analysis evidence text (Procmon CSV, tshark, strings, olevba, Sysmon
JSON, deobfuscated scripts, ...) and print them defanged.

Usage: ioc_extract.py <file> [<file>...]        (or stdin when no args / '-')
"""
import argparse
import csv
import io
import json
import re
import sys
from collections import OrderedDict

TYPES = ["ipv4", "url", "domain", "email", "md5", "sha1", "sha256", "btc", "regkey", "winpath"]

FILE_EXT_TLDS = {
    "exe", "dll", "txt", "dat", "log", "tmp", "ini", "py", "js", "vbs", "ps1",
    "bat", "cmd", "lnk", "zip", "rar", "7z", "pdf", "doc", "docx", "xls",
    "xlsx", "png", "jpg", "gif", "css", "html", "htm", "xml", "json", "cfg",
    "sys", "drv", "ocx", "cpl", "scr", "pdb", "bin", "db", "sqlite", "config", "lock", "old", "bak",
}

ALLOWLIST = {
    "microsoft.com", "windows.com", "google.com", "gstatic.com", "msftncsi.com",
    "mozilla.org", "digicert.com", "verisign.com", "symantec.com", "symcb.com",
    "symcd.com", "w3.org", "schemas.microsoft.com", "adobe.com", "apple.com",
}

RE = {
    "url": re.compile(r"\b(?:https?|ftp|hxxps?|fxp)://[^\s\"'<>]+", re.I),
    "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,24}\b"),
    "ipv4": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
    "sha256": re.compile(r"\b[a-fA-F0-9]{64}\b"),
    "sha1": re.compile(r"\b[a-fA-F0-9]{40}\b"),
    "md5": re.compile(r"\b[a-fA-F0-9]{32}\b"),
    "btc": re.compile(r"\b(?:[13][a-km-zA-HJ-NP-Z1-9]{25,34}|bc1[a-z0-9]{25,59})\b"),
    "regkey": re.compile(
        r"\b(?:HKEY_LOCAL_MACHINE|HKEY_CURRENT_USER|HKEY_CLASSES_ROOT|HKEY_USERS|"
        r"HKEY_CURRENT_CONFIG|HKLM|HKCU|HKCR|HKU)\\[^\s\"',]+"
    ),
    "winpath": re.compile(r"(?:(?<![A-Za-z0-9])[A-Za-z]:\\[^\s\"',]+|%[A-Za-z0-9_]+%[^\s\"',]*|\\\\[^\s\"',]+)"),
    "domain": re.compile(r"\b(?:[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?\.)+[A-Za-z]{2,24}\b"),
}


def is_private_ipv4(ip):
    parts = [int(p) for p in ip.split(".")]
    if any(p > 255 for p in parts):
        return None  # not a real IP
    a, b = parts[0], parts[1]
    if a == 10 or (a == 172 and 16 <= b <= 31) or (a == 192 and b == 168):
        return "private-ip"
    if a == 127:
        return "loopback"
    if a == 169 and b == 254:
        return "link-local"
    return ""


def defang(kind, value):
    # Normalise first so already-defanged input (hxxp://a[.]b) is not double-bracketed.
    value = refang(value)
    if kind == "ipv4":
        return value.replace(".", "[.]")
    if kind in ("url", "domain", "email"):
        v = re.sub(r"(?i)^https", "hxxps", value)
        v = re.sub(r"(?i)^http", "hxxp", v)
        v = re.sub(r"(?i)^ftp", "fxp", v)
        v = v.replace(".", "[.]").replace("@", "[@]")
        return v
    return value  # hashes, btc, regkey, winpath are not defanged


def refang(text):
    text = re.sub(r"(?i)hxxps", "https", text)
    text = re.sub(r"(?i)hxxp", "http", text)
    text = re.sub(r"(?i)\bfxp://", "ftp://", text)
    text = text.replace("[.]", ".").replace("(.)", ".").replace("[@]", "@")
    return text


def is_extension_like(domain, preceding_char):
    tld = domain.rsplit(".", 1)[-1].lower()
    if tld != "com":
        return tld in FILE_EXT_TLDS
    return preceding_char == "\\"  # looks like C:\...\payload.com


def extract(text, types, use_allowlist):
    """Yield (kind, raw_value, note) in first-seen order for one chunk of text."""
    spans_to_blank = []
    out = []

    def want(kind):
        return types is None or kind in types

    # url / email first, then blank their spans so domain regex doesn't
    # double-extract the hostname portion out of them.
    for kind in ("url", "email"):
        for m in RE[kind].finditer(text):
            if want(kind):
                out.append((kind, m.group(0), ""))
            spans_to_blank.append(m.span())

    if want("ipv4"):
        for m in RE["ipv4"].finditer(text):
            note = is_private_ipv4(m.group(0))
            if note is None:
                continue
            out.append(("ipv4", m.group(0), note))

    if want("btc"):
        for m in RE["btc"].finditer(text):
            out.append(("btc", m.group(0), ""))

    for kind in ("sha256", "sha1", "md5"):
        if want(kind):
            for m in RE[kind].finditer(text):
                out.append((kind, m.group(0), ""))

    if want("regkey"):
        for m in RE["regkey"].finditer(text):
            out.append(("regkey", m.group(0), ""))

    if want("winpath"):
        for m in RE["winpath"].finditer(text):
            out.append(("winpath", m.group(0), ""))

    if want("domain"):
        masked = text
        for start, end in spans_to_blank:
            masked = masked[:start] + (" " * (end - start)) + masked[end:]
        for m in RE["domain"].finditer(masked):
            domain = m.group(0)
            preceding = masked[m.start() - 1] if m.start() > 0 else ""
            if is_extension_like(domain, preceding):
                continue
            note = ""
            if use_allowlist and domain.lower() in ALLOWLIST:
                note = "allowlisted-skipped-not-shown"
                out.append(("domain", domain, note))
                continue
            out.append(("domain", domain, note))

    return out


def apply_winpath_filter(items, all_paths):
    result = []
    for kind, value, note in items:
        if kind == "winpath":
            low = value.lower()
            if not all_paths and ("\\windows\\system32" in low or "\\windows\\winsxs" in low):
                continue
        result.append((kind, value, note))
    return result


def detect_encoding(head):
    if head[:2] == b"\xff\xfe":
        return "utf-16-le"
    if head[:2] == b"\xfe\xff":
        return "utf-16-be"
    return "utf-8"


def read_text(path):
    """Read a whole file/stream as text (used by --self-test only)."""
    raw = sys.stdin.buffer.read() if path == "-" else open(path, "rb").read()
    return raw.decode(detect_encoding(raw), errors="replace")


def iter_lines(path):
    """Yield decoded lines, streaming from disk so multi-GB evidence files
    (Procmon CSV, Sysmon JSON) never have to fit in memory at once."""
    if path == "-":
        raw = sys.stdin.buffer.read()
        yield from raw.decode(detect_encoding(raw), errors="replace").splitlines()
        return
    with open(path, "rb") as f:
        enc = detect_encoding(f.read(2))
        f.seek(0)
        for line in io.TextIOWrapper(f, encoding=enc, errors="replace", newline=""):
            yield line.rstrip("\r\n")


def collect(paths, args):
    types = set(args.types.split(",")) if args.types else None
    use_allowlist = not args.no_allowlist
    counts = OrderedDict()  # (kind, normalized_value) -> [display_value, count, note]

    for path in paths:
        for line in iter_lines(path):
            if args.refang:
                line = refang(line)
            items = extract(line, types, use_allowlist)
            items = apply_winpath_filter(items, args.all_paths)
            for kind, value, note in items:
                if note == "allowlisted-skipped-not-shown":
                    key = (kind, value.lower())
                    if key not in counts:
                        counts[key] = [value, 0, note]
                    counts[key][1] += 1
                    continue
                norm = refang(value).lower() if kind in ("ipv4", "url", "domain", "email", "md5", "sha1", "sha256") else value
                key = (kind, norm)
                if key not in counts:
                    display = value if args.no_defang or args.refang else defang(kind, value)
                    counts[key] = [display, 0, note]
                counts[key][1] += 1
    return counts


def output(counts, fmt):
    rows = [(k[0], v[0], v[1], v[2]) for k, v in counts.items()]
    if fmt == "json":
        print(json.dumps(
            [{"type": t, "value": v, "count": c, "note": n} for t, v, c, n in rows],
            indent=2,
        ))
    elif fmt == "csv":
        w = csv.writer(sys.stdout)
        w.writerow(["type", "value", "count", "note"])
        for t, v, c, n in rows:
            w.writerow([t, v, c, n])
    else:
        by_type = OrderedDict()
        for t, v, c, n in rows:
            by_type.setdefault(t, []).append((v, c, n))
        for t in by_type:
            print(f"== {t} ({len(by_type[t])}) ==")
            for v, c, n in by_type[t]:
                suffix = f"  [{n}]" if n else ""
                print(f"  {v}  (x{c}){suffix}")
            print()


def self_test():
    sample = (
        "Visit http://evil[.]com/a[.]exe and hxxps://bad-tld[.]tk/x now.\n"
        "Also seen evil.tk in the config, and payload.exe is not a domain.\n"
        "Beacon to 10.0.0.5 and 8.8.8.8 and email attacker@evil.com\n"
        "Hash: " + "a" * 64 + "\n"
        "Key HKCU\\Software\\Microsoft\\Windows\\CurrentVersion\\Run\\x\n"
        "BTC 1BoatSLRHtKNngkdXEeobR76b53LETtpyT talks to microsoft.com too\n"
    )
    items = extract(sample, None, True)
    items = apply_winpath_filter(items, False)
    kinds = {(k, v) for k, v, n in items}

    assert any(k == "url" and "evil" in v for k, v in kinds), "url not found"
    live = refang("hxxps://bad-tld[.]tk/x")
    assert live == "https://bad-tld.tk/x", "refang roundtrip failed"
    assert defang("url", live) == "hxxps://bad-tld[.]tk/x", "defang(refang(x)) != x"

    priv = [n for k, v, n in items if k == "ipv4" and v == "10.0.0.5"]
    assert priv and priv[0] == "private-ip", "private ip note missing"

    domains = [v for k, v, n in items if k == "domain"]
    assert not any("payload" in d for d in domains), "filename wrongly kept as domain"
    assert any("evil.tk" == d for d in domains), "evil.tk not kept"

    assert any(k == "email" and v == "attacker@evil.com" for k, v in kinds), "email not found"
    assert any(k == "sha256" for k, v in kinds), "sha256 not found"
    assert any(k == "regkey" and v.startswith("HKCU\\") for k, v in kinds), "regkey not found"
    assert any(k == "btc" for k, v in kinds), "btc not found"

    ms_domains = [(v, n) for k, v, n in items if k == "domain" and "microsoft" in v.lower()]
    assert ms_domains and ms_domains[0][1] == "allowlisted-skipped-not-shown", "allowlist not applied"

    import tempfile
    utf16 = b"\xff\xfe" + "192.168.1.1".encode("utf-16-le")
    with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
        f.write(utf16)
        tmp_path = f.name
    text = read_text(tmp_path)
    assert "192.168.1.1" in text, "utf-16 decode failed"

    print("self-test: all checks passed")


def main():
    ap = argparse.ArgumentParser(description="Extract and defang IOCs from evidence files")
    ap.add_argument("files", nargs="*", default=["-"])
    ap.add_argument("--format", choices=["text", "csv", "json"], default="text")
    ap.add_argument("--no-defang", action="store_true")
    ap.add_argument("--types", help="comma list restricting types, e.g. ipv4,domain,url")
    ap.add_argument("--refang", action="store_true", help="input is defanged text; output live indicators")
    ap.add_argument("--no-allowlist", action="store_true")
    ap.add_argument("--all-paths", action="store_true", help="don't skip System32/WinSxS paths")
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()

    if args.self_test:
        self_test()
        return

    paths = args.files
    for p in paths:
        if p != "-":
            try:
                open(p, "rb").close()
            except OSError as e:
                print(f"ioc_extract.py: cannot read {p}: {e}", file=sys.stderr)
                sys.exit(1)

    counts = collect(paths, args)
    output(counts, args.format)


if __name__ == "__main__":
    main()