re disasm: each [esp+N]/[ebp+N] operand is named as the argument or local it is (argN, local entry-0xN, return address), from the stack depth tracked over the function's control flow; _alloca_probe frames and known callee cleanup included, inferred cleanup marked ~

harbor/reverse.py

@@ -257,6 +257,7 @@ class Reverse:
             raise HarborError(f"the function is {size} bytes; pass max_bytes (at most {MAX_BYTES}) and walk it in parts")
         code = self.read(start, size)
         md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_32)
+        md.detail = True
         lines = []
         if func and not max_bytes:
             lines.append(f"; {func[2] or 'unnamed'} 0x{func[0]:X}..0x{func[0] + func[1]:X} ({func[1]} bytes)")
@@ -264,11 +265,148 @@ class Reverse:
             if conv:
                 lines.append(f"; convention from the bytes: {conv['summary']}; registers in: "
                              f"{', '.join(conv['registers_in']) or 'none'}; ret pops {conv['ret_pops']}")
+        frames = self._stack_depths(md, func) if func else {}
+        if frames:
+            lines.append("; stack notes: argN = the Nth stack argument (entry esp + 4N), local = below the return "
+                         "address, ~ = after a call whose cleanup was inferred from the pushes")
         for ins in md.disasm(code, start):
-            note = self._annotate(ins.op_str, ins.mnemonic, start, start + size)
+            notes = [n for n in (self._frame_note(ins, frames.get(ins.address)),
-            lines.append(f"{ins.address:08X}  {ins.mnemonic:7} {ins.op_str}{('    ; ' + note) if note else ''}")
+                                 self._annotate(ins.op_str, ins.mnemonic, start, start + size)) if n]
+            lines.append(f"{ins.address:08X}  {ins.mnemonic:7} {ins.op_str}{('    ; ' + ', '.join(notes)) if notes else ''}")
         return "\n".join(lines)
 
+    def _callee_pops(self, target):
+        """Bytes a directly called function pops on return (its `ret N`), or None when not known."""
+        key = ("pops", target)
+        if key not in self._cache:
+            pops = None
+            try:
+                head = self.read(target, 5)
+                if head == b"\x3d\x00\x10\x00\x00":           # cmp eax, 0x1000: the stack probe (_alloca_probe)
+                    pops = "probe"
+                else:
+                    conv = self.convention(target) if self.function_at(target) and \
+                        self.function_at(target)[0] == target else None
+                    if conv and len(conv["ret_pops"]) == 1:
+                        pops = conv["ret_pops"][0]
+            except Exception:  # noqa: BLE001 - unknown cleanup is reported as inferred
+                pops = None
+            self._cache[key] = pops
+        return self._cache[key]
+
+    def _stack_depths(self, md, func):
+        """{address: (bytes pushed since entry or None, ebp's depth or None, inferred)} over the function's control
+        flow, so an [esp+N] or [ebp+N] operand can be named as the argument or local it is."""
+        import capstone
+        from capstone import x86
+        start, size = func[0], func[1]
+        if size > 0x8000:
+            return {}
+        insns = list(md.disasm(self.read(start, size), start))
+        index = {ins.address: i for i, ins in enumerate(insns)}
+        state = {start: (0, None, 0, False)}       # depth, ebp depth, bytes pushed since the last call, inferred
+        work = [start]
+        while work:
+            addr = work.pop()
+            i = index.get(addr)
+            while i is not None and i < len(insns):
+                ins = insns[i]
+                d, e, p, guess = state[ins.address]
+                ops, m, text = ins.operands, ins.mnemonic, ins.op_str
+                imm = ops[1].imm if len(ops) == 2 and ops[1].type == x86.X86_OP_IMM else None
+                if m == "push":
+                    d, p = (d + 4 if d is not None else None), p + 4
+                elif m == "pushal":
+                    d = d + 32 if d is not None else None
+                elif m == "pop":
+                    d, p = (d - 4 if d is not None else None), max(0, p - 4)
+                    if text == "ebp":
+                        e = None
+                elif m == "popal":
+                    d = d - 32 if d is not None else None
+                elif m == "sub" and text.startswith("esp, ") and imm is not None:
+                    d = d + imm if d is not None else None
+                elif m == "add" and text.startswith("esp, ") and imm is not None:
+                    d, p = (d - imm if d is not None else None), max(0, p - imm)
+                elif m == "and" and text.startswith("esp, "):
+                    d = None                                   # aligned: esp no longer tracks the entry
+                elif m == "mov" and text == "ebp, esp":
+                    e = d
+                elif m == "mov" and text == "esp, ebp":
+                    d = e
+                elif m == "leave":
+                    d, e = (e - 4 if e is not None else None), None
+                elif m == "lea" and text.startswith("esp, [ebp") and e is not None:
+                    d = e - ops[1].mem.disp
+                elif m == "call":
+                    pops = None
+                    if ops and ops[0].type == x86.X86_OP_IMM:
+                        pops = self._callee_pops(ops[0].imm)
+                    if pops == "probe":                        # mov eax, N / call _alloca_probe reserves N bytes
+                        prev = insns[i - 1] if i else None
+                        pv = prev.operands[1].imm if prev is not None and prev.mnemonic == "mov" and \
+                            prev.op_str.startswith("eax, ") and prev.operands[1].type == x86.X86_OP_IMM else None
+                        d = d + pv if d is not None and pv is not None else None
+                    elif pops is not None:
+                        d = d - pops if d is not None else None
+                    else:
+                        nxt = insns[i + 1] if i + 1 < len(insns) else None
+                        cleaned = nxt is not None and nxt.mnemonic == "add" and nxt.op_str.startswith("esp, ")
+                        if not cleaned and p:
+                            d, guess = (d - p if d is not None else None), True   # the callee popped what was pushed
+                    p = 0
+                succ = []
+                if ins.group(capstone.CS_GRP_RET):
+                    pass
+                elif m == "jmp":
+                    if ops and ops[0].type == x86.X86_OP_IMM and ops[0].imm in index:
+                        succ.append(ops[0].imm)
+                else:
+                    if ins.group(capstone.CS_GRP_JUMP) and ops and ops[0].type == x86.X86_OP_IMM \
+                            and ops[0].imm in index:
+                        succ.append(ops[0].imm)
+                    if i + 1 < len(insns):
+                        succ.append(insns[i + 1].address)
+                fall = insns[i + 1].address if i + 1 < len(insns) else None
+                i = None
+                for s in succ:
+                    if s in state:
+                        continue
+                    state[s] = (d, e, p, guess)
+                    if s == fall:
+                        i = index[s]
+                    else:
+                        work.append(s)
+        return state
+
+    @staticmethod
+    def _frame_note(ins, frame):
+        """'esp+0x10 = arg2' style notes for the stack operands of one instruction."""
+        if not frame:
+            return ""
+        from capstone import x86
+        d, e, _, guess = frame
+        notes = []
+        for op in ins.operands:
+            if op.type != x86.X86_OP_MEM or op.mem.index:
+                continue
+            base = ins.reg_name(op.mem.base) if op.mem.base else ""
+            if base == "esp" and d is not None:
+                off = op.mem.disp - d
+            elif base == "ebp" and e is not None:
+                off = op.mem.disp - e
+            else:
+                continue
+            if off == 0:
+                what = "return address"
+            elif off >= 4:
+                k, r = divmod(off - 4, 4)
+                what = f"arg{k + 1}" + (f"+{r}" if r else "")
+            else:
+                what = f"local entry-0x{-off:X}"
+            notes.append(("~" if guess else "") + what)
+        return ", ".join(notes)
+
     def _annotate(self, op_str, mnemonic, lo, hi):
         notes = []
         constant = re.search(r"(dword|qword) ptr \[0x([0-9a-f]{6,8})\]", op_str) if mnemonic.startswith("f") else None