#!/usr/bin/env python3
# Copyright (C) 2026 Matrox Graphics Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
IPMX USB (TR-10-14) PCAP-Driven Sender Simulator.

Reads any IPMX USB PCAP file (generated by generate_usb_test_pcap.py or
usb_pcap_to_ipmx.py) and replays the Sender side of the protocol so that
ipmx_usb_tester.py can connect to it for end-to-end testing without real
hardware.

How it works
------------
1. Parses the PCAP with usbDissector.analyze_pcap to extract Sender and
   Receiver messages for every channel.
2. Listens on the control port.  When a Receiver (tester) connects:
   a. Sends SenderConnectionInfo (from the PCAP).
   b. Receives SenderConnectionStatus → learns the Receiver's data port.
3. Connects to the Receiver's data port for each USB device in the PCAP.
   a. Sends USBStreamInfo.
   b. Receives USBStreamStatus.
4. Main loop — responds to every Submit from the tester:
   - CONTROL_SUBMIT with seqnum N → send the recorded CONTROL_SUBMIT_RETURN
     for seqnum N (matched from the PCAP).  Falls back to a generic SUCCESS
     response if the seqnum is not in the recording.
   - INTERRUPT_SUBMIT → send the next queued INTERRUPT_SUBMIT_RETURN from the
     PCAP (in recording order), cycling back to idle HID reports when exhausted.
   - BULK_SUBMIT / CANCEL_SUBMIT → generic SUCCESS Return.
5. Sends HEARTBEAT messages on the control channel at the negotiated interval.
6. Closes cleanly on Ctrl-C or when the tester disconnects.

Usage
-----
    # In terminal 1 — start the simulator from any IPMX PCAP:
    python3 pcap_sender_sim.py test_pcaps/scenario_12_realistic_typing.pcap

    # In terminal 2 — connect the live tester to it:
    python3 ipmx_usb_tester.py 127.0.0.1 --verbose --duration 30

    # With capture for later analysis:
    python3 ipmx_usb_tester.py 127.0.0.1 --verbose --capture live.pcap

    # Real PCAP sources:
    python3 pcap_sender_sim.py test_pcaps/real_keyboard.pcap
"""

from __future__ import annotations

import argparse
import socket
import struct
import sys
import threading
import time
from collections import deque
from typing import Optional

import ipmx_usb_message as usb
from ipmx_usb_message import IpmxUsbMessage, MsgType

# Import the PCAP parser / session model from the dissector
from usbDissector import (
    analyze_pcap, Session, ChannelMessage, DataChannel,
    _heartbeat_period,
)

from ipmx_pep import (
    PepParams, PepMode, IvMode,
    derive_privacy_key, compute_iv_prime,
    encrypt_usb_message, decrypt_usb_message,
    usb_ctr_advance,
)


# ---------------------------------------------------------------------------
# PEP encryption state for the simulator (Sender role)
# ---------------------------------------------------------------------------

class SimPepState:
    """Manages PEP encrypt/decrypt for one direction of one channel.

    The Sender sim:
      - Encrypts its outgoing messages with the S2R key (even SSID).
      - Decrypts incoming Receiver messages with the R2S key (odd SSID).

    Handshake SSID: first message pair on each data channel uses SSID 2/3,
    then switches to the payload SSID from USBStreamInfo.
    """

    def __init__(self, psk: bytes, mode: PepMode = PepMode.AES_128_CTR_CMAC_64_AAD,
                 sender_kv: int = 1, corrupt_nth: int = 0,
                 base_iv: int = 0, key_generator: bytes = b'\x00' * 16):
        self.psk = psk
        self.mode = mode
        self.sender_kv = sender_kv
        self.receiver_kv = 0

        self.key_generator = key_generator
        kv_bytes = sender_kv.to_bytes(4, 'big')
        self.s2r_key = derive_privacy_key(
            psk, self.key_generator, kv_bytes, key_bits=128)
        self.r2s_key = b'\x00' * 16  # derived after we learn receiver_kv

        self.base_iv: int = base_iv
        self.ctr_s2r: int = 0
        self.ctr_r2s: int = 0
        self.handshake_done = False
        self.payload_ssid: int = 0
        self.corrupt_nth: int = corrupt_nth
        self._s2r_msg_count: int = 0

    def derive_r2s_key(self, receiver_kv: int) -> None:
        self.receiver_kv = receiver_kv
        kv_bytes = receiver_kv.to_bytes(4, 'big')
        self.r2s_key = derive_privacy_key(
            self.psk, self.key_generator, kv_bytes, key_bits=128)

    def encrypt_s2r(self, raw: bytes) -> bytes:
        """Encrypt an outgoing Sender message (S2R direction).

        Patches CTR and KEYVERSION into the header before encrypting.
        """
        ssid = self.payload_ssid if self.handshake_done else 2
        iv_prime = compute_iv_prime(self.base_iv, ssid)

        hdr = bytearray(raw[:16])
        struct.pack_into('>Q', hdr, 0, self.ctr_s2r)
        struct.pack_into('>I', hdr, 8, self.sender_kv)
        patched = bytes(hdr) + raw[16:]

        enc = encrypt_usb_message(patched, self.s2r_key, iv_prime, self.mode)
        self.ctr_s2r += usb_ctr_advance(len(enc))
        self._s2r_msg_count += 1
        if (self.corrupt_nth > 0 and self.handshake_done and
                self._s2r_msg_count % self.corrupt_nth == 0):
            ba = bytearray(enc)
            ba[-1] ^= 0xFF
            enc = bytes(ba)
            print(f"  [sim pep] INJECTED MAC corruption on S2R msg #{self._s2r_msg_count}")
        return enc

    def decrypt_r2s(self, raw: bytes) -> bytes:
        """Decrypt an incoming Receiver message (R2S direction).

        On the first R2S message, reads key_version from the clear-text header
        and derives the R2S key dynamically (the Receiver picks a random kv).
        """
        if len(raw) >= 12 and self.receiver_kv == 0:
            incoming_kv = struct.unpack('>I', raw[8:12])[0]
            if incoming_kv != 0:
                self.derive_r2s_key(incoming_kv)
                print(f"  [sim pep] derived R2S key from receiver_kv={incoming_kv:#010x}")

        ssid = (self.payload_ssid | 1) if self.handshake_done else 3
        iv_prime = compute_iv_prime(self.base_iv, ssid)
        dec, mac_ok = decrypt_usb_message(raw, self.r2s_key, iv_prime, self.mode)
        self.ctr_r2s += usb_ctr_advance(len(raw))
        if not mac_ok:
            print(f"  [sim pep] WARNING: R2S MAC verification failed")
        return dec

    def finish_handshake(self, payload_ssid: int) -> None:
        self.payload_ssid = payload_ssid
        self.handshake_done = True


# ---------------------------------------------------------------------------
# Socket helpers
# ---------------------------------------------------------------------------

def _sendall(sock: socket.socket, data: bytes) -> None:
    try:
        sock.sendall(data)
    except OSError:
        pass


# Sentinel returned by _recv_msg when the peer closed the connection
# (as opposed to None which means the timeout expired with no message).
_DISCONNECTED = object()


def _recv_msg(sock: socket.socket, buf: bytearray,
              timeout: float = 10.0):
    """Block until one complete IPMX message is in *buf*, then parse it.

    Returns:
        IpmxUsbMessage  — a complete message was received.
        None            — timeout expired with no complete message.
        _DISCONNECTED   — the peer closed or reset the connection.
    """
    deadline = time.monotonic() + timeout
    while True:
        if len(buf) >= usb.HEADER_SIZE:
            length = usb.peek_length(bytes(buf))
            if length is not None and len(buf) >= length:
                raw = bytes(buf[:length])
                del buf[:length]
                return usb.parse_one(raw)
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return None
        sock.settimeout(min(1.0, remaining))
        try:
            chunk = sock.recv(65536)
        except socket.timeout:
            continue
        except OSError:
            return _DISCONNECTED
        if not chunk:          # TCP FIN — peer closed cleanly
            return _DISCONNECTED
        buf.extend(chunk)


# ---------------------------------------------------------------------------
# PCAP session loader
# ---------------------------------------------------------------------------

class SimSession:
    """
    Extracts Sender-side messages from a parsed IPMX PCAP session and
    organises them for replay.
    """

    def __init__(self, session: Session):
        self.session = session

        # Control channel: raw bytes of each Sender message, in order
        self.ctrl_msgs: list[bytes] = [
            cm.msg.raw for cm in session.control.messages
            if cm.src_ip == session.sender_ip
        ]

        # SenderConnectionInfo (first Sender control message)
        self.sci_raw: bytes = self.ctrl_msgs[0] if self.ctrl_msgs else b''

        # Heartbeat messages
        self.heartbeat_raw: bytes = usb.build_heartbeat()

        # Data channels: per substreamid maps
        # ctrl_returns[substreamid][seqnum] → raw SUBMIT_RETURN bytes
        # int_returns[substreamid]          → deque of raw INTERRUPT_SUBMIT_RETURN bytes
        # idle_int_return[substreamid]      → raw idle (all-zero) report Return
        self.ctrl_returns:    dict[int, dict[int, bytes]] = {}
        self.int_returns:     dict[int, deque] = {}
        self.bulk_returns:    dict[int, dict[int, bytes]] = {}
        self.iso_returns:     dict[int, dict[int, deque]] = {}
        self.stream_info_raw: dict[int, bytes] = {}

        for substreamid, ch in session.data_channels.items():
            self.stream_info_raw[substreamid] = b''
            self.ctrl_returns[substreamid]    = {}
            self.int_returns[substreamid]     = deque()
            self.bulk_returns[substreamid]    = {}
            self.iso_returns[substreamid]     = {}

            for cm in ch.messages:
                msg = cm.msg
                mt  = msg.msg_type_enum
                # After dissector decryption, is_encrypted is still True
                # (CTR/KV preserved) but payload fields are accessible.
                # Use the _encrypted sentinel to detect truly opaque messages.
                opaque = '_encrypted' in msg.payload

                if mt == MsgType.USB_STREAM_INFO:
                    self.stream_info_raw[substreamid] = msg.raw

                elif mt == MsgType.USB_CONTROL_SUBMIT_RETURN:
                    seqnum = msg.payload.get('seqnum', -1) if not opaque else -1
                    if seqnum >= 0:
                        self.ctrl_returns[substreamid][seqnum] = msg.raw

                elif mt == MsgType.USB_INTERRUPT_SUBMIT_RETURN:
                    self.int_returns[substreamid].append(msg.raw)

                elif mt == MsgType.USB_BULK_SUBMIT_RETURN:
                    seqnum = msg.payload.get('seqnum', -1) if not opaque else -1
                    if seqnum >= 0:
                        self.bulk_returns[substreamid][seqnum] = msg.raw

                elif mt == MsgType.USB_ISOCHRONOUS_SUBMIT_RETURN:
                    ep = msg.payload.get('endpoint', 0) if not opaque else 0
                    if ep not in self.iso_returns[substreamid]:
                        self.iso_returns[substreamid][ep] = deque()
                    self.iso_returns[substreamid][ep].append(msg.raw)

        print(f"  Loaded session: {len(session.data_channels)} data channel(s)")
        for sid, ch in session.data_channels.items():
            n_ctrl = len(self.ctrl_returns.get(sid, {}))
            n_int  = len(self.int_returns.get(sid, deque()))
            iso_ep_map = self.iso_returns.get(sid, {})
            n_iso = sum(len(q) for q in iso_ep_map.values())
            ep_detail = ", ".join(f"EP{e:#04x}={len(q)}" for e, q in sorted(iso_ep_map.items()))
            print(f"    SubstreamID=0x{sid:02X}: {n_ctrl} control returns, "
                  f"{n_int} interrupt returns, {n_iso} iso returns ({ep_detail})")


# ---------------------------------------------------------------------------
# Data channel handler (Sender side)
# ---------------------------------------------------------------------------

class DataChannelSender:
    """
    Connects to the Receiver's data port and replays one USB device's
    data stream from the PCAP.
    """

    def __init__(self, sim: SimSession, substreamid: int,
                 receiver_ip: str, receiver_data_port: int,
                 stop_event: threading.Event,
                 pep_state: Optional[SimPepState] = None):
        self._sim        = sim
        self._sid        = substreamid
        self._rip        = receiver_ip
        self._rport      = receiver_data_port
        self._stop       = stop_event
        self._pep        = pep_state

    def run(self) -> None:
        sid = self._sid
        time.sleep(0.05)   # brief pause before connecting

        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((self._rip, self._rport))
        except OSError as exc:
            print(f"  [sim data {sid:#04x}] connect failed: {exc}")
            return

        buf = bytearray()
        print(f"  [sim data {sid:#04x}] connected to Receiver {self._rip}:{self._rport}")

        # PEP helpers: wrap send/recv with encryption/decryption
        def pep_send(s: socket.socket, data: bytes) -> None:
            if self._pep:
                data = self._pep.encrypt_s2r(data)
            _sendall(s, data)

        def pep_recv(s: socket.socket, b: bytearray,
                     timeout: float = 10.0):
            msg = _recv_msg(s, b, timeout=timeout)
            if msg is _DISCONNECTED or msg is None:
                return msg
            if self._pep and msg.raw:
                dec = bytearray(self._pep.decrypt_r2s(msg.raw))
                dec[0:12] = b'\x00' * 12  # clear CTR/KV so parse_one decodes payload
                return usb.parse_one(bytes(dec))
            return msg

        # Send USBStreamInfo
        si_raw = self._sim.stream_info_raw.get(sid, b'')
        if not si_raw:
            print(f"  [sim data {sid:#04x}] no USBStreamInfo in PCAP — skipping")
            sock.close()
            return
        pep_send(sock, si_raw)
        print(f"  [sim data {sid:#04x}] sent USBStreamInfo (encrypted={self._pep is not None})")

        # Wait for USBStreamStatus
        msg = pep_recv(sock, buf, timeout=10.0)
        if msg is _DISCONNECTED:
            print(f"  [sim data {sid:#04x}] Receiver disconnected before USBStreamStatus")
            sock.close()
            return
        if msg is None or msg.msg_type_enum != MsgType.USB_STREAM_STATUS:
            print(f"  [sim data {sid:#04x}] expected USBStreamStatus, "
                  f"got {msg.msg_type_name if msg else 'None'}")
            sock.close()
            return
        print(f"  [sim data {sid:#04x}] received USBStreamStatus — stream active")

        # Switch from handshake SSID (2/3) to payload SSID
        if self._pep:
            self._pep.finish_handshake(sid)
            print(f"  [sim data {sid:#04x}] PEP handshake done, payload SSID={sid}")

        ctrl_returns  = self._sim.ctrl_returns.get(sid, {})
        ctrl_return_queue = deque(
            v for _, v in sorted(ctrl_returns.items())
        )
        int_returns   = deque(self._sim.int_returns.get(sid, deque()))
        bulk_returns  = self._sim.bulk_returns.get(sid, {})
        iso_ep_map: dict[int, deque] = {
            ep: deque(q) for ep, q in self._sim.iso_returns.get(sid, {}).items()
        }

        # Fallback idle interrupt return (all-zero HID report, seqnum patched)
        idle_report   = b'\x00' * 8

        # Track the next expected control seqnum for proper pairing
        _RETURN_TYPES = {
            MsgType.USB_CONTROL_SUBMIT,
            MsgType.USB_BULK_SUBMIT,
            MsgType.USB_INTERRUPT_SUBMIT,
            MsgType.USB_ISOCHRONOUS_SUBMIT,
            MsgType.USB_CANCEL_SUBMIT,
        }

        while not self._stop.is_set():
            msg = pep_recv(sock, buf, timeout=2.0)
            if msg is _DISCONNECTED:
                print(f"  [sim data {sid:#04x}] Receiver disconnected")
                break
            if msg is None:
                continue

            mt = msg.msg_type_enum
            if mt not in _RETURN_TYPES:
                continue

            seqnum = msg.payload.get('seqnum', -1) if not msg.is_encrypted else -1
            ep     = msg.payload.get('endpoint', 0) if not msg.is_encrypted else 0
            dirn   = msg.payload.get('direction', 0) if not msg.is_encrypted else 0

            if mt == MsgType.USB_CONTROL_SUBMIT:
                ret_raw = ctrl_return_queue.popleft() if ctrl_return_queue else None
                if ret_raw:
                    pcap_ret = usb.parse_one(ret_raw)
                    pcap_dirn = pcap_ret.payload.get('direction', dirn) if not pcap_ret.is_encrypted else dirn
                    if pcap_dirn != dirn:
                        pep_send(sock, usb.build_usb_control_submit_return(
                            seqnum, ep, dirn, 0, 0,
                        ))
                    else:
                        pcap_seq = pcap_ret.payload.get('seqnum', seqnum)
                        patched = _patch_seqnum(ret_raw, seqnum) if pcap_seq != seqnum else ret_raw
                        pep_send(sock, patched)
                else:
                    pep_send(sock, usb.build_usb_control_submit_return(
                        seqnum, ep, dirn, 0, 0,
                    ))

            elif mt == MsgType.USB_INTERRUPT_SUBMIT:
                if int_returns:
                    ret_raw = int_returns.popleft()
                    pcap_ret    = usb.parse_one(ret_raw)
                    actual_len  = pcap_ret.payload.get('actuallength', 0) if not pcap_ret.is_encrypted else 0
                    xfr_hex     = pcap_ret.payload.get('transferdata') if not pcap_ret.is_encrypted else None
                    xfr_data    = bytes.fromhex(xfr_hex) if xfr_hex else b''
                    pep_send(sock, usb.build_usb_interrupt_submit_return(
                        seqnum, ep, dirn, actual_len, 0, xfr_data,
                    ))
                else:
                    pep_send(sock, usb.build_usb_interrupt_submit_return(
                        seqnum, ep, dirn, len(idle_report), 0, idle_report,
                    ))

            elif mt == MsgType.USB_BULK_SUBMIT:
                ret_raw = bulk_returns.get(seqnum)
                if ret_raw:
                    pep_send(sock, _patch_seqnum(ret_raw, seqnum))
                else:
                    pep_send(sock, usb.build_usb_bulk_submit_return(
                        seqnum, ep, dirn, 0, 0,
                    ))

            elif mt == MsgType.USB_ISOCHRONOUS_SUBMIT:
                ep_queue = iso_ep_map.get(ep, deque())
                if ep_queue:
                    ret_raw = ep_queue.popleft()
                    pep_send(sock, _patch_seqnum(ret_raw, seqnum))
                else:
                    pep_send(sock, usb.build_usb_isochronous_submit_return(
                        seqnum, ep, dirn,
                        startframe=0, errorcount=0,
                        num_packets=0, iso_packets=[],
                    ))

            elif mt == MsgType.USB_CANCEL_SUBMIT:
                pep_send(sock, usb.build_usb_cancel_submit_return(
                    seqnum, ep, dirn, 0,
                ))

        sock.close()
        print(f"  [sim data {sid:#04x}] closed")


def _patch_seqnum(raw: bytes, new_seqnum: int) -> bytes:
    """Patch the SEQNUM field (first 3 bytes of DATA) in a submit-return message.

    TR-10-14 encodes SEQNUM as the lower 3 bytes of a big-endian 4-byte value:
      data[0..2] = struct.pack('>I', seqnum)[1:]
    Byte data[3] is the endpoint/direction byte and MUST NOT be touched.
    Using struct.pack_into with '>I' would overwrite 4 bytes and corrupt data[3].
    """
    if len(raw) < usb.HEADER_SIZE + 3:
        return raw
    out = bytearray(raw)
    seq3 = struct.pack('>I', new_seqnum & 0xFFFFFF)[1:]  # 3 bytes, big-endian
    out[usb.HEADER_SIZE:usb.HEADER_SIZE + 3] = seq3
    return bytes(out)


# ---------------------------------------------------------------------------
# Control channel + orchestration
# ---------------------------------------------------------------------------

class PcapSenderSim:
    """
    Listens on the control port, exchanges SenderConnectionInfo /
    SenderConnectionStatus, then connects data channels back to the
    Receiver.  Runs until Ctrl-C or the Receiver disconnects.
    """

    def __init__(self, sim: SimSession, ctrl_port: int,
                 psk: Optional[bytes] = None, corrupt_nth: int = 0,
                 pep_params: Optional[PepParams] = None):
        self._sim        = sim
        self._ctrl_port  = ctrl_port
        self._stop       = threading.Event()
        self._psk        = psk
        self._corrupt_nth = corrupt_nth
        self._pep_params = pep_params

    def run(self) -> None:
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind(("0.0.0.0", self._ctrl_port))
        srv.listen(4)
        srv.settimeout(1.0)
        print(f"Listening for Receiver on :{self._ctrl_port}  "
              f"({len(self._sim.session.data_channels)} data channel(s) to replay)")
        print("Press Ctrl-C to stop\n")

        while not self._stop.is_set():
            try:
                conn, addr = srv.accept()
            except socket.timeout:
                continue
            except OSError:
                break

            receiver_ip = addr[0]
            print(f"Receiver connected from {receiver_ip}:{addr[1]}")
            try:
                self._handle_receiver(conn, receiver_ip)
            except Exception as exc:
                print(f"  [sim ctrl] error: {exc}")
            finally:
                try:
                    conn.close()
                except Exception:
                    pass

        srv.close()

    def _handle_receiver(self, conn: socket.socket, receiver_ip: str) -> None:
        sim  = self._sim
        buf  = bytearray()

        # Set up PEP for the control channel (SSID 0 = S2R, 1 = R2S)
        ctrl_pep: Optional[SimPepState] = None
        if self._psk:
            pp = self._pep_params
            kv = int.from_bytes(pp.key_version, 'big') if pp else 1
            biv = pp.iv if pp else 0
            kg = pp.key_generator if pp else b'\x00' * 16
            ctrl_pep = SimPepState(self._psk, sender_kv=kv,
                                   corrupt_nth=self._corrupt_nth,
                                   base_iv=biv, key_generator=kg)
            ctrl_pep.payload_ssid = 0
            ctrl_pep.handshake_done = True  # control channel has no handshake phase

        def ctrl_send(data: bytes) -> None:
            if ctrl_pep:
                data = ctrl_pep.encrypt_s2r(data)
            _sendall(conn, data)

        def ctrl_recv(timeout: float = 15.0):
            msg = _recv_msg(conn, buf, timeout=timeout)
            if msg is _DISCONNECTED or msg is None:
                return msg
            if ctrl_pep and msg.raw:
                dec = bytearray(ctrl_pep.decrypt_r2s(msg.raw))
                dec[0:12] = b'\x00' * 12  # clear CTR/KV so parse_one decodes payload
                return usb.parse_one(bytes(dec))
            return msg

        # Step 1 — send SenderConnectionInfo
        sci = sim.sci_raw
        if not sci:
            print("  [sim ctrl] no SenderConnectionInfo in PCAP!")
            return
        ctrl_send(sci)
        print(f"  [sim ctrl] sent SenderConnectionInfo (encrypted={ctrl_pep is not None})")

        # Step 2 — receive SenderConnectionStatus to learn data port
        msg = ctrl_recv(timeout=15.0)
        if msg is _DISCONNECTED:
            print(f"  [sim ctrl] Receiver disconnected before SenderConnectionStatus")
            return
        if msg is None or msg.msg_type_enum != MsgType.SENDER_CONNECTION_STATUS:
            print(f"  [sim ctrl] expected SCS, got {msg.msg_type_name if msg else 'None'}")
            return

        data_port = msg.payload.get('port', 40000)
        hbeat_idx = msg.payload.get('hbeat', 10)
        print(f"  [sim ctrl] received SenderConnectionStatus  "
              f"data_port={data_port}  hbeat={hbeat_idx}")

        # Per-session stop: shared by all data threads for this connection.
        # Set when this Receiver disconnects so data threads exit too.
        session_stop = threading.Event()

        # Step 3 — spawn data channel connections
        data_threads = []
        for substreamid in sim.session.data_channels:
            data_pep: Optional[SimPepState] = None
            if self._psk:
                pp = self._pep_params
                kv = int.from_bytes(pp.key_version, 'big') if pp else 1
                biv = pp.iv if pp else 0
                kg = pp.key_generator if pp else b'\x00' * 16
                data_pep = SimPepState(self._psk, sender_kv=kv,
                                       corrupt_nth=self._corrupt_nth,
                                       base_iv=biv, key_generator=kg)
            t = threading.Thread(
                target=DataChannelSender(
                    sim, substreamid, receiver_ip, data_port, session_stop,
                    pep_state=data_pep,
                ).run,
                name=f"sim-data-{substreamid:#04x}",
                daemon=True,
            )
            t.start()
            data_threads.append(t)

        # Step 4 — control channel loop: send heartbeats, handle disconnect
        hbeat_period = _heartbeat_period(hbeat_idx)
        last_hbeat   = time.monotonic()
        print(f"  [sim ctrl] heartbeat period ≈ {hbeat_period:.0f} s  — ready "
              f"(Ctrl-C to stop, or reconnect after tester exits)")

        try:
            while not self._stop.is_set():
                if time.monotonic() - last_hbeat >= hbeat_period:
                    ctrl_send(sim.heartbeat_raw)
                    last_hbeat = time.monotonic()
                    print(f"  [sim ctrl] sent HEARTBEAT")

                msg = ctrl_recv(timeout=min(hbeat_period, 1.0))
                if msg is _DISCONNECTED:
                    print(f"  [sim ctrl] Receiver disconnected — waiting for next connection")
                    break
                if msg is None:
                    continue

                mt = msg.msg_type_enum
                if mt == MsgType.VENDOR_SPECIFIC_QUERY:
                    vqtype = msg.payload.get('vqtype', 255) if not msg.is_encrypted else 255
                    cid    = (bytes.fromhex(msg.payload.get('cid', '0050C2'))
                              if not msg.is_encrypted else b'\x00\x50\xC2')
                    ctrl_send(usb.build_vendor_specific_query_return(cid, vqtype, 255))
        finally:
            # Signal data threads to stop and wait briefly for them to wind down.
            session_stop.set()
            for t in data_threads:
                t.join(timeout=2.0)
            print(f"  [sim ctrl] session ended")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="IPMX USB TR-10-14 PCAP-driven Sender Simulator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Replay a synthetic scenario:
  python3 pcap_sender_sim.py test_pcaps/scenario_12_realistic_typing.pcap

  # Replay a converted real capture:
  python3 pcap_sender_sim.py test_pcaps/real_keyboard.pcap

  # Then in another terminal, connect the live tester:
  python3 ipmx_usb_tester.py 127.0.0.1 --verbose --duration 30
  python3 ipmx_usb_tester.py 127.0.0.1 --verbose --capture live_session.pcap
""",
    )
    parser.add_argument(
        'pcap_file',
        help="IPMX USB PCAP file to replay as Sender",
    )
    parser.add_argument(
        '--port', type=int, default=27502,
        metavar='PORT',
        help="Control port to listen on (default: 27502)",
    )
    parser.add_argument(
        '--psk', type=str, default=None,
        metavar='HEX',
        help="Pre-shared key (hex) for PEP encryption (16 or 32 bytes)",
    )
    parser.add_argument(
        '--pep-iv', type=str, default=None, metavar='HEX',
        help="PEP base IV (hex, 16 chars)",
    )
    parser.add_argument(
        '--pep-key-generator', type=str, default=None, metavar='HEX',
        help="PEP key_generator (hex, 32 chars)",
    )
    parser.add_argument(
        '--pep-key-version', type=str, default=None, metavar='HEX',
        help="PEP key_version (hex, 8 chars)",
    )
    parser.add_argument(
        '--corrupt-nth', type=int, default=0,
        metavar='N',
        help="Corrupt every Nth outgoing encrypted S2R message (0 = off)",
    )
    args = parser.parse_args()

    psk_bytes: Optional[bytes] = None
    pep_params: Optional[PepParams] = None
    if args.psk:
        psk_bytes = bytes.fromhex(args.psk)
        if len(psk_bytes) not in (16, 32):
            print("PSK must be 16 or 32 bytes (32 or 64 hex digits)", file=sys.stderr)
            sys.exit(1)
        pep_iv = int(args.pep_iv, 16) if args.pep_iv else 0
        pep_key_gen = bytes.fromhex(args.pep_key_generator) if args.pep_key_generator else b'\x00' * 16
        pep_kv = bytes.fromhex(args.pep_key_version) if args.pep_key_version else b'\x00\x00\x00\x01'
        pep_params = PepParams(
            psk=psk_bytes,
            mode=PepMode.AES_128_CTR_CMAC_64_AAD,
            iv=pep_iv,
            key_generator=pep_key_gen,
            key_version=pep_kv,
        )
        print(f"PEP enabled: PSK={args.psk[:8]}... ({len(psk_bytes)*8}-bit)  "
              f"iv={pep_iv:#018x}  kv={pep_kv.hex()}")

    print(f"Loading PCAP: {args.pcap_file}")
    try:
        sessions = analyze_pcap(args.pcap_file, verbose=False,
                                pep_params=pep_params)
    except SystemExit:
        sys.exit(1)

    if not sessions:
        print("No IPMX sessions found in PCAP", file=sys.stderr)
        sys.exit(1)

    # Use the first session
    session = sessions[0]
    sim     = SimSession(session)

    srv = PcapSenderSim(sim, ctrl_port=args.port, psk=psk_bytes,
                        corrupt_nth=args.corrupt_nth, pep_params=pep_params)

    try:
        srv.run()
    except KeyboardInterrupt:
        print("\nStopped by user")
        srv._stop.set()


if __name__ == '__main__':
    main()
