#!/usr/bin/env python3
"""Safe sample-handling helpers: hashing, provenance records, and IOC defanging.

Operates only on file bytes and text — it never executes a sample. Use it to compute a
sample's identity, emit a provenance record, and defang indicators before they leave the lab.
"""
from __future__ import annotations

import argparse
import datetime as _dt
import hashlib
import json
import re
import sys
from pathlib import Path

_DEFANG_SCHEME = re.compile(r"\bhttp(s?)://", re.I)
_DEFANG_DOT = re.compile(r"\.(?=[A-Za-z0-9])")
_IPV4 = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")


def compute_hashes(filepath: str | Path) -> dict:
    data = Path(filepath).read_bytes()
    return {
        "md5": hashlib.md5(data).hexdigest(),  # noqa: S324 - identification only
        "sha1": hashlib.sha1(data).hexdigest(),  # noqa: S324
        "sha256": hashlib.sha256(data).hexdigest(),
        "size": len(data),
    }


def make_record(filepath: str | Path, source: str, case: str) -> dict:
    rec = {
        "case": case,
        "source": source,
        "filename": Path(filepath).name,
        "recorded_at": _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    }
    rec.update(compute_hashes(filepath))
    return rec


def defang(text: str) -> str:
    """Render URLs/IPs non-clickable for safe sharing."""
    def _ip(m: re.Match) -> str:
        return m.group(0).replace(".", "[.]")
    text = _IPV4.sub(_ip, text)
    text = _DEFANG_SCHEME.sub(lambda m: "hxxp" + ("s" if m.group(1) else "") + "://", text)
    # bracket the dot in domains within URLs only (lines containing hxxp)
    out_lines = []
    for line in text.splitlines():
        if "hxxp" in line:
            line = _DEFANG_DOT.sub("[.]", line)
        out_lines.append(line)
    return "\n".join(out_lines)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Safe sample-handling helpers")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_hash = sub.add_parser("hash", help="compute file hashes")
    p_hash.add_argument("file")

    p_rec = sub.add_parser("record", help="emit a provenance JSON record")
    p_rec.add_argument("file")
    p_rec.add_argument("--source", required=True)
    p_rec.add_argument("--case", required=True)

    p_def = sub.add_parser("defang", help="defang URLs/IPs in a text file")
    p_def.add_argument("file")

    args = parser.parse_args(argv)

    if args.cmd == "hash":
        print(json.dumps(compute_hashes(args.file), indent=2))
    elif args.cmd == "record":
        print(json.dumps(make_record(args.file, args.source, args.case), indent=2))
    elif args.cmd == "defang":
        print(defang(Path(args.file).read_text(encoding="utf-8", errors="replace")))
    return 0


if __name__ == "__main__":
    sys.exit(main())