#!/usr/bin/env python3
"""run_decoder.py - real-time protocol decoder for Paralogue (Fallen Doll).

Attaches Frida to Paralogue-Win64-Shipping.exe, loads protocol_decoder.js
(a read-only onLeave hook on the encrypt/decrypt wrapper at RVA 0x49A2EB0),
reassembles the [2B length-prefix][PacketHead][body] framing, and then:

  * parses each PacketHead to obtain msg_id and body_len
  * identifies msg_id -> protobuf message-type by trying all 239 known types
    from paralogue_net_full_pb2 and scoring the parses
  * writes one JSON line per message to captures/session_<ts>.jsonl
  * prints a human-readable conversation log

Usage
    python run_decoder.py                    # decode only
    python run_decoder.py --redirect         # also load redirect_connect.js

Author: protocol research tool; instrumentation is read-only.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from dataclasses import dataclass, field
from typing import Optional

import frida

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

import paralogue_net_full_pb2 as pb_mod  # noqa: E402

from google.protobuf.message import DecodeError  # noqa: E402

PROCESS_NAME = "Paralogue-Win64-Shipping.exe"
DECODER_JS = os.path.join(SCRIPT_DIR, "protocol_decoder.js")
REDIRECT_JS = os.path.join(SCRIPT_DIR, "redirect_connect.js")
OUT_DIR = os.path.join(SCRIPT_DIR, "captures")
os.makedirs(OUT_DIR, exist_ok=True)
KEYSTREAM_PATH = os.path.join(OUT_DIR, "keystream.bin")

# User-editable known-good mappings discovered from research.
# Lines are loaded on startup and used as strong hints (any type here
# outranks heuristic matches). This file is rewritten when the decoder
# discovers new high-confidence mappings.
KNOWN_IDS_PATH = os.path.join(SCRIPT_DIR, "known_msg_ids.json")


# ---------------------------------------------------------------------------
# protobuf catalog + msg_id identification
# ---------------------------------------------------------------------------

def _build_catalog():
    """Return (list of (name, descriptor, message_class)) for all types."""
    fd = pb_mod.DESCRIPTOR
    out = []
    for name, desc in fd.message_types_by_name.items():
        cls = getattr(pb_mod, name, None)
        if cls is None:
            continue
        out.append((name, desc, cls))
    return out


CATALOG = _build_catalog()
PACKET_HEAD_CLS = pb_mod.PacketHead


def _score_fit(msg, body: bytes, desc) -> int:
    """Score how well `msg` (parsed from `body`) fits the descriptor.

    Strategy:
    * require at least one field set (wire bytes are not empty yet type has
      no matching fields -> protobuf stored them as unknown -> score 0)
    * re-serialize; exact byte-equality = huge bonus (canonical round-trip)
    * length-equality = smaller bonus
    * validate that declared `string` fields are valid UTF-8 (protobuf's
      pure-python decoder does not always enforce this; the C++ decoder does,
      but we check anyway for cross-platform consistency)
    * penalise large unknown-field byte counts
    """
    try:
        set_fields = msg.ListFields()
    except Exception:
        return 0
    if not set_fields:
        return 0

    score = len(set_fields) * 2

    # 1. Round-trip check
    try:
        rt = msg.SerializeToString()
    except Exception:
        rt = b""
    if rt == body:
        score += 100  # canonical fit -- almost certainly the right type
    elif rt and len(rt) == len(body):
        score += 20   # same size but different bytes (ordering differences)
    elif rt:
        # how close is the length?
        ratio = min(len(rt), len(body)) / max(len(rt), len(body))
        score += int(ratio * 10)

    # 2. Unknown fields: wire bytes we couldn't assign to any declared field.
    # If serializing loses lots of bytes, the type is a poor fit.
    if rt:
        missing = max(0, len(body) - len(rt))
        score -= missing // 4

    # 3. String validation: for every field declared `string`, require UTF-8.
    TYPE_STRING = 9   # FieldDescriptor.TYPE_STRING
    TYPE_BYTES = 12
    NUMERIC_TYPES = {
        1, 2,        # double, float
        3, 4,        # int64, uint64
        5,           # int32
        6, 7,        # fixed64, fixed32
        13, 14,      # uint32, enum
        15, 16,      # sfixed32, sfixed64
        17, 18,      # sint32, sint64
    }
    LABEL_REPEATED = 3

    for f, v in set_fields:
        if f.type == TYPE_STRING:
            vals = v if f.label == LABEL_REPEATED else [v]
            for s in vals:
                if not isinstance(s, str):
                    continue
                try:
                    s.encode("utf-8")
                except Exception:
                    score -= 50
                    continue
                if s and sum(1 for c in s if c.isprintable()) / len(s) >= 0.8:
                    score += 5  # strongly string-like
                elif s and sum(1 for c in s if c.isprintable()) / len(s) < 0.5:
                    score -= 10

        # Repeated-numeric field whose values all look like ASCII? That's
        # a packed-int parse of what was really a string. Penalise hard.
        if f.label == LABEL_REPEATED and f.type in NUMERIC_TYPES and isinstance(v, (list, tuple)) and v:
            try:
                ints = [int(x) for x in v]
            except Exception:
                ints = []
            if ints and all(32 <= x <= 126 for x in ints) and len(ints) >= 3:
                score -= 40

    return score


def identify_body(body: bytes, preferred_name: Optional[str] = None):
    """Try to parse `body` against every known message type.

    Returns (best_name, best_msg, best_score, candidates, top) where
    `top` is a list of (name, score) for the top 3 candidates (for the log).
    """
    if not body:
        return None, None, 0, 0, []

    scored: list[tuple[int, str, object]] = []

    for name, desc, cls in CATALOG:
        try:
            m = cls()
            m.MergeFromString(body)
        except (DecodeError, ValueError, Exception):
            continue
        score = _score_fit(m, body, desc)
        if score > 0:
            scored.append((score, name, m))

    if not scored:
        return None, None, 0, 0, []

    # preferred (cached) name wins ties
    scored.sort(key=lambda t: (-t[0], 0 if t[1] == preferred_name else 1, t[1]))
    if preferred_name:
        for s, n, m in scored:
            if n == preferred_name:
                top = [(nn, ss) for ss, nn, _ in scored[:3]]
                return n, m, s, len(scored), top

    best_score, best_name, best_msg = scored[0]
    top = [(nn, ss) for ss, nn, _ in scored[:3]]
    return best_name, best_msg, best_score, len(scored), top


def msg_to_dict(msg) -> dict:
    """Protobuf message -> plain python dict (best-effort, handles bytes)."""
    from google.protobuf import json_format
    try:
        return json_format.MessageToDict(
            msg,
            preserving_proto_field_name=True,
            including_default_value_fields=False,
        )
    except Exception:
        # Fallback: field-by-field
        out = {}
        for f, v in msg.ListFields():
            if isinstance(v, bytes):
                out[f.name] = v.hex()
            else:
                out[f.name] = repr(v)
        return out


# ---------------------------------------------------------------------------
# per-stream framing reassembler
# ---------------------------------------------------------------------------

@dataclass
class Stream:
    """One reassembled byte stream (one per (this_id, direction))."""

    this_id: str
    direction: str  # "C->S" or "S->C"
    chunks: dict = field(default_factory=dict)  # offset -> bytes
    consumed: int = 0  # bytes already emitted as decoded messages
    frontier: int = 0  # highest contiguous byte available
    _cache_bytes: bytearray = field(default_factory=bytearray)

    def add(self, off: int, data: bytes) -> None:
        # On the very first chunk, initialise consumed/frontier to its offset
        # so we don't scan through thousands of zero-fill bytes from offset 0.
        if not self.chunks:
            self.consumed = off
            self.frontier = off
        # Grow the cache to cover the end of this chunk, zero-filled.
        end = off + len(data)
        if end > len(self._cache_bytes):
            self._cache_bytes.extend(b"\x00" * (end - len(self._cache_bytes)))
        self._cache_bytes[off:end] = data
        self.chunks[off] = len(data)
        # Update contiguous frontier: walk forward while chunks cover it.
        f = self.frontier
        while True:
            # Find chunk that starts at or before f and extends past f.
            advanced = False
            for co, cl in self.chunks.items():
                if co <= f < co + cl:
                    f = co + cl
                    advanced = True
                    break
            if not advanced:
                break
        self.frontier = f

    def available(self) -> int:
        return self.frontier - self.consumed

    def peek(self, n: int) -> bytes:
        return bytes(self._cache_bytes[self.consumed:self.consumed + n])

    def consume(self, n: int) -> bytes:
        out = bytes(self._cache_bytes[self.consumed:self.consumed + n])
        self.consumed += n
        return out


# ---------------------------------------------------------------------------
# framing decoder
# ---------------------------------------------------------------------------

def _load_known_ids() -> dict:
    if os.path.isfile(KNOWN_IDS_PATH):
        try:
            with open(KNOWN_IDS_PATH, "r", encoding="utf-8") as f:
                raw = json.load(f)
            # keys are ints (but JSON stores them as strings)
            return {int(k): v for k, v in raw.items()}
        except Exception:
            return {}
    return {}


def _save_known_ids(m: dict) -> None:
    try:
        with open(KNOWN_IDS_PATH, "w", encoding="utf-8") as f:
            json.dump({str(k): v for k, v in sorted(m.items())}, f, indent=2)
    except Exception:
        pass


class Decoder:
    """Pull messages out of a reassembled Stream.

    Framing (from RESEARCH.md):
        [2B BE length-prefix N][N bytes PacketHead protobuf][body_len bytes body]

    The N-byte blob is parseable as a PacketHead and gives msg_id + body_len.
    """

    # score threshold above which we consider an identification strong enough
    # to auto-promote into the persistent known_msg_ids.json cache
    STRONG_SCORE = 60

    def __init__(self, log_fp, seq_start: int = 1):
        self.log_fp = log_fp
        self.seq = seq_start
        self.msg_id_cache: dict[int, str] = _load_known_ids()
        self._user_supplied_ids = set(self.msg_id_cache.keys())

    def process_chunk(self, direction: str, off: int, data: bytes) -> None:
        """Each chunk is a self-contained packet.

        Strategy:
          1. Try to parse the leading bytes (4-10 bytes) as a PacketHead; if
             body_len matches the chunk's total size (TOTAL-bytes semantics)
             or the chunk has just a header-worth of bytes, emit.
          2. If leading-bytes parse fails, skip a byte and retry up to 4 times
             (handles a leading length-prefix or envelope byte).
          3. Otherwise emit an "unparsed" record so the chunk isn't dropped.
        """
        n = len(data)

        # 1. Try to find a clean PacketHead at a small offset.
        for skip in range(0, min(4, n - 3)):
            for hlen in range(4, min(16, n - skip + 1)):
                hb = data[skip:skip + hlen]
                head = PACKET_HEAD_CLS()
                try:
                    head.MergeFromString(hb)
                except DecodeError:
                    continue
                try:
                    rt = head.SerializeToString()
                except Exception:
                    continue
                if len(rt) != hlen:
                    continue
                mid = int(head.msg_id)
                blen = int(head.body_len)
                if mid <= 0 or mid > 100000:
                    continue
                # Accept when body_len equals chunk size (TOTAL semantics)
                # or when the post-header bytes == body_len (body-only).
                total_match = (blen == n - skip)
                bodyonly_match = (skip + hlen + blen == n)
                if not (total_match or bodyonly_match):
                    continue
                body = (data[skip + hlen:skip + blen] if total_match
                        else data[skip + hlen:skip + hlen + blen])
                self._emit_packet(direction, off + skip, head, body, skip)
                return

        # 2. Fallback: emit as raw/unparsed
        self._emit_raw(direction, off, data)

    def _emit_raw(self, direction: str, off: int, data: bytes) -> None:
        record = {
            "seq": self.seq,
            "direction": direction,
            "stream_offset": off,
            "msg_id": None,
            "msg_type": "<raw-chunk>",
            "body_len": len(data),
            "fields": {"_raw_hex": data.hex()},
        }
        self.seq += 1
        self.log_fp.write(json.dumps(record, ensure_ascii=False) + "\n")
        self.log_fp.flush()
        # Don't spam stdout for short control-plane chunks.

    def _emit_packet(self, direction: str, off: int, head, body: bytes, skipped: int) -> None:
        mid = int(head.msg_id)
        preferred = self.msg_id_cache.get(mid)
        name, parsed, score, cand, top = identify_body(body, preferred)

        if name and not preferred and score >= self.STRONG_SCORE:
            self.msg_id_cache[mid] = name
            if mid not in self._user_supplied_ids:
                _save_known_ids(self.msg_id_cache)

        if parsed is not None:
            fields_obj = msg_to_dict(parsed)
        else:
            fields_obj = {"_raw_hex": body.hex()}

        record = {
            "seq": self.seq,
            "direction": direction,
            "stream_offset": off,
            "msg_id": mid,
            "msg_type": name or "<unknown>",
            "packet_id": int(head.packet_id),
            "flag": int(head.flag),
            "body_len": int(head.body_len),
            "seq_id": int(head.seq_id),
            "total_pack_count": int(head.total_pack_count),
            "identify_score": score,
            "identify_candidates": cand,
            "top_candidates": [{"type": n, "score": s} for n, s in top],
            "skipped_bytes": skipped,
            "fields": fields_obj,
        }
        self.seq += 1

        self.log_fp.write(json.dumps(record, ensure_ascii=False) + "\n")
        self.log_fp.flush()

        short = _short_repr(fields_obj)
        print(
            f"[{record['seq']:04d}] {direction} "
            f"msg_id={mid:<5d} "
            f"{record['msg_type']:<28s} "
            f"body={record['body_len']:>5d}B  {short}"
        )

    def drain(self, stream: Stream) -> int:
        """Emit as many complete messages as possible.

        Framing (confirmed against live plaintext captures):
          [2B BE prefix N][N-byte PacketHead][body_len bytes body]

        Where body_len is the PacketHead.body_len field (body-only size).
        The PacketHead-length prefix lets us parse the header unambiguously.
        We also tolerate a stray pattern where the prefix is replaced by a
        sentinel `00 00` ; in that case we advance one byte and retry.
        """
        emitted = 0
        while True:
            avail = stream.available()
            if avail < 2:
                return emitted
            prefix = stream.peek(2)
            hlen = (prefix[0] << 8) | prefix[1]
            if hlen == 0 or hlen > 32:
                # not a sane header length -- skip one byte and retry
                stream.consume(1)
                continue
            if avail < 2 + hlen:
                return emitted
            head_bytes = bytes(
                stream._cache_bytes[stream.consumed + 2:
                                    stream.consumed + 2 + hlen]
            )
            head = PACKET_HEAD_CLS()
            try:
                head.MergeFromString(head_bytes)
            except DecodeError:
                stream.consume(1)
                continue
            try:
                rt = head.SerializeToString()
            except Exception:
                stream.consume(1)
                continue
            if len(rt) != hlen:
                # header parsed but canonical length differs -- resync
                stream.consume(1)
                continue
            mid = int(head.msg_id)
            blen = int(head.body_len)
            if mid <= 0 or mid > 100000 or blen < 0 or blen > (1 << 20):
                stream.consume(1)
                continue
            if avail < 2 + hlen + blen:
                return emitted

            stream.consume(2 + hlen)
            body = stream.consume(blen)
            self._emit(stream, head, body, skipped=0)
            emitted += 1
        # (unreachable) -- kept to keep the legacy scan below inert
        return emitted  # type: ignore[unreachable]

    def _unused_scan_path(self, stream: Stream) -> int:
        emitted = 0
        while True:
            avail = stream.available()
            if avail < 5:
                return emitted

            base = stream.consumed
            window = bytes(stream._cache_bytes[base:base + min(avail, 4096)])

            # Score every viable (pos, hlen) candidate and pick the best.
            best = None   # (score, pos, hlen, head, body)
            for pos in range(0, min(len(window) - 4, 16)):
                for hlen in range(5, 24):
                    if pos + hlen > len(window):
                        break
                    hb = window[pos:pos + hlen]
                    head = PACKET_HEAD_CLS()
                    try:
                        head.MergeFromString(hb)
                    except DecodeError:
                        continue
                    try:
                        rt = head.SerializeToString()
                    except Exception:
                        continue
                    if len(rt) != hlen:
                        continue
                    mid = int(head.msg_id)
                    blen = int(head.body_len)
                    if mid <= 0 or mid > 100000:
                        continue
                    if blen < 0 or blen > (1 << 20):
                        continue
                    # `body_len` appears to be body-only for login-phase
                    # packets but total-size for ping keepalives. Prefer
                    # body-only (matches RESEARCH.md observations on login)
                    # and fall back to total-size when body-only exceeds
                    # what's available in the current stream window.
                    bodyonly_ok = pos + hlen + blen <= len(window)
                    total_ok = blen >= hlen and pos + blen <= len(window)
                    if bodyonly_ok:
                        body = window[pos + hlen:pos + hlen + blen]
                        consumed_len = pos + hlen + blen
                    elif total_ok:
                        body = window[pos + hlen:pos + blen]
                        consumed_len = pos + blen
                    else:
                        continue

                    # Score: prefer headers that explicitly set body_len and
                    # have more set fields; prefer low `pos` (earlier match);
                    # prefer non-empty bodies.
                    score = len(head.ListFields()) * 10
                    if head.HasField("body_len") if False else (blen > 0 or hlen >= 7):
                        score += 20
                    if blen > 0:
                        score += 30
                    score -= pos * 2
                    if best is None or score > best[0]:
                        best = (score, pos, hlen, head, body, consumed_len)
            if best is not None:
                _score, pos, hlen, head, body, consumed_len = best

            if best is None:
                # no packet found in the lookahead; if the window is full,
                # drop the first byte and try again to resync
                if avail >= 4096:
                    stream.consume(1)
                    continue
                return emitted

            stream.consume(consumed_len)
            self._emit(stream, head, body, skipped=pos)
            emitted += 1

    def _emit(self, stream: Stream, head, body: bytes, skipped: int = 0) -> None:
        mid = int(head.msg_id)
        preferred = self.msg_id_cache.get(mid)
        name, parsed, score, cand, top = identify_body(body, preferred)

        # Promote high-confidence first-time identifications into the cache.
        if name and not preferred and score >= self.STRONG_SCORE:
            self.msg_id_cache[mid] = name
            if mid not in self._user_supplied_ids:
                _save_known_ids(self.msg_id_cache)

        if parsed is not None:
            fields_obj = msg_to_dict(parsed)
        else:
            fields_obj = {"_raw_hex": body.hex()}

        record = {
            "seq": self.seq,
            "direction": stream.direction,
            "stream_offset": stream.consumed - len(body),
            "msg_id": mid,
            "msg_type": name or "<unknown>",
            "packet_id": int(head.packet_id),
            "flag": int(head.flag),
            "body_len": int(head.body_len),
            "seq_id": int(head.seq_id),
            "total_pack_count": int(head.total_pack_count),
            "identify_score": score,
            "identify_candidates": cand,
            "top_candidates": [{"type": n, "score": s} for n, s in top],
            "skipped_bytes": skipped,
            "fields": fields_obj,
        }
        self.seq += 1

        # JSONL for machine consumption
        self.log_fp.write(json.dumps(record, ensure_ascii=False) + "\n")
        self.log_fp.flush()

        # Human-readable line
        short_fields = _short_repr(fields_obj)
        print(
            f"[{record['seq']:04d}] {record['direction']} "
            f"msg_id={record['msg_id']:<5d} "
            f"{record['msg_type']:<28s} "
            f"body={record['body_len']:>5d}B  {short_fields}"
        )


def _short_repr(d) -> str:
    if isinstance(d, dict):
        parts = []
        for k, v in list(d.items())[:4]:
            if isinstance(v, (dict, list)):
                parts.append(f"{k}=<{type(v).__name__}>")
            else:
                s = str(v)
                if len(s) > 40:
                    s = s[:37] + "..."
                parts.append(f"{k}={s}")
        extra = "" if len(d) <= 4 else f" (+{len(d)-4} more)"
        return "{" + ", ".join(parts) + "}" + extra
    return str(d)


# ---------------------------------------------------------------------------
# session manager
# ---------------------------------------------------------------------------

class Session:
    def __init__(self, log_path: str, raw_log_path: Optional[str] = None):
        self.log_path = log_path
        self.log_fp = open(log_path, "w", encoding="utf-8")
        self._raw_fp = open(raw_log_path, "w", encoding="utf-8") if raw_log_path else None
        self.streams: dict[tuple[str, int], Stream] = {}
        self.decoder = Decoder(self.log_fp)
        self.total_chunks = 0

    def close(self):
        for fp in (self.log_fp, self._raw_fp):
            if fp is None:
                continue
            try:
                fp.close()
            except Exception:
                pass

    def on_chunk(self, meta: dict, data: bytes) -> None:
        self.total_chunks += 1
        this_id = meta["this_id"]
        enc = int(meta.get("enc", 0))
        off = int(meta.get("off", 0))

        direction = "C->S" if enc == 1 else "S->C"

        if self._raw_fp is not None:
            rec = {
                "n": self.total_chunks,
                "dir": direction,
                "this": this_id[-8:],
                "off": off,
                "len": len(data),
                "hex": data.hex(),
            }
            self._raw_fp.write(json.dumps(rec) + "\n")
            self._raw_fp.flush()

        # Chunks are fragments of a larger plaintext stream. Accumulate by
        # (this_id, enc_flag) and drain full messages as framing allows.
        key = (this_id, enc)
        stream = self.streams.get(key)
        if stream is None:
            stream = Stream(this_id=this_id, direction=direction)
            self.streams[key] = stream
        stream.add(off, data)
        self.decoder.drain(stream)


# ---------------------------------------------------------------------------
# frida plumbing
# ---------------------------------------------------------------------------

def attach_and_run(redirect: bool, raw: bool = False) -> None:
    ts = int(time.time())
    log_path = os.path.join(OUT_DIR, f"session_{ts}.jsonl")
    raw_path = os.path.join(OUT_DIR, f"session_{ts}.raw.jsonl") if raw else None
    session = Session(log_path, raw_path)

    print("=" * 64)
    print(f"  Paralogue protocol decoder")
    print(f"  log: {log_path}")
    print("=" * 64)

    with open(DECODER_JS, "r", encoding="utf-8") as f:
        decoder_js = f.read()

    redirect_js = None
    if redirect:
        with open(REDIRECT_JS, "r", encoding="utf-8") as f:
            redirect_js = f.read()

    print(f"[*] waiting for {PROCESS_NAME} ...")
    frida_session = None
    while frida_session is None:
        try:
            for proc in frida.get_local_device().enumerate_processes():
                if proc.name == PROCESS_NAME:
                    print(f"[*] attaching to PID {proc.pid}")
                    frida_session = frida.attach(proc.pid)
                    break
        except Exception as e:
            print(f"[!] {e}")
        if frida_session is None:
            time.sleep(1)

    scripts = []
    if redirect_js:
        s = frida_session.create_script(redirect_js)
        s.on("message", lambda msg, data: _on_redirect_msg(msg))
        s.load()
        scripts.append(s)
        print("[*] redirect hook loaded")

    decoder = frida_session.create_script(decoder_js)

    def on_message(message, data):
        if message.get("type") == "error":
            print(f"[!] js error: {message.get('description','')}")
            return
        payload = message.get("payload")
        if not isinstance(payload, dict):
            return
        t = payload.get("t")
        if t == "ready":
            print(f"[*] {payload.get('msg','')}")
        elif t == "status":
            print(f"[*] hook calls so far: {payload.get('calls',0)} | "
                  f"chunks processed: {session.total_chunks} | "
                  f"mapped msg_ids: {len(session.decoder.msg_id_cache)}")
        elif t == "error":
            print(f"[!] {payload.get('msg','')}")
        elif t == "chunk":
            if data is not None:
                session.on_chunk(payload, bytes(data))
        elif t == "keystream":
            if data is not None:
                _write_keystream(int(payload.get("off", 0)), bytes(data))

    decoder.on("message", on_message)
    decoder.load()
    scripts.append(decoder)
    print("[*] decoder hook loaded. Connect to a server in-game.")
    print("[*] Ctrl+C to stop.\n")

    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass
    finally:
        print("\n[*] detaching...")
        try:
            frida_session.detach()
        except Exception:
            pass
        # Summary
        print(f"\n[=] session log: {log_path}")
        print(f"[=] chunks processed: {session.total_chunks}")
        print(f"[=] messages decoded: {session.decoder.seq - 1}")
        if session.decoder.msg_id_cache:
            print(f"[=] msg_id -> type mapping discovered ({len(session.decoder.msg_id_cache)}):")
            for mid, name in sorted(session.decoder.msg_id_cache.items()):
                print(f"       {mid:<6d}  {name}")
        session.close()


_KS_INIT_DONE = False


def _write_keystream(off: int, data: bytes) -> None:
    """Write keystream bytes at the game's absolute stream offset.

    The proxy and the in-process wrapper count the same encrypted payload
    byte stream (post-handshake), so absolute offsets align between them.
    Truncates the file on first call so stale cross-session bytes don't
    poison the proxy.
    """
    global _KS_INIT_DONE
    if not _KS_INIT_DONE:
        try:
            open(KEYSTREAM_PATH, "wb").close()
        except Exception:
            pass
        _KS_INIT_DONE = True
    try:
        with open(KEYSTREAM_PATH, "r+b") as f:
            f.seek(off)
            f.write(data)
    except FileNotFoundError:
        with open(KEYSTREAM_PATH, "wb") as f:
            f.seek(off)
            f.write(data)


def _on_redirect_msg(msg):
    if msg.get("type") == "error":
        print(f"[!] redirect js error: {msg.get('description','')}")
        return
    p = msg.get("payload") or {}
    t = p.get("t")
    if t == "ready":
        print(f"[*] redirect: {p.get('msg','')}")
    elif t == "redirect":
        print(f"[*] redirect: {p.get('from','?')}:{p.get('port','?')} -> {p.get('to','?')}")
    elif t == "error":
        print(f"[!] redirect: {p.get('msg','')}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Paralogue protocol decoder")
    ap.add_argument(
        "--redirect",
        action="store_true",
        help="also load redirect_connect.js to redirect port 11001 to 127.0.0.1",
    )
    ap.add_argument(
        "--raw",
        action="store_true",
        help="also write session_<ts>.raw.jsonl with every chunk's hex payload",
    )
    args = ap.parse_args()

    attach_and_run(redirect=args.redirect, raw=args.raw)


if __name__ == "__main__":
    main()