tdkpin/original/tools/inject_original_state.py

#!/usr/bin/env python3
"""Inject a controlled ball state into the running original for collision probes."""

from __future__ import annotations

import argparse
import os
import signal
import struct
import time

from trace_original_state import (
    BALL_SUSPENDED,
    BALL_X,
    BALL_Y,
    DATA_READ_SIZE,
    NEXT_BALL_X,
    NEXT_BALL_Y,
    OBJECT_BALL_X,
    OBJECT_BALL_Y,
    PHYSICS_SUBSTEPS,
    PREVIOUS_BALL_X,
    PREVIOUS_BALL_Y,
    SPIN_REAL48,
    VELOCITY_X,
    VELOCITY_Y,
    ProcessMemory,
    find_winevdm_pid,
    locate_data_segment,
    locate_game_object,
    verify_original,
)

BALL_ACTIVE = 0x0BDD
PLAYER_ENTRY_OPEN = 0x0BDE
COLLISION_DISABLED = 0x0BD8
IMPULSE_X = 0x0BCA
IMPULSE_Y = 0x0BCE
IMPULSE_MODE = 0x0BD2


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pid", type=int, help="Linux PID of the TDKPIN winevdm.exe process")
    parser.add_argument(
        "--object-base",
        type=lambda value: int(value, 0),
        help="known live object base; otherwise locate it from the waiting ball",
    )
    parser.add_argument("--x", type=float, required=True, help="ball center x in logical pixels")
    parser.add_argument("--y", type=float, required=True, help="ball center y in logical pixels")
    parser.add_argument("--vx", type=float, default=0.0, help="x velocity in pixels per timer tick")
    parser.add_argument("--vy", type=float, default=0.0, help="y velocity in pixels per timer tick")
    parser.add_argument(
        "--substeps",
        type=int,
        choices=range(1, 6),
        help="temporarily force the original physics substep count",
    )
    parser.add_argument(
        "--stop-process",
        action="store_true",
        help="stop all Wine threads during writes; unsafe for repeated Win16 mutex probes",
    )
    return parser.parse_args()


def process_is_stopped(pid: int) -> bool:
    with open(f"/proc/{pid}/status", encoding="ascii") as stream:
        for line in stream:
            if line.startswith("State:"):
                return line.split()[1] in {"T", "t"}
    raise RuntimeError(f"cannot find process state for PID {pid}")


def stop_process(pid: int) -> bool:
    already_stopped = process_is_stopped(pid)
    if already_stopped:
        return False
    os.kill(pid, signal.SIGSTOP)
    deadline = time.monotonic() + 1.0
    while time.monotonic() < deadline:
        if process_is_stopped(pid):
            return True
        time.sleep(0.005)
    raise RuntimeError(f"PID {pid} did not stop")


def pack_millipixels(value: float) -> bytes:
    millipixels = round(value * 1000.0)
    if not -(2**31) <= millipixels < 2**31:
        raise RuntimeError(f"fixed-point value is out of range: {value}")
    return struct.pack("<i", millipixels)


def main() -> None:
    args = parse_args()
    verify_original()
    pid = args.pid if args.pid is not None else find_winevdm_pid()
    resume_afterward = stop_process(pid) if args.stop_process else False
    try:
        with ProcessMemory(pid, writable=True) as memory:
            data_base = locate_data_segment(memory)
            data = memory.read(data_base, DATA_READ_SIZE)
            object_base = (
                args.object_base if args.object_base is not None else locate_game_object(memory, data)
            )
            x = pack_millipixels(args.x)
            y = pack_millipixels(args.y)
            velocity_x = pack_millipixels(args.vx)
            velocity_y = pack_millipixels(args.vy)
            next_x = pack_millipixels(args.x + args.vx)
            next_y = pack_millipixels(args.y + args.vy)

            for offset, value in (
                (BALL_X, x),
                (BALL_Y, y),
                (NEXT_BALL_X, next_x),
                (NEXT_BALL_Y, next_y),
            ):
                memory.write(data_base + offset, value)
            if args.substeps is not None:
                memory.write(data_base + PHYSICS_SUBSTEPS, struct.pack("<H", args.substeps))
            for offset, value in (
                (VELOCITY_X, velocity_x),
                (VELOCITY_Y, velocity_y),
                (PREVIOUS_BALL_X, x),
                (PREVIOUS_BALL_Y, y),
                (OBJECT_BALL_X, x),
                (OBJECT_BALL_Y, y),
            ):
                memory.write(object_base + offset, value)

            memory.write(object_base + IMPULSE_X, bytes(4))
            memory.write(object_base + IMPULSE_Y, bytes(4))
            memory.write(object_base + IMPULSE_MODE, bytes(2))
            memory.write(object_base + COLLISION_DISABLED, bytes(1))
            memory.write(object_base + BALL_ACTIVE, bytes([1]))
            memory.write(object_base + PLAYER_ENTRY_OPEN, bytes(1))
            memory.write(object_base + BALL_SUSPENDED, bytes(1))
            print(
                f"TDKPIN pid={pid} data={data_base:#x} object={object_base:#x} "
                f"ball=({args.x:.3f},{args.y:.3f}) velocity=({args.vx:.3f},{args.vy:.3f}) "
                f"spin={int.from_bytes(memory.read(object_base + SPIN_REAL48, 6), 'little'):#014x}"
            )
    finally:
        if resume_afterward:
            os.kill(pid, signal.SIGCONT)


if __name__ == "__main__":
    try:
        main()
    except (OSError, RuntimeError) as error:
        raise SystemExit(f"inject_original_state.py: {error}") from error