NJL/day-zero
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
day-zero/day-zero.py
T
NJL 8484b78838 Add day-zero.py
2026-07-06 18:50:37 +00:00

2561 lines
90 KiB
Python
Raw Permalink Blame History

import argparse
import bisect
import json
import struct
import sys
from collections import defaultdict
from dataclasses import dataclass, field, asdict
import pefile
from capstone import (Cs, CS_ARCH_X86, CS_MODE_64, CS_MODE_32,
CS_GRP_JUMP, CS_GRP_CALL, CS_GRP_RET, CS_GRP_INT,
CS_OP_REG, CS_OP_IMM, CS_OP_MEM)
from capstone import x86 as csx86
from capstone.x86 import X86_OP_IMM, X86_OP_MEM, X86_OP_REG, X86_REG_RIP
_TRACE = False
def set_trace(on):
global _TRACE
_TRACE = on
def log(msg, indent=0):
if _TRACE:
print((" " * indent) + msg, flush=True)
@dataclass
class FoundString:
encoding: str
file_offset: int
vaddr: int | None
section: str | None
text: str
def _section_for_offset(pe, file_offset):
for s in pe.sections:
start = s.PointerToRawData
end = start + s.SizeOfRawData
if start <= file_offset < end:
name = s.Name.rstrip(b"\x00").decode("latin-1", "replace")
rva = s.VirtualAddress + (file_offset - start)
vaddr = pe.OPTIONAL_HEADER.ImageBase + rva
return name, vaddr
return None, None
def extract_strings(pe, data, min_len=4):
results = []
run_start = None
for i, b in enumerate(data):
printable = 0x20 <= b <= 0x7E
if printable and run_start is None:
run_start = i
elif not printable and run_start is not None:
if i - run_start >= min_len:
sec, va = _section_for_offset(pe, run_start)
results.append(FoundString(
"ascii", run_start, va, sec,
data[run_start:i].decode("ascii")))
run_start = None
if run_start is not None and len(data) - run_start >= min_len:
sec, va = _section_for_offset(pe, run_start)
results.append(FoundString(
"ascii", run_start, va, sec, data[run_start:].decode("ascii")))
i = 0
n = len(data)
while i < n - 1:
if 0x20 <= data[i] <= 0x7E and data[i + 1] == 0x00:
start = i
chars = []
while i < n - 1 and 0x20 <= data[i] <= 0x7E and data[i + 1] == 0x00:
chars.append(chr(data[i]))
i += 2
if len(chars) >= min_len:
sec, va = _section_for_offset(pe, start)
results.append(FoundString(
"utf-16le", start, va, sec, "".join(chars)))
else:
i += 1
return results
@dataclass
class Function:
name: str | None
vaddr: int
rva: int
size: int | None
source: str
insn_count: int | None = None
decoded_ok: bool = True
def _rva_to_offset(pe, rva):
try:
return pe.get_offset_from_rva(rva)
except Exception:
return None
def collect_exports(pe):
out = {}
if not hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
return out
for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
if exp.address == 0:
continue
name = exp.name.decode("latin-1") if exp.name else f"ordinal_{exp.ordinal}"
out[exp.address] = name
return out
def collect_pdata_functions(pe):
ranges = []
try:
entry = pe.OPTIONAL_HEADER.DATA_DIRECTORY[
pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXCEPTION"]]
except Exception:
return ranges
if not entry or entry.VirtualAddress == 0 or entry.Size == 0:
return ranges
off = _rva_to_offset(pe, entry.VirtualAddress)
if off is None:
return ranges
raw = bytes(pe.__data__[off:off + entry.Size])
for i in range(0, len(raw) - 11, 12):
begin, end, _unwind = struct.unpack_from("<III", raw, i)
if begin == 0 and end == 0:
continue
ranges.append((begin, end))
return ranges
def _md_for(pe, detail=False):
machine = pe.FILE_HEADER.Machine
if machine == 0x8664:
md = Cs(CS_ARCH_X86, CS_MODE_64)
elif machine == 0x14C:
md = Cs(CS_ARCH_X86, CS_MODE_32)
else:
return None
md.detail = detail
return md
IMAGE_SCN_MEM_EXECUTE = 0x20000000
def executable_ranges(pe):
base = pe.OPTIONAL_HEADER.ImageBase
ranges = []
for s in pe.sections:
if not (s.Characteristics & IMAGE_SCN_MEM_EXECUTE):
continue
start = base + s.VirtualAddress
raw = s.SizeOfRawData
size = min(s.Misc_VirtualSize, raw) if raw else s.Misc_VirtualSize
ranges.append((start, start + size, s.PointerToRawData, raw))
return ranges
def _make_va_reader(pe):
data = bytes(pe.__data__)
ranges = executable_ranges(pe)
def in_exec(va):
for start, end, _off, _raw in ranges:
if start <= va < end:
return True
return False
def read(va, n):
for start, end, off, raw in ranges:
if start <= va < end:
delta = va - start
avail = min(n, raw - delta, end - va)
if avail <= 0:
return b""
return data[off + delta:off + delta + avail]
return b""
return in_exec, read, ranges
def recover_functions(pe, disassemble=True, discover=True):
exports = collect_exports(pe)
pdata = collect_pdata_functions(pe)
image_base = pe.OPTIONAL_HEADER.ImageBase
md = _md_for(pe) if disassemble else None
pdata_by_start = {b: e for b, e in pdata}
starts = {}
for begin, end in pdata:
starts[begin] = Function(
name=exports.get(begin),
vaddr=image_base + begin, rva=begin,
size=end - begin,
source="export+pdata" if begin in exports else "pdata")
for rva, name in exports.items():
if rva not in starts:
size = None
end = pdata_by_start.get(rva)
if end:
size = end - rva
starts[rva] = Function(
name=name, vaddr=image_base + rva, rva=rva,
size=size, source="export")
discovery_stats = None
if discover:
seed_list = [starts[k] for k in sorted(starts)]
new_only, discovery_stats = discover_functions(pe, seed_list)
pdata_ends = sorted(e for _b, e in pdata)
for va, src in new_only.items():
rva = va - image_base
size = _size_from_next_start(rva, starts, pdata_by_start, pdata_ends, image_base)
starts[rva] = Function(
name=None, vaddr=va, rva=rva, size=size, source=src)
functions = [starts[k] for k in sorted(starts)]
if md:
for fn in functions:
off = _rva_to_offset(pe, fn.rva)
if off is None:
fn.decoded_ok = False
continue
size = fn.size if fn.size else 64
code = bytes(pe.__data__[off:off + size])
try:
count = sum(1 for _ in md.disasm(code, fn.vaddr))
fn.insn_count = count
fn.decoded_ok = count > 0
except Exception:
fn.decoded_ok = False
return functions, discovery_stats
def _size_from_next_start(rva, starts, pdata_by_start, pdata_ends, image_base):
if rva in pdata_by_start:
return pdata_by_start[rva] - rva
import bisect
idx = bisect.bisect_right(pdata_ends, rva)
later = [s for s in starts if s > rva]
next_start = min(later) if later else None
if next_start is not None:
return next_start - rva
return None
X64_PROLOGUES = [
b"\x55\x48\x8b\xec", # push rbp; mov rbp, rsp
b"\x48\x89\x5c\x24", # mov [rsp+X], rbx
b"\x48\x83\xec", # sub rsp, imm8
b"\x48\x81\xec", # sub rsp, imm32
b"\x40\x53", # push rbx (REX)
b"\x40\x55", # push rbp (REX)
b"\x40\x56", # push rsi (REX)
b"\x40\x57", # push rdi (REX)
b"\x48\x89\x4c\x24", # mov [rsp+X], rcx
]
X86_PROLOGUES = [
b"\x55\x8b\xec", # push ebp; mov ebp, esp
b"\x53\x56\x57", # push ebx; push esi; push edi
b"\x8b\xff\x55\x8b\xec", # mov edi,edi; push ebp; mov ebp,esp (hotpatch)
b"\x83\xec", # sub esp, imm8
b"\x81\xec", # sub esp, imm32
]
def _traverse(md, read, in_exec, seeds, max_insns_per_fn=100000, tag=""):
discovered = {}
worklist = list(seeds)
visited_starts = set()
processed = 0
while worklist:
va = worklist.pop()
if va in visited_starts or not in_exec(va):
continue
visited_starts.add(va)
processed += 1
if _TRACE and processed % 1000 == 0:
log(f"[traverse{tag}] {processed} functions processed, "
f"{len(discovered)} mapped, worklist {len(worklist)}", 2)
pending = [va]
seen_blocks = set()
insn_count = 0
while pending:
addr = pending.pop()
if addr in seen_blocks or not in_exec(addr):
continue
seen_blocks.add(addr)
code = read(addr, 4096)
if not code:
continue
for ins in md.disasm(code, addr):
insn_count += 1
if insn_count > max_insns_per_fn:
break
g = set(ins.groups)
if CS_GRP_CALL in g:
t = _direct_target(ins)
if t is not None and in_exec(t) and t not in visited_starts:
worklist.append(t)
pending.append(ins.address + ins.size)
break
if CS_GRP_RET in g:
break
if CS_GRP_JUMP in g:
t = _direct_target(ins)
if ins.mnemonic == "jmp":
if t is not None and in_exec(t):
pending.append(t)
break
else:
if t is not None and in_exec(t):
pending.append(t)
pending.append(ins.address + ins.size)
break
if CS_GRP_INT in g:
pending.append(ins.address + ins.size)
break
else:
continue
discovered[va] = insn_count
return discovered
def _prologue_scan(pe, read, ranges, known_starts):
machine = pe.FILE_HEADER.Machine
sigs = X64_PROLOGUES if machine == 0x8664 else X86_PROLOGUES
data = bytes(pe.__data__)
hits = set()
for start, end, off, raw in ranges:
if raw <= 0:
continue
blob = data[off:off + raw]
for sig in sigs:
idx = blob.find(sig)
while idx != -1:
va = start + idx
if va not in known_starts and (va % 16 == 0 or blob[idx - 1:idx] in (b"\xcc", b"\x90", b"")):
hits.add(va)
idx = blob.find(sig, idx + 1)
return hits
def discover_functions(pe, seed_functions):
md = _md_for(pe, detail=True)
if md is None:
return {}, {"error": "unsupported arch"}
in_exec, read, ranges = _make_va_reader(pe)
base = pe.OPTIONAL_HEADER.ImageBase
log("[discover] seeding from entry point + exports + .pdata starts", 1)
seeds = set()
ep = base + pe.OPTIONAL_HEADER.AddressOfEntryPoint
if in_exec(ep):
seeds.add(ep)
log(f"[discover] entry point seed {hex(ep)}", 2)
for fn in seed_functions:
if in_exec(fn.vaddr):
seeds.add(fn.vaddr)
log(f"[discover] {len(seeds)} seeds total; starting recursive traversal", 2)
known = {fn.vaddr for fn in seed_functions}
reachable = _traverse(md, read, in_exec, seeds)
log(f"[discover] recursive traversal reached {len(reachable)} functions", 2)
log("[discover] scanning executable sections for function prologues", 1)
prologue_hits = _prologue_scan(pe, read, ranges, known | set(reachable))
log(f"[discover] {len(prologue_hits)} prologue candidates; validating", 2)
prologue_reached = _traverse(md, read, in_exec, prologue_hits)
found = {}
for va in reachable:
found[va] = "recursive" if va not in known else "seed"
for va in prologue_reached:
if va not in found:
found[va] = "prologue"
new_only = {va: src for va, src in found.items() if va not in known}
log(f"[discover] {len(new_only)} NEW functions beyond seeds "
f"({sum(1 for v in found.values() if v=='recursive')} recursive, "
f"{sum(1 for v in found.values() if v=='prologue')} prologue)", 1)
stats = {
"seeds": len(seeds),
"seed_functions": len(known),
"recursive_discovered": sum(1 for v in found.values() if v == "recursive"),
"prologue_discovered": sum(1 for v in found.values() if v == "prologue"),
"new_total": len(new_only),
}
if known:
confirmed = sum(1 for va in reachable if va in known)
stats["seed_confirmed_by_recursion"] = confirmed
stats["seed_coverage_pct"] = round(100 * confirmed / len(known), 1)
return new_only, stats
@dataclass
class BasicBlock:
start: int
end: int
insn_count: int
succ: list = field(default_factory=list)
kind: str = "fallthrough"
@dataclass
class FunctionCFG:
func_vaddr: int
blocks: dict = field(default_factory=dict)
calls: list = field(default_factory=list)
icalls: list = field(default_factory=list)
@property
def block_count(self):
return len(self.blocks)
@property
def edge_count(self):
return sum(len(b.succ) for b in self.blocks.values())
def build_iat_map(pe):
iat = {}
if not hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
return iat
for entry in pe.DIRECTORY_ENTRY_IMPORT:
dll = entry.dll.decode("latin-1")
for imp in entry.imports:
if imp.address is None:
continue
nm = imp.name.decode("latin-1") if imp.name else f"ordinal_{imp.ordinal}"
iat[imp.address] = f"{dll}!{nm}"
return iat
def _direct_target(ins):
if ins.operands and ins.operands[0].type == X86_OP_IMM:
return ins.operands[0].imm
return None
def _rip_mem_target(ins):
for op in ins.operands:
if op.type == X86_OP_MEM and op.mem.base == X86_REG_RIP and op.mem.index == 0:
return ins.address + ins.size + op.mem.disp
return None
def build_cfg_for_function(md, code, func_va, func_starts, iat):
insns = list(md.disasm(code, func_va))
if not insns:
return None, False
end_va = insns[-1].address + insns[-1].size
addr_set = {i.address for i in insns}
leaders = {func_va}
flow = {}
cfg = FunctionCFG(func_vaddr=func_va)
for idx, ins in enumerate(insns):
g = set(ins.groups)
nxt = insns[idx + 1].address if idx + 1 < len(insns) else end_va
kind = "fallthrough"
succ = []
call_target = None
icall = None
if CS_GRP_CALL in g:
kind = "call"
t = _direct_target(ins)
if t is not None:
call_target = t
else:
slot = _rip_mem_target(ins)
icall = iat.get(slot, "?") if slot is not None else "?"
succ = [nxt]
elif CS_GRP_RET in g:
kind = "ret"
elif CS_GRP_INT in g:
kind = "int"
succ = [nxt]
elif CS_GRP_JUMP in g:
t = _direct_target(ins)
if ins.mnemonic == "jmp":
kind = "jmp"
if t is not None and t in addr_set:
succ = [t]
elif t is not None and t in func_starts:
kind = "tailcall"
call_target = t
else:
kind = "cond"
succ = [nxt]
if t is not None and t in addr_set:
succ.append(t)
if t is not None and t in addr_set:
leaders.add(t)
if nxt in addr_set:
leaders.add(nxt)
else:
succ = [nxt] if nxt in addr_set else []
if kind in ("call", "ret", "int") and nxt in addr_set:
leaders.add(nxt)
flow[ins.address] = (kind, succ, call_target, icall)
if call_target is not None and kind != "tailcall":
cfg.calls.append(call_target)
if icall is not None:
cfg.icalls.append(icall)
leader_set = set(leaders)
cur = None
for idx, ins in enumerate(insns):
if ins.address in leader_set or cur is None:
cur = BasicBlock(start=ins.address, end=ins.address, insn_count=0)
cfg.blocks[cur.start] = cur
cur.insn_count += 1
cur.end = ins.address + ins.size
kind, succ, _ct, _ic = flow[ins.address]
nxt = insns[idx + 1].address if idx + 1 < len(insns) else None
if nxt is None or nxt in leader_set:
cur.kind = kind
cur.succ = [s for s in succ if s in leader_set]
return cfg, True
def build_cfgs(pe, functions, cap_functions=None):
md = _md_for(pe, detail=True)
if md is None:
return {}, {"direct": {}, "indirect": {}}, {"error": "unsupported arch"}
iat = build_iat_map(pe)
func_starts = {f.vaddr for f in functions}
targets = functions if cap_functions is None else functions[:cap_functions]
cfgs = {}
callgraph = defaultdict(set)
icall_edges = defaultdict(list)
failed = 0
for fn in targets:
if fn.size is None:
continue
off = _rva_to_offset(pe, fn.rva)
if off is None:
continue
code = bytes(pe.__data__[off:off + fn.size])
cfg, ok = build_cfg_for_function(md, code, fn.vaddr, func_starts, iat)
if not ok:
failed += 1
continue
cfgs[fn.vaddr] = cfg
for t in cfg.calls:
if t in func_starts:
callgraph[fn.vaddr].add(t)
icall_edges[fn.vaddr] = cfg.icalls
stats = {
"functions_analyzed": len(cfgs),
"functions_failed": failed,
"total_blocks": sum(c.block_count for c in cfgs.values()),
"total_edges": sum(c.edge_count for c in cfgs.values()),
"direct_call_edges": sum(len(v) for v in callgraph.values()),
"indirect_call_sites": sum(len(v) for v in icall_edges.values()),
}
return cfgs, {"direct": callgraph, "indirect": icall_edges}, stats
def build_reverse_callgraph(callgraph, callback_edges=None):
rev = defaultdict(dict)
for caller, callees in callgraph.get("direct", {}).items():
for callee in callees:
rev[callee][caller] = "direct"
if callback_edges:
for caller, callees in callback_edges.items():
for callee in callees:
rev[callee].setdefault(caller, "callback")
return rev
ATTACK_SURFACE_PATTERNS = [
"dispatch", "ioctl", "deviceio", "irp", "recv", "receive", "inbound",
"parse", "decode", "unmarshal", "deserialize", "fromwire", "read",
"callout", "classify", "callback", "handler", "onpacket", "wfp",
"ndis", "netbuffer", "indicate", "complete", "worker", "process",
"input", "request", "query", "setinformation",
]
TRUSTED_HINT_PATTERNS = [
"init", "alloc", "free", "cleanup", "destroy", "release", "reference",
"lookup", "hash", "insert", "remove", "acquire", "lock", "unlock",
]
def address_taken_edges(pe, functions):
func_starts = {f.vaddr for f in functions}
edges = defaultdict(set)
md = _md_for(pe, detail=True)
if md is None:
return edges
data = bytes(pe.__data__)
for fn in functions:
if fn.size is None:
continue
off = _rva_to_offset(pe, fn.rva)
if off is None:
continue
code = data[off:off + fn.size]
for ins in md.disasm(code, fn.vaddr):
if ins.mnemonic == "lea" and ins.operands and len(ins.operands) == 2:
op = ins.operands[1]
if op.type == X86_OP_MEM and op.mem.base == X86_REG_RIP:
tgt = ins.address + ins.size + op.mem.disp
if tgt in func_starts and tgt != fn.vaddr:
edges[fn.vaddr].add(tgt)
return edges
def _address_taken_functions(pe, functions):
edges = address_taken_edges(pe, functions)
taken = set()
for targets in edges.values():
taken |= targets
return taken
def _looks_like_symbol(text):
if not (6 <= len(text) <= 96):
return False
if not (text[0].isalpha() or text[0] == "_"):
return False
ok = sum(1 for c in text if c.isalnum() or c in "_")
if ok / len(text) < 0.95:
return False
has_upper = any(c.isupper() for c in text)
has_lower = any(c.islower() for c in text)
return has_upper and has_lower and " " not in text
def build_name_hints(pe, functions, strings):
md = _md_for(pe, detail=True)
if md is None:
return {}
import_names = set()
if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
for entry in pe.DIRECTORY_ENTRY_IMPORT:
for imp in entry.imports:
if imp.name:
import_names.add(imp.name.decode("latin-1"))
sym_by_va = {}
for s in strings:
if s["encoding"] == "ascii" and s["vaddr"] and _looks_like_symbol(s["text"]):
sym_by_va[s["vaddr"]] = s["text"]
if not sym_by_va:
return {}
data = bytes(pe.__data__)
counts = defaultdict(lambda: defaultdict(int))
name_refcount = defaultdict(set)
for fn in functions:
if fn.size is None:
continue
off = _rva_to_offset(pe, fn.rva)
if off is None:
continue
code = data[off:off + fn.size]
for ins in md.disasm(code, fn.vaddr):
if ins.mnemonic == "lea" and ins.operands and len(ins.operands) == 2:
op = ins.operands[1]
if op.type == X86_OP_MEM and op.mem.base == X86_REG_RIP:
tgt = ins.address + ins.size + op.mem.disp
if tgt in sym_by_va:
nm = sym_by_va[tgt]
counts[fn.vaddr][nm] += 1
name_refcount[nm].add(fn.vaddr)
hints = {}
for va, namecounts in counts.items():
candidates = [(nm, c) for nm, c in namecounts.items()
if nm not in import_names and len(name_refcount[nm]) <= 3]
if not candidates:
continue
best = max(candidates, key=lambda kv: (kv[1], -len(name_refcount[kv[0]])))
hints[va] = best[0]
log(f"[reach] name hints: {len(hints)} functions named "
f"(excluded {len(import_names)} import names + shared-string refs)", 1)
return hints
def identify_attack_surface(pe, functions, callgraph, strings, name_hint=None):
base = pe.OPTIONAL_HEADER.ImageBase
func_starts = sorted(f.vaddr for f in functions)
surface = {}
ep = base + pe.OPTIONAL_HEADER.AddressOfEntryPoint
if ep in set(func_starts):
surface[ep] = "entry-point"
exports = collect_exports(pe)
for rva in exports:
va = base + rva
if va in set(func_starts):
surface.setdefault(va, "export")
name_of = name_hint or {}
for va, nm in name_of.items():
low = nm.lower()
if any(p in low for p in ATTACK_SURFACE_PATTERNS):
surface[va] = f"named:{nm}"
taken = _address_taken_functions(pe, functions)
for va in taken:
surface.setdefault(va, "address-taken(callback)")
log(f"[reach] attack surface: {len(surface)} entry functions "
f"({sum(1 for v in surface.values() if v.startswith('named'))} named, "
f"{sum(1 for v in surface.values() if 'callback' in v)} callbacks)", 1)
return surface
def compute_reachability(finding_funcs, surface, rev_callgraph, max_hops=40):
surface_set = set(surface)
result = {}
for target in finding_funcs:
if target in surface_set:
result[target] = {"reachable": True, "via": [surface[target]],
"entries": [target], "hops": 0, "used_callback": False}
continue
seen = {target}
frontier = {target: False}
entries = {}
hops = 0
found_hop = None
while frontier and hops < max_hops:
hops += 1
nxt = {}
for node, via_cb in frontier.items():
for caller, etype in rev_callgraph.get(node, {}).items():
path_cb = via_cb or (etype == "callback")
if caller in surface_set:
if caller not in entries or (entries[caller] and not path_cb):
entries[caller] = path_cb
if found_hop is None:
found_hop = hops
if caller not in seen:
seen.add(caller)
nxt[caller] = path_cb
frontier = nxt
if entries and hops >= (found_hop or 0) + 2:
break
result[target] = {
"reachable": bool(entries),
"entries": sorted(entries)[:8],
"via": sorted({surface[e] for e in entries})[:6],
"hops": found_hop,
"used_callback": bool(entries) and all(entries.values()),
}
return result
def build_arg_flow_edges(pe, functions, iat, cfgs=None, trace=False):
eng = InterpEngine(pe, functions, iat, cfgs=cfgs, max_depth=1)
edges = defaultdict(lambda: defaultdict(set))
def on_call(site, target, iname, args, machine, stack):
if target is None or target not in eng.func_starts:
return
caller = stack[0] if stack else None
if caller is None:
return
for i, r in enumerate(X64_ARG_REGS):
v = args[r]
srcs = v.input_sources() if isinstance(v, V) else set()
for s in srcs:
if s.startswith("arg"):
try:
caller_argi = int(s[3:])
except ValueError:
continue
edges[caller][target].add((caller_argi, i))
import time
t0 = time.time()
n = 0
for fn in functions:
if fn.size is None:
continue
try:
eng.run_function(fn.vaddr, on_call=on_call)
except Exception:
pass
n += 1
if trace and n % 1000 == 0:
log(f"[flow] {n} functions processed for arg-flow edges", 2)
total = sum(len(v) for cm in edges.values() for v in cm.values())
log(f"[flow] {total} caller-arg -> callee-arg flow edges over "
f"{len(edges)} callers", 1)
return edges
def compute_input_flow(finding_targets, surface, callgraph, arg_flow, max_hops=30):
surface_set = set(surface)
tainted = defaultdict(set)
for e in surface_set:
tainted[e] = {0, 1, 2, 3}
direct = callgraph.get("direct", {})
changed = True
passes = 0
while changed and passes < max_hops:
changed = False
passes += 1
for caller, callees in direct.items():
ct = tainted.get(caller)
if not ct:
continue
flowmap = arg_flow.get(caller, {})
for callee in callees:
pairs = flowmap.get(callee)
if not pairs:
continue
for (caller_argi, callee_argi) in pairs:
if caller_argi in ct and callee_argi not in tainted[callee]:
tainted[callee].add(callee_argi)
changed = True
result = {}
for t in finding_targets:
result[t] = {"tainted_args": sorted(tainted.get(t, set())),
"any": bool(tainted.get(t))}
return result
def build_callsite_contracts(pe, functions, iat, cfgs=None, trace=False,
time_budget=180.0):
eng = InterpEngine(pe, functions, iat, cfgs=cfgs, max_depth=1)
contracts = defaultdict(lambda: defaultdict(list))
def on_any_call(site, target, iname, args, state, stack):
caller = stack[0] if stack else None
if caller is None:
return
conds = state.conds
for i, r in enumerate(X64_ARG_REGS):
v = args[r]
if not isinstance(v, V) or v.kind in ("const", "local"):
continue
if not v.contains_input():
continue
ctrl = assess_control(v, conds, role="size")
passed_argis = sorted(int(s[3:]) for s in v.input_sources()
if s.startswith("arg") and s[3:].isdigit())
key = (target, i)
existing = contracts[caller][key]
if any(e["site"] == site for e in existing):
continue
contracts[caller][key].append({
"site": site,
"caller_bounds": ctrl["upper_bounded"],
"status": ctrl["status"],
"from_args": passed_argis,
"carries_input": v.contains_input(),
"expr": repr(v),
})
import time
explorer = PathExplorer(eng, on_any_call=on_any_call, max_paths=32,
per_func_secs=3.0, clock=time.time)
total_fns = sum(1 for fn in functions if fn.size is not None)
log(f"[collide] building call-site contracts over {total_fns} functions "
f"(budget {time_budget}s)", 1)
t0 = time.time()
last = t0
n = 0
for fn in functions:
if fn.size is None:
continue
explorer.paths_done = 0
try:
explorer.explore(fn.vaddr)
except Exception:
pass
n += 1
now = time.time()
if trace and (n % 200 == 0 or now - last >= 3.0):
nc = sum(len(v) for cm in contracts.values() for v in cm.values())
log(f"[collide] {n}/{total_fns} fns, {nc} contracts, {now - t0:.0f}s", 2)
last = now
if now - t0 > time_budget:
log(f"[collide] time budget {time_budget}s reached at fn {n}", 2)
break
total = sum(len(v) for cm in contracts.values() for v in cm.values())
log(f"[collide] {total} call-site arg contracts over "
f"{len(contracts)} callers", 1)
return contracts
def detect_collisions(findings, contracts, owner_func, rev_callgraph, name_hints,
surface=None, trace=False):
callee_unmet = {}
skipped_norole = skipped_bounded = skipped_noargs = 0
for f in findings:
if f.arg_role not in ("size", "len"):
skipped_norole += 1
continue
cstat = f.verdict.get("control", {}).get("status")
if cstat not in ("unguarded", "guarded-not-upper"):
skipped_bounded += 1
continue
srcs = f.verdict.get("sources", [])
argis = sorted(int(s[3:]) for s in srcs if s.startswith("arg") and s[3:].isdigit())
if not argis:
skipped_noargs += 1
continue
fnva = owner_func(f.site)
if fnva is None:
continue
for ai in argis:
callee_unmet.setdefault((fnva, ai), []).append(f)
log(f"[collide] unmet-assumption callee slots: {len(callee_unmet)} "
f"(skipped: {skipped_norole} non-size, {skipped_bounded} bounded, "
f"{skipped_noargs} no-arg-source)", 1)
forward = defaultdict(list)
for caller, cm in contracts.items():
for (callee, callee_argi), entries in cm.items():
for e in entries:
if e["caller_bounds"]:
continue
for src_argi in e["from_args"]:
forward[(callee, callee_argi)].append({
"caller": caller,
"caller_argi": src_argi,
"site": e["site"],
"status": e["status"],
"carries_input": e.get("carries_input", False),
"value": e["expr"],
})
matched_callers = 0
slots_with_callers = 0
collisions = []
for (callee, callee_argi), fs in callee_unmet.items():
callers = rev_callgraph.get(callee, {})
if callers:
slots_with_callers += 1
for caller in callers:
entries = contracts.get(caller, {}).get((callee, callee_argi))
if not entries:
continue
matched_callers += 1
for e in entries:
if e["caller_bounds"]:
continue
chain = _extend_chain(caller, e, forward, name_hints, surface)
collisions.append({
"callee": hex(callee),
"callee_name": name_hints.get(callee),
"callee_argi": callee_argi,
"caller": hex(caller),
"caller_name": name_hints.get(caller),
"call_site": hex(e["site"]),
"caller_status": e["status"],
"passed_from": e["from_args"],
"carries_input": e.get("carries_input", False),
"value": e["expr"],
"sink_kind": fs[0].sink_kind,
"sink": fs[0].sink,
"sink_site": hex(fs[0].site),
"hops": len(chain),
"chain": chain,
"reaches_entry": bool(chain and chain[-1]["is_entry"]),
})
log(f"[collide] {slots_with_callers} unmet slots have callers; "
f"{matched_callers} caller/slot pairs matched a call-site contract", 1)
seen = set()
uniq = []
for c in collisions:
k = (c["callee"], c["callee_argi"], c["caller"], c["call_site"])
if k in seen:
continue
seen.add(k)
uniq.append(c)
uniq.sort(key=lambda c: (0 if c.get("carries_input") else 1,
0 if c.get("reaches_entry") else 1,
-c.get("hops", 0),
0 if c["caller_status"] == "unguarded" else 1))
multi = sum(1 for c in uniq if c["hops"] > 1)
log(f"[collide] {len(uniq)} collisions ({multi} multi-hop chains)", 1)
return uniq
def _extend_chain(start_caller, start_edge, forward, name_hints, surface,
max_hops=8):
surface = surface or {}
chain = [{
"func": hex(start_caller),
"name": name_hints.get(start_caller),
"site": hex(start_edge["site"]),
"status": start_edge["status"],
"value": start_edge.get("value") or start_edge.get("expr", "?"),
"is_entry": start_caller in surface,
}]
cur = start_caller
seen = {cur}
for _ in range(max_hops):
if cur in surface:
break
best = None
for ai in range(4):
cand = forward.get((cur, ai))
if cand:
best = cand[0]
break
if not best:
break
nxt = best["caller"]
if nxt in seen:
break
seen.add(nxt)
chain.append({
"func": hex(nxt),
"name": name_hints.get(nxt),
"site": hex(best["site"]),
"status": best["status"],
"value": best["value"],
"is_entry": nxt in surface,
})
cur = nxt
return chain
X64_GP = {
csx86.X86_REG_RAX: "rax", csx86.X86_REG_RBX: "rbx", csx86.X86_REG_RCX: "rcx",
csx86.X86_REG_RDX: "rdx", csx86.X86_REG_RSI: "rsi", csx86.X86_REG_RDI: "rdi",
csx86.X86_REG_RBP: "rbp", csx86.X86_REG_RSP: "rsp",
csx86.X86_REG_R8: "r8", csx86.X86_REG_R9: "r9", csx86.X86_REG_R10: "r10",
csx86.X86_REG_R11: "r11", csx86.X86_REG_R12: "r12", csx86.X86_REG_R13: "r13",
csx86.X86_REG_R14: "r14", csx86.X86_REG_R15: "r15",
}
_SUBREG = {
csx86.X86_REG_EAX: "rax", csx86.X86_REG_AX: "rax", csx86.X86_REG_AL: "rax", csx86.X86_REG_AH: "rax",
csx86.X86_REG_EBX: "rbx", csx86.X86_REG_BX: "rbx", csx86.X86_REG_BL: "rbx", csx86.X86_REG_BH: "rbx",
csx86.X86_REG_ECX: "rcx", csx86.X86_REG_CX: "rcx", csx86.X86_REG_CL: "rcx", csx86.X86_REG_CH: "rcx",
csx86.X86_REG_EDX: "rdx", csx86.X86_REG_DX: "rdx", csx86.X86_REG_DL: "rdx", csx86.X86_REG_DH: "rdx",
csx86.X86_REG_ESI: "rsi", csx86.X86_REG_SI: "rsi", csx86.X86_REG_SIL: "rsi",
csx86.X86_REG_EDI: "rdi", csx86.X86_REG_DI: "rdi", csx86.X86_REG_DIL: "rdi",
csx86.X86_REG_EBP: "rbp", csx86.X86_REG_BP: "rbp",
csx86.X86_REG_ESP: "rsp", csx86.X86_REG_SP: "rsp",
csx86.X86_REG_R8D: "r8", csx86.X86_REG_R8W: "r8", csx86.X86_REG_R8B: "r8",
csx86.X86_REG_R9D: "r9", csx86.X86_REG_R9W: "r9", csx86.X86_REG_R9B: "r9",
csx86.X86_REG_R10D: "r10", csx86.X86_REG_R10W: "r10", csx86.X86_REG_R10B: "r10",
csx86.X86_REG_R11D: "r11", csx86.X86_REG_R11W: "r11", csx86.X86_REG_R11B: "r11",
csx86.X86_REG_R12D: "r12", csx86.X86_REG_R12W: "r12", csx86.X86_REG_R12B: "r12",
csx86.X86_REG_R13D: "r13", csx86.X86_REG_R13W: "r13", csx86.X86_REG_R13B: "r13",
csx86.X86_REG_R14D: "r14", csx86.X86_REG_R14W: "r14", csx86.X86_REG_R14B: "r14",
csx86.X86_REG_R15D: "r15", csx86.X86_REG_R15W: "r15", csx86.X86_REG_R15B: "r15",
}
def reg_name(reg):
if reg in X64_GP:
return X64_GP[reg]
return _SUBREG.get(reg)
X64_ARG_REGS = ["rcx", "rdx", "r8", "r9"]
class V:
__slots__ = ("kind", "a", "b", "op", "size", "tag")
def __init__(self, kind, a=None, b=None, op=None, size=None, tag=None):
self.kind = kind
self.a = a
self.b = b
self.op = op
self.size = size
self.tag = tag
def __repr__(self):
k = self.kind
if k == "const":
return hex(self.a) if isinstance(self.a, int) else str(self.a)
if k == "input":
return f"in:{self.tag}"
if k == "local":
return f"loc:{self.tag}"
if k == "field":
return f"{self.a!r}.[{hex(self.b) if isinstance(self.b, int) else self.b}]"
if k == "load":
return f"*({self.a!r})"
if k == "binop":
return f"({self.a!r} {self.op} {self.b!r})"
if k == "call_ret":
return f"ret:{self.tag}"
if k == "top":
return "?"
return f"V({k})"
def key(self):
k = self.kind
if k == "const":
return ("const", self.a)
if k == "input":
return ("input", self.tag)
if k == "local":
return ("local", self.tag)
if k == "call_ret":
return ("call_ret", self.tag)
if k == "top":
return ("top", id(self))
if k == "load":
return ("load", self.a.key())
if k == "field":
return ("field", self.a.key(), self.b)
if k == "binop":
return ("binop", self.op, self.a.key(), self.b.key())
return ("v", id(self))
def contains_input(self):
if self.kind in ("input",):
return True
if self.kind == "call_ret":
return False
for child in (self.a, self.b):
if isinstance(child, V) and child.contains_input():
return True
return False
def input_sources(self):
out = set()
stack = [self]
while stack:
v = stack.pop()
if not isinstance(v, V):
continue
if v.kind == "input":
out.add(v.tag)
for c in (v.a, v.b):
if isinstance(c, V):
stack.append(c)
return out
def const(n):
return V("const", a=n)
def top():
return V("top")
def inp(tag):
return V("input", tag=tag)
def local(tag):
return V("local", tag=tag)
def _split_base_off(addr):
if addr.kind == "binop" and addr.op in ("+", "-") and addr.b.kind == "const" \
and isinstance(addr.b.a, int):
off = addr.b.a if addr.op == "+" else -addr.b.a
return addr.a, off
if addr.kind in ("input", "field", "load", "local"):
return addr, 0
return None, 0
def mk_binop(op, a, b):
if a.kind == "const" and b.kind == "const" and isinstance(a.a, int) and isinstance(b.a, int):
if op == "+":
return const(a.a + b.a)
if op == "-":
return const(a.a - b.a)
if op == "*":
return const(a.a * b.a)
if op == "&":
return const(a.a & b.a)
if op == "^":
return const(a.a ^ b.a)
if op == "|":
return const(a.a | b.a)
if op == "<<":
return const(a.a << b.a)
if op == "^" and a.key() == b.key():
return const(0)
if op in ("+", "-") and b.kind == "const" and b.a == 0:
return a
if op == "*" and b.kind == "const" and b.a == 1:
return a
return V("binop", a=a, b=b, op=op)
def mk_load(addr, size=None):
return V("load", a=addr, size=size)
def mk_field(base, off):
return V("field", a=base, b=off)
class Machine:
def __init__(self, image_base=0, arg_regs=X64_ARG_REGS):
self.regs = {}
self.mem = {}
self.image_base = image_base
self.arg_regs = arg_regs
self.flags_cmp = None
def clone(self):
m = Machine(self.image_base, self.arg_regs)
m.regs = dict(self.regs)
m.mem = dict(self.mem)
m.flags_cmp = self.flags_cmp
return m
def seed_params(self):
for i, r in enumerate(self.arg_regs):
self.regs[r] = inp(f"arg{i}")
def get_reg(self, name):
if name is None:
return top()
if name in self.regs:
return self.regs[name]
return local(name)
def set_reg(self, name, val):
if name is None:
return
self.regs[name] = val
def _mem_key(self, addr):
return addr.key()
def load_mem(self, addr, size=None):
k = self._mem_key(addr)
if k in self.mem:
return self.mem[k]
base, off = _split_base_off(addr)
if base is not None and base.kind in ("input", "field", "load"):
return mk_field(base, off)
return mk_load(addr, size)
def store_mem(self, addr, val):
self.mem[self._mem_key(addr)] = val
def eval_mem_addr(self, op):
m = op.mem
parts = []
if m.base != 0:
if m.base == X86_REG_RIP:
return const(self.image_base + 0)
bn = reg_name(m.base)
parts.append(self.get_reg(bn))
if m.index != 0:
idx = self.get_reg(reg_name(m.index))
scale = const(m.scale if m.scale else 1)
parts.append(mk_binop("*", idx, scale))
acc = None
for p in parts:
acc = p if acc is None else mk_binop("+", acc, p)
if m.disp:
acc = const(m.disp) if acc is None else mk_binop("+", acc, const(m.disp))
if acc is None:
acc = const(0)
return acc
def eval_operand(self, ins, op):
if op.type == X86_OP_REG:
return self.get_reg(reg_name(op.reg))
if op.type == X86_OP_IMM:
return const(op.imm)
if op.type == X86_OP_MEM:
if op.mem.base == X86_REG_RIP:
addr = ins.address + ins.size + op.mem.disp
return mk_load(const(addr), op.size)
addr = self.eval_mem_addr(op)
return self.load_mem(addr, op.size)
return top()
def write_operand(self, ins, op, val):
if op.type == X86_OP_REG:
self.set_reg(reg_name(op.reg), val)
elif op.type == X86_OP_MEM:
if op.mem.base == X86_REG_RIP:
addr = const(ins.address + ins.size + op.mem.disp)
else:
addr = self.eval_mem_addr(op)
self.store_mem(addr, val)
_BINOP_MNEM = {
"add": "+", "sub": "-", "imul": "*", "and": "&",
"or": "|", "xor": "^", "shl": "<<", "sar": ">>", "shr": ">>",
}
def step(machine, ins):
m = ins.mnemonic
ops = ins.operands
if m == "mov" or m == "movzx" or m == "movsx" or m == "movsxd":
if len(ops) == 2:
machine.write_operand(ins, ops[0], machine.eval_operand(ins, ops[1]))
return
if m == "lea":
if len(ops) == 2 and ops[1].type == X86_OP_MEM:
if ops[1].mem.base == X86_REG_RIP:
addr = const(ins.address + ins.size + ops[1].mem.disp)
else:
addr = machine.eval_mem_addr(ops[1])
machine.write_operand(ins, ops[0], addr)
return
if m == "push":
return
if m == "pop":
if ops:
machine.write_operand(ins, ops[0], top())
return
if m in _BINOP_MNEM and len(ops) == 2:
dst = machine.eval_operand(ins, ops[0])
src = machine.eval_operand(ins, ops[1])
machine.write_operand(ins, ops[0], mk_binop(_BINOP_MNEM[m], dst, src))
return
if m == "cmp" and len(ops) == 2:
machine.flags_cmp = (machine.eval_operand(ins, ops[0]),
machine.eval_operand(ins, ops[1]))
return
if m == "test" and len(ops) == 2:
machine.flags_cmp = (machine.eval_operand(ins, ops[0]),
machine.eval_operand(ins, ops[1]))
return
if m == "xor" and len(ops) == 2 and ops[0].type == X86_OP_REG and ops[1].type == X86_OP_REG \
and reg_name(ops[0].reg) == reg_name(ops[1].reg):
machine.write_operand(ins, ops[0], const(0))
return
if m in ("inc", "dec") and ops:
v = machine.eval_operand(ins, ops[0])
machine.write_operand(ins, ops[0], mk_binop("+" if m == "inc" else "-", v, const(1)))
return
if m in ("call",):
return
for op in ops:
if op.type in (X86_OP_REG,):
machine.set_reg(reg_name(op.reg), top())
X64_VOLATILE = ["rax", "rcx", "rdx", "r8", "r9", "r10", "r11"]
class InterpEngine:
def __init__(self, pe, functions, iat, cfgs=None, max_depth=8, max_insns=20000):
self.pe = pe
self.image_base = pe.OPTIONAL_HEADER.ImageBase
self.iat = iat
self.md = _md_for(pe, detail=True)
self.in_exec, self.read, _ = _make_va_reader(pe)
self.func_starts = {f.vaddr for f in functions}
self.func_size = {f.vaddr: f.size for f in functions}
self.cfgs = cfgs or {}
self.max_depth = max_depth
self.max_insns = max_insns
self._insns_cache = {}
self.deadline = None
self.clock = None
self.pulse = None
def expired(self):
if self.deadline is None or self.clock is None:
return False
return self.clock() > self.deadline
def _decode(self, va):
if va in self._insns_cache:
return self._insns_cache[va]
size = self.func_size.get(va)
if size is None:
code = self.read(va, 1024)
else:
code = self.read(va, size)
insns = list(self.md.disasm(code, va)) if code else []
idx = {i.address: i for i in insns}
self._insns_cache[va] = (insns, idx)
return insns, idx
def run_function(self, func_va, machine=None, depth=0, stack=None, budget=None, on_call=None):
stack = stack or ()
if func_va in stack or depth > self.max_depth:
return top()
stack = stack + (func_va,)
if budget is None:
budget = [self.max_insns]
insns, idx = self._decode(func_va)
if not insns:
return top()
if machine is None:
machine = Machine(image_base=self.image_base)
machine.seed_params()
succ = self._linear_succ(insns, idx)
pc = func_va
visited = set()
ret_val = top()
while pc is not None and pc in idx:
if budget[0] <= 0:
break
if pc in visited:
break
if (budget[0] & 0xFF) == 0:
if self.pulse is not None:
self.pulse.tick(func_va, -1, -1, depth)
if self.expired():
break
visited.add(pc)
ins = idx[pc]
budget[0] -= 1
g = ins.groups
if CS_GRP_CALL in g:
target = _direct_target(ins)
callee_ret = self._handle_call(ins, machine, target, depth, stack, budget, on_call)
for r in X64_VOLATILE:
machine.regs[r] = top()
machine.regs["rax"] = callee_ret
pc = succ.get(pc)
continue
if CS_GRP_RET in g:
ret_val = machine.get_reg("rax")
break
if CS_GRP_JUMP in g:
t = _direct_target(ins)
if ins.mnemonic == "jmp":
if t in self.func_starts and t not in idx:
ret_val = self.run_function(t, machine, depth + 1, stack, budget, on_call)
break
pc = t if t in idx else succ.get(pc)
continue
pc = succ.get(pc)
continue
step(machine, ins)
pc = succ.get(pc)
return ret_val
def _handle_call(self, ins, machine, target, depth, stack, budget, on_call):
if on_call is not None:
args = {r: machine.get_reg(r) for r in X64_ARG_REGS}
slot = _rip_mem_target(ins)
iname = self.iat.get(slot) if slot is not None else None
on_call(ins.address, target, iname, args, machine, stack)
if target is not None and target in self.func_starts:
callee = Machine(image_base=self.image_base)
for r in X64_ARG_REGS:
callee.regs[r] = machine.get_reg(r)
callee.mem = machine.mem
return self.run_function(target, callee, depth + 1, stack, budget, on_call)
return top()
@staticmethod
def _linear_succ(insns, idx):
succ = {}
for i, ins in enumerate(insns):
nxt = insns[i + 1].address if i + 1 < len(insns) else None
succ[ins.address] = nxt
return succ
_JCC_TRUE = {
"je": "==", "jz": "==", "jne": "!=", "jnz": "!=",
"jb": "u<", "jnae": "u<", "jc": "u<",
"jae": "u>=", "jnb": "u>=", "jnc": "u>=",
"jbe": "u<=", "jna": "u<=", "ja": "u>", "jnbe": "u>",
"jl": "s<", "jnge": "s<", "jge": "s>=", "jnl": "s>=",
"jle": "s<=", "jng": "s<=", "jg": "s>", "jnle": "s>",
"js": "sign", "jns": "nsign",
}
_NEG_REL = {
"==": "!=", "!=": "==",
"u<": "u>=", "u>=": "u<", "u<=": "u>", "u>": "u<=",
"s<": "s>=", "s>=": "s<", "s<=": "s>", "s>": "s<=",
"sign": "nsign", "nsign": "sign",
}
@dataclass
class Constraint:
rel: str
a: object
b: object
site: int
def __repr__(self):
return f"[{self.a!r} {self.rel} {self.b!r}]"
def vars(self):
out = set()
for v in (self.a, self.b):
if isinstance(v, V):
out |= {k for k in _all_subkeys(v)}
return out
def _all_subkeys(v):
out = set()
stack = [v]
while stack:
node = stack.pop()
if not isinstance(node, V):
continue
out.add(node.key())
for c in (node.a, node.b):
if isinstance(c, V):
stack.append(c)
return out
class PathState:
__slots__ = ("machine", "conds", "block_visits", "insns_run")
def __init__(self, machine, conds=None, block_visits=None, insns_run=0):
self.machine = machine
self.conds = conds if conds is not None else []
self.block_visits = block_visits if block_visits is not None else {}
self.insns_run = insns_run
def fork(self):
return PathState(self.machine.clone(), list(self.conds),
dict(self.block_visits), self.insns_run)
class Pulse:
def __init__(self, clock, interval=2.0):
self.clock = clock
self.interval = interval
self.last = clock()
def tick(self, fn_va, paths, bdepth, call_depth):
now = self.clock()
if now - self.last >= self.interval:
self.last = now
loc = hex(fn_va) if fn_va is not None else "?"
print(f" ... still working in {loc}: {paths} paths, "
f"branch-depth {bdepth}, call-depth {call_depth}", flush=True)
class PathExplorer:
def __init__(self, engine, on_sink=None, max_paths=48, max_block_visits=2,
max_insns=6000, max_depth=6, max_sink_obs=400,
per_func_secs=None, clock=None, max_branch_depth=120,
on_any_call=None):
self.eng = engine
self.on_sink = on_sink
self.on_any_call = on_any_call
self.max_paths = max_paths
self.max_block_visits = max_block_visits
self.max_insns = max_insns
self.max_depth = max_depth
self.max_sink_obs = max_sink_obs
self.per_func_secs = per_func_secs
self.clock = clock
self.max_branch_depth = max_branch_depth
self.paths_done = 0
self.timed_out = False
self._deadline = None
self._sink_seen = set()
self._sink_obs = 0
self.pulse = None
self._cur_fn = None
def explore(self, func_va):
insns, idx = self.eng._decode(func_va)
if not insns:
return
self._sink_seen = set()
self._sink_obs = 0
self.timed_out = False
self._cur_fn = func_va
if self.per_func_secs and self.clock:
self._deadline = self.clock() + self.per_func_secs
self.eng.deadline = self._deadline
self.eng.clock = self.clock
else:
self._deadline = None
self.eng.deadline = None
m = Machine(image_base=self.eng.image_base)
m.seed_params()
start = PathState(m)
self._succ = self.eng._linear_succ(insns, idx)
self._steps = 0
self._walk(func_va, func_va, start, idx, (func_va,), 0, 0)
def _expired(self):
if self._deadline is None:
return False
if self.clock() > self._deadline:
self.timed_out = True
return True
return False
def _heartbeat(self, bdepth, depth):
self._steps += 1
if (self._steps & 0xFF) == 0:
if self.pulse is not None:
self.pulse.tick(self._cur_fn, self.paths_done, bdepth, depth)
if self._expired():
return True
return False
def _walk(self, func_va, pc, state, idx, stack, depth, bdepth):
succ = self._succ
while pc is not None and pc in idx:
if self.paths_done >= self.max_paths or self.timed_out:
return
if self._heartbeat(bdepth, depth):
return
if state.insns_run >= self.max_insns or bdepth > self.max_branch_depth:
return
ins = idx[pc]
state.insns_run += 1
g = ins.groups
if CS_GRP_CALL in g:
self._do_call(ins, state, depth, stack)
pc = succ.get(pc)
continue
if CS_GRP_RET in g:
self.paths_done += 1
return
if CS_GRP_JUMP in g:
mnem = ins.mnemonic
t = _direct_target(ins)
if mnem == "jmp":
if t in self.eng.func_starts and t not in idx:
self.paths_done += 1
return
nb = t if t in idx else succ.get(pc)
pc = self._advance_block(state, nb)
if pc is None:
return
continue
rel = _JCC_TRUE.get(mnem)
cmpv = state.machine.flags_cmp
taken_pc = t if t in idx else None
fall_pc = succ.get(pc)
self._branch(func_va, state, idx, stack, depth,
rel, cmpv, ins.address, taken_pc, fall_pc, bdepth)
return
step(state.machine, ins)
pc = succ.get(pc)
return
def _branch(self, func_va, state, idx, stack, depth, rel, cmpv, site, taken_pc, fall_pc, bdepth):
branches = []
if taken_pc is not None:
branches.append((taken_pc, rel, False))
if fall_pc is not None:
branches.append((fall_pc, rel, True))
for i, (npc, r, negate) in enumerate(branches):
if self.paths_done >= self.max_paths or self.timed_out:
return
child = state if i == len(branches) - 1 else state.fork()
if cmpv is not None and r is not None:
a, b = cmpv
use_rel = _NEG_REL.get(r, r) if negate else r
child.conds.append(Constraint(use_rel, a, b, site))
start_pc = self._advance_block(child, npc)
if start_pc is not None:
self._walk(func_va, start_pc, child, idx, stack, depth, bdepth + 1)
def _advance_block(self, state, npc):
if npc is None:
return None
c = state.block_visits.get(npc, 0)
if c >= self.max_block_visits:
return None
state.block_visits[npc] = c + 1
return npc
def _emit_sink(self, site, target, iname, args, state, stack):
if self.on_sink is None or _sink_name(iname) is None:
return
if self._sink_obs >= self.max_sink_obs:
return
argkey = tuple(args[r].key() for r in X64_ARG_REGS)
condkey = frozenset(c.rel + str(c.a.key()) for c in state.conds
if isinstance(c.a, V))
k = (site, iname, argkey, condkey)
if k in self._sink_seen:
return
self._sink_seen.add(k)
self._sink_obs += 1
self.on_sink(site, target, iname, args, state, stack)
def _do_call(self, ins, state, depth, stack):
m = state.machine
target = _direct_target(ins)
slot = _rip_mem_target(ins)
iname = self.eng.iat.get(slot) if slot is not None else None
args = {r: m.get_reg(r) for r in X64_ARG_REGS}
self._emit_sink(ins.address, target, iname, args, state, stack)
if self.on_any_call is not None and target is not None \
and target in self.eng.func_starts:
self.on_any_call(ins.address, target, iname, args, state, stack)
ret = top()
if target is not None and target in self.eng.func_starts and depth < self.max_depth \
and target not in stack:
ret = self._call_summary(target, m, depth, stack, state)
for r in X64_VOLATILE:
m.regs[r] = top()
m.regs["rax"] = ret
def _call_summary(self, target, caller_m, depth, stack, state):
callee = Machine(image_base=self.eng.image_base)
for r in X64_ARG_REGS:
callee.regs[r] = caller_m.get_reg(r)
callee.mem = caller_m.mem
budget = [self.eng.max_insns]
return self.eng.run_function(target, callee, depth + 1, stack, budget,
on_call=self._nested_sink(state))
def _nested_sink(self, state):
def cb(site, target, iname, args, machine, cstack):
self._emit_sink(site, target, iname, args, state, cstack)
return cb
SINKS = {
"ExAllocatePool2": {"kind": "alloc", "args": {2: "size"}},
"ExAllocatePool3": {"kind": "alloc", "args": {2: "size"}},
"ExAllocatePoolWithTag": {"kind": "alloc", "args": {1: "size"}},
"ExAllocatePoolWithQuotaTag": {"kind": "alloc", "args": {1: "size"}},
"MmAllocateContiguousMemory": {"kind": "alloc", "args": {0: "size"}},
"RtlCopyMemory": {"kind": "copy", "args": {0: "dest", 1: "src", 2: "len"}},
"memcpy": {"kind": "copy", "args": {0: "dest", 1: "src", 2: "len"}},
"memmove": {"kind": "copy", "args": {0: "dest", 1: "src", 2: "len"}},
"RtlMoveMemory": {"kind": "copy", "args": {0: "dest", 1: "src", 2: "len"}},
"memset": {"kind": "fill", "args": {0: "dest", 2: "len"}},
"RtlZeroMemory": {"kind": "fill", "args": {0: "dest", 1: "len"}},
"RtlFillMemory": {"kind": "fill", "args": {0: "dest", 1: "len"}},
"strcpy_s": {"kind": "copy", "args": {0: "dest", 1: "len", 2: "src"}},
"wcscpy_s": {"kind": "copy", "args": {0: "dest", 1: "len", 2: "src"}},
"ProbeForRead": {"kind": "probe", "args": {0: "ptr", 1: "len"}},
"ProbeForWrite": {"kind": "probe", "args": {0: "ptr", 1: "len"}},
"MmMapLockedPagesSpecifyCache": {"kind": "map", "args": {0: "mdl"}},
"NdisGetDataBuffer": {"kind": "parse", "args": {1: "len"}},
}
def _sink_name(iname):
if not iname:
return None
base = iname.split("!", 1)[-1]
return base if base in SINKS else None
def classify_value(v):
if v is None:
return {"class": "unknown", "risk": 1, "sources": [], "depth": 0, "expr": "?"}
kind = v.kind
sources = sorted(v.input_sources())
depth = _deref_depth(v)
if kind == "const":
return {"class": "const", "risk": 0, "sources": [], "depth": 0, "expr": repr(v)}
if not v.contains_input():
if _has_call_ret(v):
cls, risk = "call_result", 2
elif _has_kind(v, "local"):
cls, risk = "local", 0
else:
cls, risk = "unknown", 1
return {"class": cls, "risk": risk, "sources": [], "depth": depth, "expr": repr(v)}
if depth >= 1:
risk = min(5, 3 + depth)
cls = "referenced-field-chain" if depth >= 2 else "referenced-field"
else:
risk = 3
cls = "input-direct"
return {"class": cls, "risk": risk, "sources": sources, "depth": depth, "expr": repr(v)}
def _deref_depth(v):
best = 0
stack = [(v, 0)]
while stack:
node, d = stack.pop()
if not isinstance(node, V):
continue
nd = d
if node.kind in ("load", "field"):
nd = d + 1
best = max(best, nd)
for c in (node.a, node.b):
if isinstance(c, V):
stack.append((c, nd))
return best
def _has_call_ret(v):
return _has_kind(v, "call_ret")
def _has_kind(v, kind):
stack = [v]
while stack:
node = stack.pop()
if not isinstance(node, V):
continue
if node.kind == kind:
return True
for c in (node.a, node.b):
if isinstance(c, V):
stack.append(c)
return False
@dataclass
class Finding:
sink: str
sink_kind: str
arg_role: str
site: int
entry: int
chain: tuple
verdict: dict
def key(self):
return (self.sink, self.site, self.arg_role, self.verdict["expr"])
def mine_findings(pe, functions, iat, cfgs=None, entry_cap=None,
max_depth=8, time_budget=120.0):
eng = InterpEngine(pe, functions, iat, cfgs=cfgs, max_depth=max_depth)
seen = {}
raw_count = [0]
def on_call(site, target, iname, args, machine, stack):
sink = _sink_name(iname)
if sink is None:
return
spec = SINKS[sink]
for idx, role in spec["args"].items():
if idx >= len(X64_ARG_REGS):
continue
val = args[X64_ARG_REGS[idx]]
verdict = classify_value(val)
if role in ("size", "len") and verdict["risk"] < 3:
continue
if role in ("dest", "ptr", "src", "mdl") and verdict["class"] in ("const",):
continue
raw_count[0] += 1
f = Finding(sink, spec["kind"], role, site, stack[0] if stack else 0,
stack, verdict)
k = (sink, site, role, verdict["expr"])
prev = seen.get(k)
if prev is None or verdict["depth"] > prev.verdict["depth"]:
seen[k] = f
targets = functions if entry_cap is None else functions[:entry_cap]
import time
t0 = time.time()
scanned = 0
for fn in targets:
if fn.size is None:
continue
try:
eng.run_function(fn.vaddr, on_call=on_call)
except Exception:
pass
scanned += 1
if time.time() - t0 > time_budget:
break
findings = list(seen.values())
findings.sort(key=lambda f: (-f.verdict["risk"], -f.verdict["depth"], f.site))
stats = {
"entries_scanned": scanned,
"raw_sink_observations": raw_count[0],
"unique_findings": len(findings),
"elapsed_sec": round(time.time() - t0, 1),
}
return findings, stats
def _internal_mask_bound(v):
stack = [v]
while stack:
node = stack.pop()
if not isinstance(node, V):
continue
if node.kind == "binop" and node.op == "&":
for side in (node.a, node.b):
if isinstance(side, V) and side.kind == "const" and isinstance(side.a, int):
return side.a
for c in (node.a, node.b):
if isinstance(c, V):
stack.append(c)
return None
_MEANINGFUL_KINDS = ("input", "field", "load")
def _meaningful_keys(keys):
return {k for k in keys if k[0] in _MEANINGFUL_KINDS}
def _is_trivial_constraint(c):
if isinstance(c.a, V) and isinstance(c.b, V):
if c.a.key() == c.b.key():
return True
return False
def _constraint_bounds(c, subkeys):
if _is_trivial_constraint(c):
return False
a_in = isinstance(c.a, V) and (_meaningful_keys(_all_subkeys(c.a)) & subkeys)
b_in = isinstance(c.b, V) and (_meaningful_keys(_all_subkeys(c.b)) & subkeys)
if not (a_in or b_in):
return False
other = c.b if a_in else c.a
if isinstance(other, V) and other.kind == "const":
return True
if a_in and b_in:
return True
return isinstance(other, V) and not other.contains_input()
_UPPER_WHEN_LHS = {"u<", "u<=", "s<", "s<="}
_UPPER_WHEN_RHS = {"u>", "u>=", "s>", "s>="}
_NONBOUNDING_REL = {"==", "!=", "sign", "nsign"}
def _constraint_upper_bounds(c, subkeys):
if _is_trivial_constraint(c):
return False
if c.rel in _NONBOUNDING_REL:
return False
a_in = isinstance(c.a, V) and bool(_meaningful_keys(_all_subkeys(c.a)) & subkeys)
b_in = isinstance(c.b, V) and bool(_meaningful_keys(_all_subkeys(c.b)) & subkeys)
if a_in and not b_in:
return c.rel in _UPPER_WHEN_LHS
if b_in and not a_in:
return c.rel in _UPPER_WHEN_RHS
return False
def assess_control(val, conds, role=None):
subkeys = _meaningful_keys(_all_subkeys(val))
mask = _internal_mask_bound(val)
guards = []
for c in conds:
if _is_trivial_constraint(c):
continue
cvars = _meaningful_keys(c.vars())
if cvars & subkeys:
guards.append(c)
bounding = [c for c in guards if _constraint_bounds(c, subkeys)]
upper = [c for c in guards if _constraint_upper_bounds(c, subkeys)]
size_like = role in ("size", "len")
if size_like:
if upper or mask is not None:
status = "bounded"
elif bounding:
status = "guarded-not-upper"
elif guards:
status = "related-guard"
else:
status = "unguarded"
else:
if bounding:
status = "guarded"
elif guards:
status = "related-guard"
elif mask is not None:
status = "masked"
else:
status = "unguarded"
return {
"status": status,
"mask": hex(mask) if mask is not None else None,
"upper_bounded": bool(upper) or (size_like and mask is not None),
"guards": [repr(c) for c in (upper or bounding)[:6]] or [repr(c) for c in guards[:4]],
"guard_count": len(guards),
"bounding_count": len(bounding),
}
_RISK_LABEL = {5: "CRIT", 4: "HIGH", 3: "MED", 2: "LOW", 1: "INFO", 0: "-"}
def _short(s, n=64):
return s if len(s) <= n else s[:n - 3] + "..."
def mine_findings_guarded(pe, functions, iat, cfgs=None, entry_cap=None,
time_budget=180.0, max_paths=48, trace=False,
per_func_secs=3.0, surface=None, rev_callgraph=None,
name_hints=None, callgraph=None, arg_flow=None):
eng = InterpEngine(pe, functions, iat, cfgs=cfgs, max_depth=6)
seen = {}
raw_count = [0]
starts_sorted = sorted(f.vaddr for f in functions if f.size is not None)
def owner_func(va):
i = bisect.bisect_right(starts_sorted, va) - 1
return starts_sorted[i] if i >= 0 else None
def on_sink(site, target, iname, args, state, stack):
sink = _sink_name(iname)
if sink is None:
return
spec = SINKS[sink]
conds = state.conds
depth = len(stack)
for idx, role in spec["args"].items():
if idx >= len(X64_ARG_REGS):
continue
val = args[X64_ARG_REGS[idx]]
verdict = classify_value(val)
if trace and verdict["class"] not in ("const", "local"):
indent = " " * min(depth, 6)
print(f"{indent}→ {sink}({role}) @ {hex(site)} "
f"= {_short(verdict['expr'])} [{verdict['class']}]", flush=True)
if role in ("size", "len") and verdict["risk"] < 3:
continue
if role in ("dest", "ptr", "src", "mdl") and verdict["class"] == "const":
continue
control = assess_control(val, conds, role=role)
raw_count[0] += 1
risk = verdict["risk"]
cstat = control["status"]
if role in ("size", "len"):
if cstat in ("unguarded", "related-guard", "guarded-not-upper"):
risk = min(5, risk + 1)
if cstat == "bounded":
risk = max(0, risk - 3)
else:
if cstat in ("guarded", "masked"):
risk = max(0, risk - 2)
f = Finding(sink, spec["kind"], role, site, stack[0] if stack else 0,
stack, dict(verdict, risk=risk, control=control))
k = (sink, site, role, verdict["expr"], control["status"])
prev = seen.get(k)
if prev is None or verdict["depth"] > prev.verdict["depth"]:
if prev is None and trace:
rl = _RISK_LABEL.get(risk, "?")
print(f" ⤷ FINDING [{rl}] control={control['status']} "
f"src={','.join(verdict['sources']) or '-'}", flush=True)
seen[k] = f
import time
explorer = PathExplorer(eng, on_sink=on_sink, max_paths=max_paths,
per_func_secs=per_func_secs, clock=time.time)
if trace:
pulse = Pulse(time.time, interval=2.0)
explorer.pulse = pulse
eng.pulse = pulse
targets = functions if entry_cap is None else functions[:entry_cap]
t0 = time.time()
scanned = 0
total = sum(1 for fn in targets if fn.size is not None)
last_beat = t0
if trace:
print(f" [findings] exploring {total} functions "
f"(branch-aware, interprocedural)...", flush=True)
for fn in targets:
if fn.size is None:
continue
explorer.paths_done = 0
fstart = time.time()
try:
explorer.explore(fn.vaddr)
except Exception:
pass
now = time.time()
fdur = now - fstart
scanned += 1
if trace and (fdur > 1.0 or explorer.timed_out):
flag = " (BAILED: per-func budget)" if explorer.timed_out else ""
print(f" [findings] slow fn {hex(fn.vaddr)} ({fn.size}b) "
f"took {fdur:.1f}s, {explorer.paths_done} paths{flag}", flush=True)
if trace and (scanned % 100 == 0 or now - last_beat >= 5.0):
rate = scanned / (now - t0) if now > t0 else 0
eta = (total - scanned) / rate if rate > 0 else 0
print(f" [findings] {scanned}/{total} fns, {len(seen)} findings, "
f"{now - t0:.0f}s elapsed, ~{rate:.0f} fn/s, ETA {eta:.0f}s",
flush=True)
last_beat = now
if now - t0 > time_budget:
if trace:
print(f" [findings] time budget {time_budget}s reached", flush=True)
break
findings = list(seen.values())
reachable_count = 0
confirmed_count = 0
if surface is not None and rev_callgraph is not None:
if trace:
print(f" [reach] computing attacker-reachability for "
f"{len(findings)} findings...", flush=True)
target_funcs = {owner_func(f.site) for f in findings}
target_funcs.discard(None)
reach = compute_reachability(target_funcs, surface, rev_callgraph)
flow = {}
if arg_flow is not None and callgraph is not None:
flow = compute_input_flow(target_funcs, surface, callgraph, arg_flow)
name_hints = name_hints or {}
for f in findings:
fnva = owner_func(f.site)
r = reach.get(fnva, {"reachable": False, "entries": [], "via": [],
"hops": None, "used_callback": False})
entry_names = []
for e in r["entries"]:
nm = name_hints.get(e, hex(e))
if nm not in entry_names:
entry_names.append(nm)
f.verdict["reach"] = {
"reachable": r["reachable"],
"hops": r["hops"],
"entries": entry_names[:6],
"via": r["via"],
"used_callback": r.get("used_callback", False),
"sink_func": hex(fnva) if fnva else None,
"sink_func_name": name_hints.get(fnva),
}
fl = flow.get(fnva, {"tainted_args": [], "any": False})
tainted_args = set(fl["tainted_args"])
value_argis = set()
for s in f.verdict.get("sources", []):
if s.startswith("arg"):
try:
value_argis.add(int(s[3:]))
except ValueError:
pass
value_is_input = f.verdict["class"] not in ("local", "const", "unknown")
value_taint = value_is_input and bool(value_argis & tainted_args)
confirmed = bool(r["reachable"] and value_taint)
f.verdict["input_flow"] = {
"flows": fl["any"],
"tainted_args": sorted(tainted_args),
"value_tainted": value_taint,
"confirmed": confirmed,
}
if r["reachable"]:
reachable_count += 1
else:
f.verdict["risk"] = max(0, f.verdict["risk"] - 2)
if confirmed:
confirmed_count += 1
elif r["reachable"] and arg_flow is not None and not fl["any"]:
f.verdict["risk"] = max(0, f.verdict["risk"] - 1)
if (f.arg_role in ("size", "len") and f.verdict["class"] == "input-direct"
and f.verdict["depth"] == 0):
f.verdict["risk"] = max(0, f.verdict["risk"] - 1)
f.verdict["trivial_forward"] = True
_DANGER = {"unguarded": 2, "guarded-not-upper": 2,
"related-guard": 1, "masked": 1}
def sort_key(f):
reach = f.verdict.get("reach", {})
flow = f.verdict.get("input_flow", {})
cstat = f.verdict.get("control", {}).get("status", "")
danger = _DANGER.get(cstat, 0)
confirmed = 1 if flow.get("confirmed") else 0
reachable = 1 if reach.get("reachable") else 0
trivial = 1 if f.verdict.get("trivial_forward") else 0
return (trivial, -danger, -confirmed, -reachable, -f.verdict["risk"],
-f.verdict["depth"], f.site)
findings.sort(key=sort_key)
unguarded = sum(1 for f in findings
if f.verdict.get("control", {}).get("status") == "unguarded")
stats = {
"entries_scanned": scanned,
"raw_sink_observations": raw_count[0],
"unique_findings": len(findings),
"unguarded_findings": unguarded,
"reachable_findings": reachable_count if surface is not None else None,
"input_confirmed_findings": (confirmed_count
if arg_flow is not None else None),
"elapsed_sec": round(time.time() - t0, 1),
}
return findings, stats
def analyze(path, min_str_len=4, disassemble=True, cfg=True, cfg_cap=None,
discover=True, findings=False, findings_cap=None, trace=False,
collisions=False):
log(f"[parse] loading PE: {path}")
pe = pefile.PE(path, fast_load=False)
data = bytes(pe.__data__)
log(f"[parse] {len(data):,} bytes, machine "
f"{hex(pe.FILE_HEADER.Machine)}, {len(pe.sections)} sections")
machine_map = {0x8664: "x86-64", 0x14C: "x86", 0xAA64: "arm64"}
info = {
"path": path,
"size": len(data),
"machine": machine_map.get(pe.FILE_HEADER.Machine,
hex(pe.FILE_HEADER.Machine)),
"image_base": pe.OPTIONAL_HEADER.ImageBase,
"entry_point_rva": pe.OPTIONAL_HEADER.AddressOfEntryPoint,
"sections": [{
"name": s.Name.rstrip(b"\x00").decode("latin-1", "replace"),
"vaddr": pe.OPTIONAL_HEADER.ImageBase + s.VirtualAddress,
"vsize": s.Misc_VirtualSize,
"raw_size": s.SizeOfRawData,
"characteristics": s.Characteristics,
} for s in pe.sections],
}
imports = []
if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
for entry in pe.DIRECTORY_ENTRY_IMPORT:
dll = entry.dll.decode("latin-1")
for imp in entry.imports:
nm = imp.name.decode("latin-1") if imp.name else f"ordinal_{imp.ordinal}"
imports.append({"dll": dll, "name": nm})
log(f"[imports] {len(imports)} imported symbols")
log("[strings] extracting ASCII + UTF-16LE strings")
strings = extract_strings(pe, data, min_len=min_str_len)
log(f"[strings] {len(strings):,} strings found")
log("[functions] recovering function boundaries "
"(.pdata + exports" + (" + discovery)" if discover else ")"))
functions, discovery_stats = recover_functions(
pe, disassemble=disassemble, discover=discover)
log(f"[functions] {len(functions):,} functions recovered")
report = {
"info": info,
"import_count": len(imports),
"imports": imports,
"string_count": len(strings),
"strings": [asdict(s) for s in strings],
"function_count": len(functions),
"functions": [asdict(f) for f in functions],
}
if discovery_stats is not None:
report["discovery_stats"] = discovery_stats
if cfg:
cfgs, callgraph, cfg_stats = build_cfgs(pe, functions, cap_functions=cfg_cap)
icount = defaultdict(int)
for names in callgraph["indirect"].values():
for nm in names:
if nm != "?":
icount[nm] += 1
report["cfg_stats"] = cfg_stats
report["top_indirect_callees"] = sorted(
icount.items(), key=lambda kv: -kv[1])[:40]
report["_cfgs"] = cfgs
report["_callgraph"] = callgraph
if findings:
iat = build_iat_map(pe)
cfgs = report.get("_cfgs")
callgraph = report.get("_callgraph")
if callgraph is None:
log("[reach] building call graph for reachability")
cfgs, callgraph, _ = build_cfgs(pe, functions, cap_functions=cfg_cap)
log("[reach] building name hints + attack surface")
string_dicts = report["strings"]
name_hints = build_name_hints(pe, functions, string_dicts)
surface = identify_attack_surface(pe, functions, callgraph,
string_dicts, name_hint=name_hints)
cb_edges = address_taken_edges(pe, functions)
log(f"[reach] {sum(len(v) for v in cb_edges.values())} callback edges "
f"(address-taken) added to reachability")
rev_cg = build_reverse_callgraph(callgraph, callback_edges=cb_edges)
log("[flow] building attacker-input arg-flow edges")
arg_flow = build_arg_flow_edges(pe, functions, iat, cfgs=cfgs, trace=trace)
if trace:
print("\n == live analysis trace ==", flush=True)
found, fstats = mine_findings_guarded(
pe, functions, iat, cfgs=cfgs, entry_cap=findings_cap, trace=trace,
surface=surface, rev_callgraph=rev_cg, name_hints=name_hints,
callgraph=callgraph, arg_flow=arg_flow)
report["findings_stats"] = fstats
report["attack_surface_size"] = len(surface)
report["findings"] = [{
"sink": f.sink, "sink_kind": f.sink_kind, "arg": f.arg_role,
"site": f.site, "entry": f.entry, "chain_depth": len(f.chain),
"chain": [hex(x) for x in f.chain],
"class": f.verdict["class"], "risk": f.verdict["risk"],
"sources": f.verdict["sources"], "deref_depth": f.verdict["depth"],
"expr": f.verdict["expr"],
"control": f.verdict.get("control"),
"reach": f.verdict.get("reach"),
"input_flow": f.verdict.get("input_flow"),
} for f in found]
if collisions:
log("[collide] building call-site contracts for collision detection")
contracts = build_callsite_contracts(
pe, functions, iat, cfgs=cfgs, trace=trace,
time_budget=180.0)
starts_sorted = sorted(f.vaddr for f in functions if f.size is not None)
def _owner(va):
i = bisect.bisect_right(starts_sorted, va) - 1
return starts_sorted[i] if i >= 0 else None
cols = detect_collisions(found, contracts, _owner, rev_cg,
name_hints, surface=surface, trace=trace)
report["collisions"] = cols
report["collision_count"] = len(cols)
return report
def _executable_sections(report):
IMAGE_SCN_MEM_EXECUTE = 0x20000000
return {s["name"] for s in report["info"]["sections"]
if s.get("characteristics", 0) & IMAGE_SCN_MEM_EXECUTE}
def _meaningful_strings(report):
exec_secs = _executable_sections(report)
out = []
for s in report["strings"]:
if not s["vaddr"] or s["section"] is None or s["section"] in exec_secs:
continue
t = s["text"]
letters = sum(c.isalnum() or c in " ._/:%-\\@" for c in t)
if len(t) >= 6 and letters / len(t) > 0.85:
out.append(s)
return out
def print_summary(report, max_items=20):
info = report["info"]
print(f"== {info['path']} ==")
print(f" machine : {info['machine']}")
print(f" size : {info['size']:,} bytes")
print(f" image base : {hex(info['image_base'])}")
print(f" entry point : {hex(info['image_base'] + info['entry_point_rva'])}")
print(f" sections : {len(info['sections'])}")
for s in info["sections"]:
print(f" {s['name']:<8} @ {hex(s['vaddr'])} vsize={s['vsize']:,}")
print()
print(f" imports : {report['import_count']}")
print(f" strings : {report['string_count']}")
print(f" functions : {report['function_count']}")
named = [f for f in report["functions"] if f["name"]]
print(f" of which named/exported: {len(named)}")
decoded = [f for f in report["functions"] if f["decoded_ok"]]
print(f" decoded cleanly : {len(decoded)}")
by_src = defaultdict(int)
for f in report["functions"]:
by_src[f["source"]] += 1
print(f" by source : " +
", ".join(f"{k}={v}" for k, v in sorted(by_src.items())))
if "discovery_stats" in report:
ds = report["discovery_stats"]
parts = [f"recursive={ds.get('recursive_discovered', 0)}",
f"prologue={ds.get('prologue_discovered', 0)}"]
if "seed_coverage_pct" in ds:
parts.append(f"seed-recall={ds['seed_coverage_pct']}%")
print(f" discovery : " + ", ".join(parts))
print()
interesting = _meaningful_strings(report)
print(f" -- first {max_items} meaningful strings "
f"({len(interesting):,} of {report['string_count']:,} total) --")
for s in interesting[:max_items]:
va = hex(s["vaddr"]) if s["vaddr"] else "??"
txt = s["text"] if len(s["text"]) <= 60 else s["text"][:57] + "..."
print(f" [{s['encoding']:>8}] {va:>12} {s['section'] or '-':<8} {txt!r}")
print()
print(f" -- first {max_items} functions --")
for f in report["functions"][:max_items]:
nm = f["name"] or "(unnamed)"
sz = f"{f['size']}b" if f["size"] else "?"
ic = f["insn_count"] if f["insn_count"] is not None else "?"
print(f" {hex(f['vaddr']):>12} {sz:>7} {ic:>4} insn [{f['source']}] {nm}")
if "cfg_stats" in report:
cs = report["cfg_stats"]
print()
print(f" -- control-flow / call graph --")
print(f" functions analyzed : {cs.get('functions_analyzed', 0):,}")
print(f" basic blocks : {cs.get('total_blocks', 0):,}")
print(f" cfg edges : {cs.get('total_edges', 0):,}")
print(f" direct call edges : {cs.get('direct_call_edges', 0):,}")
print(f" indirect callsites : {cs.get('indirect_call_sites', 0):,}")
if report.get("top_indirect_callees"):
print()
print(f" -- top {max_items} imported callees (by indirect call count) --")
for nm, n in report["top_indirect_callees"][:max_items]:
print(f" {n:>6}x {nm}")
if "findings" in report:
fs = report["findings_stats"]
print()
print(f" == assumption findings ==")
print(f" entries scanned : {fs['entries_scanned']:,}")
print(f" sink observations : {fs['raw_sink_observations']:,}")
print(f" unique findings : {fs['unique_findings']:,} ({fs['elapsed_sec']}s)")
if "unguarded_findings" in fs:
print(f" unguarded : {fs['unguarded_findings']:,}")
if fs.get("reachable_findings") is not None:
print(f" attacker-reachable : {fs['reachable_findings']:,}"
f" (surface: {report.get('attack_surface_size', 0)} entries)")
if fs.get("input_confirmed_findings") is not None:
print(f" input-flow CONFIRM : {fs['input_confirmed_findings']:,}"
f" (attacker input provably reaches the sink arg)")
print()
risk_label = {5: "CRIT", 4: "HIGH", 3: "MED", 2: "LOW", 1: "INFO", 0: "-"}
print(f" -- top {max_items} findings (most-dangerous control first, then confirmed) --")
for f in report["findings"][:max_items]:
rl = risk_label.get(f["risk"], "?")
src = ",".join(f["sources"]) if f["sources"] else "-"
expr = f["expr"] if len(f["expr"]) <= 52 else f["expr"][:49] + "..."
ctrl = f.get("control") or {}
cstatus = ctrl.get("status", "?")
reach = f.get("reach") or {}
flow = f.get("input_flow") or {}
if flow.get("confirmed"):
rtag = "INPUT-CONFIRMED"
elif reach.get("reachable"):
rtag = "reachable(no-flow)"
else:
rtag = "unreachable"
print(f" [{rl:>4}] {f['sink']}({f['arg']}) @ {hex(f['site'])} "
f"deref={f['deref_depth']} control={cstatus} {rtag}")
print(f" {f['class']:<22} src={src} {expr}")
if reach.get("reachable"):
fn = reach.get("sink_func_name") or reach.get("sink_func") or "?"
ents = ", ".join(reach.get("entries", [])[:2]) or "?"
cb = " via callback" if reach.get("used_callback") else ""
ta = flow.get("tainted_args")
tainfo = f", tainted args {ta}" if ta else ""
print(f" in {fn}, {reach.get('hops')} hops{cb} from: {ents}{tainfo}")
if ctrl.get("guards"):
print(f" guards: {'; '.join(ctrl['guards'][:2])}")
if "collisions" in report:
cols = report["collisions"]
print()
print(f" == contract collisions ({len(cols)}) ==")
print(f" caller passes an unbounded value into a callee arg that the")
print(f" callee then uses unchecked -- neither side bounds it")
print()
for c in cols[:max_items]:
caller = c.get("caller_name") or c["caller"]
callee = c.get("callee_name") or c["callee"]
tainted = " [attacker-tainted]" if c.get("carries_input") else ""
entry = " [reaches entry]" if c.get("reaches_entry") else ""
hops = c.get("hops", 1)
print(f" {caller} -> {callee}(arg{c['callee_argi']}) "
f"{hops}-hop chain{tainted}{entry}")
print(f" passes {_short(c['value'], 56)} ({c['caller_status']})")
print(f" -> {c['sink']} unchecked @ {c['sink_site']}")
chain = c.get("chain") or []
if len(chain) > 1:
for h in reversed(chain):
nm = h.get("name") or h["func"]
mark = " (entry)" if h.get("is_entry") else ""
print(f" via {nm}{mark} passes unbounded @ {h['site']}")
def _json_safe(report):
return {k: v for k, v in report.items() if not k.startswith("_")}
def main(argv=None):
ap = argparse.ArgumentParser()
ap.add_argument("binary")
ap.add_argument("--min-str", type=int, default=4)
ap.add_argument("--no-disasm", action="store_true")
ap.add_argument("--no-cfg", action="store_true")
ap.add_argument("--no-discover", action="store_true")
ap.add_argument("--findings", action="store_true")
ap.add_argument("--collisions", action="store_true")
ap.add_argument("--trace", action="store_true")
ap.add_argument("--findings-cap", type=int, default=None)
ap.add_argument("--cfg-cap", type=int, default=None)
ap.add_argument("--json", metavar="FILE")
ap.add_argument("--max-items", type=int, default=20)
args = ap.parse_args(argv)
set_trace(args.trace)
report = analyze(args.binary, min_str_len=args.min_str,
disassemble=not args.no_disasm,
cfg=not args.no_cfg, cfg_cap=args.cfg_cap,
discover=not args.no_discover,
findings=args.findings or args.trace or args.collisions,
findings_cap=args.findings_cap, trace=args.trace,
collisions=args.collisions)
print_summary(report, max_items=args.max_items)
if args.json:
with open(args.json, "w", encoding="utf-8") as fh:
json.dump(_json_safe(report), fh, indent=2)
print(f"\n full report written to {args.json}")
return 0
if __name__ == "__main__":
sys.exit(main())
Reference in New Issue View Git Blame Copy Permalink