#!/usr/bin/env python3
"""Race-window: dump TNG executable maps and hunt kill+SVC after launch."""
import re
import subprocess
import sys
import time
from pathlib import Path

ADB = r"C:\Users\Administrator\AppData\Local\Android\Sdk\platform-tools\adb.exe"
PKG = "my.com.tngdigital.ewallet"
OUT = Path(__file__).resolve().parents[1] / "dumps" / "tng_runtime"
OUT.mkdir(parents=True, exist_ok=True)

# aarch64 movz x8/w8,#129/130/131 + svc #0
KILL_IMMS = [
    bytes.fromhex("281080d2"),
    bytes.fromhex("28108052"),
    bytes.fromhex("481080d2"),
    bytes.fromhex("48108052"),
    bytes.fromhex("681080d2"),
    bytes.fromhex("68108052"),
]
SVC0 = bytes.fromhex("010000d4")
EXIT_IMM = bytes.fromhex("c80b80d2")  # movz x8,#94 exit_group
EXIT2 = bytes.fromhex("ba0b80d2")  # movz x8,#93 exit


def adb(*args, check=False):
    r = subprocess.run([ADB, *args], capture_output=True)
    out = (r.stdout or b"") + (r.stderr or b"")
    if check and r.returncode != 0:
        raise RuntimeError(out.decode("utf-8", "ignore"))
    return r.returncode, out.decode("utf-8", "replace")


def su(cmd):
    return adb("shell", f"su -c '{cmd}'")


def main():
    adb("shell", "am", "force-stop", PKG)
    time.sleep(0.5)
    adb("shell", "monkey", "-p", PKG, "-c", "android.intent.category.LAUNCHER", "1")

    pid = None
    for _ in range(40):
        time.sleep(0.25)
        _, out = adb("shell", "pidof", PKG)
        toks = out.strip().split()
        if toks:
            pid = toks[0]
            break
    if not pid:
        print("FAIL: no pid")
        return 1
    print(f"pid={pid}")

    # pull maps
    _, maps = su(f"cat /proc/{pid}/maps")
    maps_path = OUT / f"maps_{pid}.txt"
    maps_path.write_text(maps, encoding="utf-8", errors="replace")
    print(f"maps -> {maps_path} lines={len(maps.splitlines())}")

    targets = []
    for line in maps.splitlines():
        if "r-xp" not in line and "r-x" not in line:
            # also rw-p with execute rarely; keep x
            if "x" not in line.split()[1] if len(line.split()) > 1 else "":
                continue
        if "tngdigital" in line or "libnative" in line or "ewallet" in line.lower():
            parts = line.split()
            rng = parts[0]
            start_s, end_s = rng.split("-")
            start, end = int(start_s, 16), int(end_s, 16)
            path = parts[-1] if len(parts) >= 6 else ""
            targets.append((start, end, path, line))

    print(f"target segments={len(targets)}")
    for start, end, path, line in targets[:12]:
        print(f"  {hex(start)}-{hex(end)} {path}")

    # dump via dd from /proc/pid/mem
    remote = f"/data/local/tmp/tng_rt_{pid}.bin"
    su(f"rm -f {remote}")
    total = 0
    for i, (start, end, path, _) in enumerate(targets):
        size = end - start
        if size <= 0 or size > 32 * 1024 * 1024:
            continue
        # append dump
        cmd = (
            f"dd if=/proc/{pid}/mem bs=4096 skip={start // 4096} "
            f"count={(size + 4095) // 4096} 2>/dev/null >> {remote}"
        )
        # dd skip is in blocks from file start — wrong for /proc/pid/mem!
        # Use busybox dd with seek on output and skip via python on device instead.
        cmd = (
            f"toybox dd if=/proc/{pid}/mem of={remote}.p{i} "
            f"bs=1 skip={start} count={size} 2>/dev/null"
        )
        code, _ = su(cmd)
        if code == 0:
            total += size
            print(f"  dumped p{i} size={size} from {path}")
        else:
            # fallback: python on device
            py = (
                f"python3 -c \"import sys;f=open('/proc/{pid}/mem','rb');"
                f"f.seek({start});d=f.read({size});open('{remote}.p{i}','wb').write(d)\""
            )
            code2, out2 = su(py)
            if code2 == 0:
                total += size
                print(f"  dumped p{i} via python size={size}")
            else:
                print(f"  FAIL dump p{i}: {out2[:120]}")

    # pull pieces and scan
    local_dir = OUT / f"mem_{pid}"
    local_dir.mkdir(exist_ok=True)
    kill_hits = 0
    exit_hits = 0
    for i, (start, end, path, _) in enumerate(targets):
        rem = f"{remote}.p{i}"
        loc = local_dir / f"seg_{i}_{start:x}.bin"
        code, _ = adb("shell", f"su -c 'test -f {rem} && echo OK'")
        if "OK" not in _:
            continue
        adb("pull", rem, str(loc))
        if not loc.exists():
            continue
        data = loc.read_bytes()
        for imm in KILL_IMMS:
            pos = 0
            while True:
                j = data.find(imm, pos)
                if j < 0:
                    break
                win = data[j : j + 36]
                if SVC0 in win:
                    kill_hits += 1
                    delta = win.find(SVC0)
                    print(f"KILL+SVC seg{i} file+0x{j:x} va=0x{start+j:x} delta={delta} path={path}")
                pos = j + 4
        for imm in (EXIT_IMM, EXIT2):
            pos = 0
            while True:
                j = data.find(imm, pos)
                if j < 0:
                    break
                win = data[j : j + 36]
                if SVC0 in win:
                    exit_hits += 1
                    if exit_hits <= 15:
                        print(f"EXIT+SVC seg{i} file+0x{j:x} va=0x{start+j:x} path={path}")
                pos = j + 4
        # also count raw svc
        print(f"  seg{i} svc0={data.count(SVC0)} size={len(data)}")

    print(f"DONE kill+svc={kill_hits} exit+svc={exit_hits} dumped_bytes~={total}")
    _, alive = adb("shell", "pidof", PKG)
    print(f"still alive? {alive.strip() or 'NO'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())