115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
#!/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())
|