"""Sweep `push imm32` (0x68) operands inside a VA range of a PE binary,
dereference each into a data section (.rdata/.data), and decode any that
resolve to a printable string literal -- UTF-16LE (PStringBase<unsigned
short>, e.g. ECM_UI notice text) or narrow ASCII (PStringBase<char>, e.g.
the ClientCommunicationSystem::Help* command-help family) alike.

Built for the CH2 REJECT-review rework (BLOCKER 2,
docs/research/2026-08-09-ch2-review-findings.md) to re-derive
ClientCommunicationSystem::HandleFailureEvent (@0x00571990)'s 344-row
display-string table from ground truth instead of the Binary Ninja
pseudo-C's ~33-char inline preview -- the same class of problem
check_exe_pdb.py and dump_pdb_info.py solve for PDB metadata, applied to
string literal recovery. Not tied to WeenieError specifically: any VA
range in any PDB-paired PE binary works.

Generalized for Campaign CH user-gate round 2, item 3 (2026-08-09): the
retail help command family (ClientCommunicationSystem::Help*, e.g.
HelpAllGroup/HelpAllegiancesGroup/HelpChannelsGroup/HelpChattingGroup/
HelpDeathGroup/HelpStatusGroup/HelpTextGroup) constructs its strings via
PStringBase<char> (narrow 8-bit ASCII), NOT PStringBase<unsigned short>
(wide UTF-16LE) like the WeenieError table or ECM_UI notices. Both
encodings now share one sweep: each push-imm32 hit is decoded as UTF-16LE
first (a real wide string reads back false as ASCII almost immediately --
every other byte is 0x00, which read_ascii_cstr rejects as a control
character), then as ASCII if that fails. The reported tuple carries which
encoding matched so callers can tell narrow help text apart from wide
notice text at a glance.

ALWAYS run check_exe_pdb.py first to confirm the candidate .exe pairs with
the PDB you're cross-referencing addresses against -- a mismatched binary
will produce confident-looking garbage.

Usage:
    py tools/pdb-extract/sweep_weenie_strings.py <exe_path> --range 0x571990 0x575480 [--min-len 4]
    py tools/pdb-extract/sweep_weenie_strings.py <exe_path> --anchor 0x005750a5 [--window 64]
    py tools/pdb-extract/sweep_weenie_strings.py <exe_path> --deref 0x007d2ee8

--range LO HI   sweep every string-valued push imm32 in [LO, HI)
--anchor VA     search backward `--window` bytes from a case-body/call-site
                VA (taken from the pseudo-C) for the nearest string-valued
                push -- use when you already know roughly where a specific
                case lives and just need its untruncated literal
--deref VA      dereference one known data pointer directly (e.g. a
                `data_XXXXXXXX` symbol name from the pseudo-C, which
                directly encodes its own VA in hex)
--min-len N     minimum decoded string length to report (default 3);
                raise this to cut noise from short accidental hits
--ascii-only    skip the UTF-16LE attempt entirely (narrow-string ranges
                run faster and cannot false-positive against wide data)
"""
import argparse
import struct


class PeImage:
    def __init__(self, path):
        with open(path, "rb") as f:
            self.data = f.read()
        if self.data[0:2] != b"MZ":
            raise ValueError("not a PE file (no MZ)")
        e_lfanew = struct.unpack_from("<I", self.data, 0x3C)[0]
        if self.data[e_lfanew:e_lfanew + 4] != b"PE\0\0":
            raise ValueError("no PE signature")
        coff_off = e_lfanew + 4
        machine, num_sections, ts, symtab, numsym, opt_hdr_size, characteristics = \
            struct.unpack_from("<HHIIIHH", self.data, coff_off)
        opt_off = coff_off + 20
        magic = struct.unpack_from("<H", self.data, opt_off)[0]
        if magic != 0x10B:
            raise ValueError(f"unexpected optional header magic 0x{magic:04x} (want PE32)")
        self.image_base = struct.unpack_from("<I", self.data, opt_off + 28)[0]
        sec_off = opt_off + opt_hdr_size
        self.sections = []
        for i in range(num_sections):
            rec = self.data[sec_off + i * 40: sec_off + (i + 1) * 40]
            name = rec[0:8].rstrip(b"\0").decode("ascii", "replace")
            virt_size, virt_addr, raw_size, raw_ptr = struct.unpack_from("<IIII", rec, 8)
            self.sections.append({
                "name": name,
                "va": self.image_base + virt_addr,
                "vsize": virt_size,
                "raw_ptr": raw_ptr,
                "raw_size": raw_size,
            })

    def section_for_va(self, va):
        for s in self.sections:
            if s["va"] <= va < s["va"] + max(s["vsize"], s["raw_size"]):
                return s
        return None

    def va_to_off(self, va):
        s = self.section_for_va(va)
        if s is None:
            return None
        off = s["raw_ptr"] + (va - s["va"])
        if off < 0 or off >= len(self.data):
            return None
        return off

    def read_bytes(self, va, n):
        off = self.va_to_off(va)
        if off is None:
            return None
        return self.data[off:off + n]

    def read_utf16_cstr(self, va, max_chars=400):
        off = self.va_to_off(va)
        if off is None:
            return None
        out = []
        for i in range(max_chars):
            chunk = self.data[off + i * 2: off + i * 2 + 2]
            if len(chunk) < 2:
                break
            code = struct.unpack("<H", chunk)[0]
            if code == 0:
                return "".join(out)
            # Reject control chars other than the ones AC strings legitimately
            # use (\n, \t) -- anything else means we've wandered off a real
            # string into unrelated data and should not report a hit.
            if code < 0x20 and code not in (0x0A, 0x09):
                return None
            if code > 0x2FFF:
                return None
            out.append(chr(code))
        return None  # ran off the end without a NUL -- not a bounded literal

    def read_ascii_cstr(self, va, max_chars=800):
        """Decode a narrow (8-bit) NUL-terminated C string -- the
        PStringBase<char> literal shape the Help* command family uses,
        distinct from read_utf16_cstr's PStringBase<unsigned short>
        shape."""
        off = self.va_to_off(va)
        if off is None:
            return None
        out = []
        for i in range(max_chars):
            chunk = self.data[off + i: off + i + 1]
            if len(chunk) < 1:
                break
            code = chunk[0]
            if code == 0:
                return "".join(out)
            # Same control-char allowance as read_utf16_cstr (\n, \t only);
            # anything else (including high bytes outside printable ASCII)
            # means this isn't a real narrow literal.
            if code < 0x20 and code not in (0x0A, 0x09):
                return None
            if code > 0x7E:
                return None
            out.append(chr(code))
        return None  # ran off the end without a NUL -- not a bounded literal

    def is_data_section(self, va):
        s = self.section_for_va(va)
        return s is not None and s["name"] in (".rdata", ".data")

    def sweep_push_imm32(self, lo, hi, min_len=3, ascii_only=False):
        """Scan [lo, hi) for `push imm32` (opcode 0x68) whose operand VA
        dereferences to a printable string literal in .rdata/.data --
        UTF-16LE tried first (unless ascii_only), narrow ASCII as the
        fallback (a real wide string's alternating 0x00 bytes make the
        ASCII decode reject it as a control character almost immediately,
        so the two encodings do not cross-contaminate each other's hits).
        Returns a list of (instr_va, target_va, text, encoding)."""
        hits = []
        off_lo = self.va_to_off(lo)
        off_hi = self.va_to_off(hi)
        if off_lo is None or off_hi is None:
            raise ValueError("range not mapped")
        i = off_lo
        while i < off_hi - 4:
            if self.data[i] == 0x68:
                operand = struct.unpack_from("<I", self.data, i + 1)[0]
                if self.is_data_section(operand):
                    text = None
                    encoding = None
                    if not ascii_only:
                        text = self.read_utf16_cstr(operand)
                        encoding = "utf16" if text is not None else None
                    if text is None:
                        text = self.read_ascii_cstr(operand)
                        encoding = "ascii" if text is not None else None
                    if text is not None and len(text) >= min_len:
                        instr_va = lo + (i - off_lo)
                        hits.append((instr_va, operand, text, encoding))
            i += 1
        return hits

    def find_push_before(self, anchor_va, window=64, min_len=3, ascii_only=False):
        """Search backward from anchor_va (a call-site VA taken from the
        pseudo-C) for the nearest preceding `push imm32` whose operand
        dereferences to a printable string literal."""
        lo = anchor_va - window
        return self.sweep_push_imm32(
            lo, anchor_va + 2, min_len=min_len, ascii_only=ascii_only)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("exe")
    ap.add_argument("--range", nargs=2, metavar=("LO", "HI"))
    ap.add_argument("--anchor", action="append", default=[])
    ap.add_argument("--window", type=int, default=64)
    ap.add_argument("--deref", action="append", default=[])
    ap.add_argument("--min-len", type=int, default=3)
    ap.add_argument("--ascii-only", action="store_true")
    args = ap.parse_args()

    pe = PeImage(args.exe)
    print(f"ImageBase=0x{pe.image_base:08x}  sections:")
    for s in pe.sections:
        print(f"  {s['name']:<9} VA=0x{s['va']:08x} vsize=0x{s['vsize']:06x} "
              f"rawptr=0x{s['raw_ptr']:08x} rawsize=0x{s['raw_size']:06x}")
    print()

    if args.range:
        lo = int(args.range[0], 16)
        hi = int(args.range[1], 16)
        hits = pe.sweep_push_imm32(
            lo, hi, min_len=args.min_len, ascii_only=args.ascii_only)
        print(f"# sweep 0x{lo:08x}-0x{hi:08x}: {len(hits)} string-valued push imm32 sites")
        for instr_va, target_va, text, encoding in hits:
            print(f"0x{instr_va:08x}  -> data_0x{target_va:08x}  [{encoding}]  {text!r}")

    for a in args.anchor:
        anchor = int(a, 16)
        hits = pe.find_push_before(
            anchor, window=args.window, min_len=args.min_len,
            ascii_only=args.ascii_only)
        print(f"\n# anchor 0x{anchor:08x} (window={args.window}): {len(hits)} hits")
        for instr_va, target_va, text, encoding in hits:
            print(f"0x{instr_va:08x}  -> data_0x{target_va:08x}  [{encoding}]  {text!r}")

    for d in args.deref:
        target = int(d, 16)
        text = pe.read_utf16_cstr(target)
        if text is None:
            text = pe.read_ascii_cstr(target)
            kind = "ascii"
        else:
            kind = "utf16"
        print(f"\n# deref 0x{target:08x}: [{kind}] {text!r}")


if __name__ == "__main__":
    main()