#!/usr/bin/env python3
"""Structurally diff two PE samples: section hashes/entropy, import sets, similarity.

Parses both PEs' section tables and import name sets, reports added/removed/changed sections and
imports, and computes a byte-level similarity ratio (plus ssdeep/tlsh if installed). Uses pefile
for imports when available; otherwise imports are best-effort. Reads bytes statically; executes
nothing.
"""
from __future__ import annotations

import argparse
import difflib
import hashlib
import json
import math
import struct
import sys
from collections import Counter
from pathlib import Path

try:
    import pefile  # type: ignore
    HAVE_PEFILE = True
except ImportError:
    HAVE_PEFILE = False

try:
    import ssdeep  # type: ignore
    HAVE_SSDEEP = True
except ImportError:
    HAVE_SSDEEP = False


def _entropy(chunk: bytes) -> float:
    if not chunk:
        return 0.0
    counts = Counter(chunk)
    total = len(chunk)
    return round(-sum((c / total) * math.log2(c / total) for c in counts.values()), 3)


def _sections(data: bytes) -> dict[str, dict]:
    out: dict[str, dict] = {}
    if data[:2] != b"MZ":
        return out
    e_lfanew = struct.unpack_from("<I", data, 0x3c)[0]
    if data[e_lfanew:e_lfanew + 4] != b"PE\x00\x00":
        return out
    num = struct.unpack_from("<H", data, e_lfanew + 6)[0]
    opt = struct.unpack_from("<H", data, e_lfanew + 20)[0]
    sect_off = e_lfanew + 24 + opt
    for i in range(num):
        base = sect_off + i * 40
        raw = data[base:base + 40]
        if len(raw) < 40:
            break
        name = raw[:8].rstrip(b"\x00").decode("latin-1", "ignore")
        _, _, rawsize, rawptr = struct.unpack_from("<IIII", raw, 8)
        body = data[rawptr:rawptr + rawsize]
        out[name] = {"sha256": hashlib.sha256(body).hexdigest(), "entropy": _entropy(body)}
    return out


def _imports(data: bytes) -> set[str]:
    if not HAVE_PEFILE:
        return set()
    try:
        pe = pefile.PE(data=data, fast_load=True)
        pe.parse_data_directories(
            directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"]])
    except Exception:
        return set()
    names = set()
    for entry in getattr(pe, "DIRECTORY_ENTRY_IMPORT", []) or []:
        for imp in entry.imports:
            if imp.name:
                names.add(imp.name.decode("latin-1", "ignore"))
    return names


def diff(path_a: str, path_b: str) -> dict:
    a = Path(path_a).read_bytes()
    b = Path(path_b).read_bytes()
    sa, sb = _sections(a), _sections(b)
    changed = sorted(n for n in sa.keys() & sb.keys() if sa[n]["sha256"] != sb[n]["sha256"])
    ia, ib = _imports(a), _imports(b)
    result = {
        "size_a": len(a), "size_b": len(b),
        "sections_added": sorted(sb.keys() - sa.keys()),
        "sections_removed": sorted(sa.keys() - sb.keys()),
        "sections_changed": changed,
        "imports_added": sorted(ib - ia)[:100],
        "imports_removed": sorted(ia - ib)[:100],
        "byte_similarity_ratio": round(difflib.SequenceMatcher(None, a, b).quick_ratio(), 4),
        "import_diff_available": HAVE_PEFILE,
    }
    if HAVE_SSDEEP:
        result["ssdeep_score"] = ssdeep.compare(ssdeep.hash(a), ssdeep.hash(b))
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Diff two PE samples structurally")
    sub = parser.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("diff")
    p.add_argument("file_a")
    p.add_argument("file_b")
    args = parser.parse_args(argv)
    print(json.dumps(diff(args.file_a, args.file_b), indent=2))
    return 0


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