#!/usr/bin/env python3
"""Dump the VT1200 firmware heap free lists over rdebug `peek`.

The firmware has a real bit-addressed allocator (malloc `0x07EE0246`, free
`0x07EE0488`, realloc near `0x07EE0B84`).  It keeps:

  * a free-bytes counter at bit `0x1E8BA0` (32-bit, in **bits**)
  * a global free-list head at bit `0x1E8AC0`
  * three arenas at `0x1E8930` / `0x1E89B0` / `0x1E8A30`, each with its free-list
    head at `+0x60`

Each free block is a 0x40-bit (8-byte) header:

  +0x00  size in bits (includes the 8-byte header)
  +0x20  next block, a **bit** address

This tool walks them, so payload placement can be checked against what the OS
may actually hand out.  Note that the firmware's payload scratch addresses
(bit `0x01800000` = byte `0x00300000`) sit **inside** the largest free block on
a default boot -- see `dwt/notes/HEAP.md`.

The **global** list is the one that matters for payload placement and is
reliable.  The per-arena lists are optional (`--arenas`); following one has
been observed to kill the rdebug process, so each list is read on its own
connection.

    python3 dwt/heapmap.py --host 192.168.99.2
    python3 dwt/heapmap.py --host 192.168.99.2 --at 0x01800000
    python3 dwt/heapmap.py --host 192.168.99.2 --blocks --arenas
"""
from __future__ import annotations

import argparse
import struct
import sys

import exploit

FREE_COUNTER = 0x1E8BA0
GLOBAL = ("global", 0x1E8AC0)
ARENAS = (
    ("arena0", 0x1E8990),
    ("arena1", 0x1E8A10),
    ("arena2", 0x1E8A90),
)
RAM_BITS = 0x02000000          # 4 MiB, as bit addresses
HDR_NEXT = 0x20                # bits
HDR_BYTES = 8


def u32(rd: exploit.RDebug, bit: int) -> int:
    return struct.unpack("<I", rd.peek(bit, 4))[0]


def walk(rd: exploit.RDebug, head: int, limit: int = 400):
    p = u32(rd, head)
    out, seen = [], set()
    while p and len(out) < limit and p not in seen and 8 <= p < RAM_BITS:
        seen.add(p)
        sz = u32(rd, p)
        nxt = u32(rd, p + HDR_NEXT)
        out.append((p, sz, nxt))
        if not (8 <= nxt < RAM_BITS):
            break
        p = nxt
    return out


def read_list(host, port, timeout, name, head):
    """Read one list on its own connection; a bad list must not abort the rest."""
    rd = exploit.RDebug(host, port, timeout, False)
    rd.connect()
    try:
        return walk(rd, head)
    finally:
        rd.close()


def main(argv):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--host", default="192.168.99.2")
    ap.add_argument("--port", type=int, default=771)
    ap.add_argument("--timeout", type=float, default=8.0)
    ap.add_argument("--at", type=lambda s: int(s, 0), default=None,
                    help="bit address to check (e.g. a payload scratch address)")
    ap.add_argument("--blocks", action="store_true", help="list every block")
    ap.add_argument("--arenas", action="store_true",
                    help="also try the per-arena lists (can kill the rdebug process)")
    a = ap.parse_args(argv[1:])

    rd = exploit.RDebug(a.host, a.port, a.timeout, False)
    rd.connect()
    try:
        free_bits = u32(rd, FREE_COUNTER)
    finally:
        rd.close()
    print(f"free counter: {free_bits} bits = {free_bits // 8} bytes "
          f"({free_bits // 8000} KB decimal)")

    lists = (GLOBAL,) + (ARENAS if a.arenas else ())
    all_blocks = []
    for name, head in lists:
        try:
            blocks = read_list(a.host, a.port, a.timeout, name, head)
        except Exception as exc:                      # noqa: BLE001
            print(f"{name:6s}: unreadable ({type(exc).__name__}: {exc})")
            continue
        if not blocks:
            print(f"{name:6s}: empty")
            continue
        lo = min(b for b, _, _ in blocks)
        hi = max(b + s for b, s, _ in blocks)
        usable = sum(s // 8 - HDR_BYTES for _, s, _ in blocks
                     if s >= HDR_BYTES * 8)
        print(f"{name:6s}: {len(blocks)} blocks, byte 0x{lo>>3:06X}..0x{hi>>3:06X}, "
              f"{usable} bytes usable")
        all_blocks += [(name, b, s) for b, s, _ in blocks]
        if a.blocks:
            for b, s, nxt in blocks:
                print(f"        bit 0x{b:08X} byte 0x{b>>3:06X}  "
                      f"size {s//8:>9} B  end 0x{(b+s)>>3:06X}  next 0x{nxt:08X}")

    if a.at is not None:
        hits = [(n, b, s) for n, b, s in all_blocks if b <= a.at < b + s]
        print(f"check bit 0x{a.at:08X} (byte 0x{a.at>>3:06X}): "
              + ("INSIDE a free block" if hits else "outside every free block"))
        for n, b, s in hits:
            print(f"        {n}: block bit 0x{b:08X}, "
                  f"{(b+s-a.at)//8} bytes of headroom to its end")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv))
    except (exploit.ProtocolError, OSError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(1)