116 lines
4.3 KiB
Python
116 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Static PE analysis: headers, sections, entropy, and import classification.
|
|
|
|
Uses ``pefile`` when available; falls back to a minimal built-in PE header walk so the tool
|
|
still reports sections and entropy without third-party packages. Never executes the sample.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import struct
|
|
import sys
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
_IMPORT_BUCKETS = {
|
|
"injection": {"virtualallocex", "writeprocessmemory", "createremotethread",
|
|
"ntunmapviewofsection", "queueuserapc", "setthreadcontext"},
|
|
"network": {"internetopena", "internetopenw", "httpsendrequesta", "winhttpconnect",
|
|
"wsastartup", "connect", "send", "urldownloadtofilea"},
|
|
"crypto": {"cryptencrypt", "cryptdecrypt", "cryptacquirecontexta",
|
|
"cryptgenkey", "bcryptencrypt"},
|
|
"persistence": {"regsetvalueexa", "regsetvalueexw", "createservicea",
|
|
"createservicew", "schtasks"},
|
|
"evasion": {"isdebuggerpresent", "checkremotedebuggerpresent",
|
|
"gettickcount", "loadlibrarya", "getprocaddress"},
|
|
}
|
|
|
|
|
|
def entropy(data: bytes) -> float:
|
|
if not data:
|
|
return 0.0
|
|
counts = Counter(data)
|
|
n = len(data)
|
|
return -sum((c / n) * math.log2(c / n) for c in counts.values())
|
|
|
|
|
|
def analyze_with_pefile(path: str) -> dict:
|
|
import pefile # type: ignore
|
|
pe = pefile.PE(path, fast_load=True)
|
|
pe.parse_data_directories()
|
|
sections = [{
|
|
"name": s.Name.rstrip(b"\x00").decode("ascii", "replace"),
|
|
"vsize": s.Misc_VirtualSize,
|
|
"rawsize": s.SizeOfRawData,
|
|
"entropy": round(s.get_entropy(), 3),
|
|
} for s in pe.sections]
|
|
imports: list[str] = []
|
|
if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
|
|
for entry in pe.DIRECTORY_ENTRY_IMPORT:
|
|
for imp in entry.imports:
|
|
if imp.name:
|
|
imports.append(imp.name.decode("ascii", "replace"))
|
|
buckets = _classify(imports)
|
|
return {
|
|
"is_dll": pe.is_dll(),
|
|
"machine": hex(pe.FILE_HEADER.Machine),
|
|
"timestamp": pe.FILE_HEADER.TimeDateStamp,
|
|
"sections": sections,
|
|
"import_count": len(imports),
|
|
"import_buckets": buckets,
|
|
"tls_callbacks": hasattr(pe, "DIRECTORY_ENTRY_TLS"),
|
|
}
|
|
|
|
|
|
def _classify(imports: list[str]) -> dict[str, list[str]]:
|
|
low = [i.lower() for i in imports]
|
|
return {bucket: sorted({i for i in low if i in names})
|
|
for bucket, names in _IMPORT_BUCKETS.items()
|
|
if any(i in names for i in low)}
|
|
|
|
|
|
def analyze_minimal(path: str) -> dict:
|
|
"""Fallback header walk without pefile: sections + entropy only."""
|
|
data = Path(path).read_bytes()
|
|
if data[:2] != b"MZ":
|
|
raise ValueError("not a PE (no MZ header)")
|
|
e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
|
|
if data[e_lfanew:e_lfanew + 4] != b"PE\x00\x00":
|
|
raise ValueError("not a PE (no PE signature)")
|
|
num_sections = struct.unpack_from("<H", data, e_lfanew + 6)[0]
|
|
opt_size = struct.unpack_from("<H", data, e_lfanew + 20)[0]
|
|
sec_off = e_lfanew + 24 + opt_size
|
|
sections = []
|
|
for i in range(num_sections):
|
|
base = sec_off + i * 40
|
|
name = data[base:base + 8].rstrip(b"\x00").decode("ascii", "replace")
|
|
raw_size = struct.unpack_from("<I", data, base + 16)[0]
|
|
raw_ptr = struct.unpack_from("<I", data, base + 20)[0]
|
|
blob = data[raw_ptr:raw_ptr + raw_size]
|
|
sections.append({"name": name, "rawsize": raw_size,
|
|
"entropy": round(entropy(blob), 3)})
|
|
return {"sections": sections, "note": "pefile not installed; minimal parse"}
|
|
|
|
|
|
def analyze(path: str) -> dict:
|
|
try:
|
|
return analyze_with_pefile(path)
|
|
except ImportError:
|
|
return analyze_minimal(path)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description="Static PE analysis")
|
|
sub = parser.add_subparsers(dest="cmd", required=True)
|
|
p = sub.add_parser("analyze")
|
|
p.add_argument("file")
|
|
args = parser.parse_args(argv)
|
|
print(json.dumps(analyze(args.file), indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|