neoLegacy/tools/performance-monitor/bots.py

"""
Bot player manager for server load testing.

Each bot connects to the game server as a real client using the same
connection flow as the server-monitor tool (PreLogin -> Login -> cipher
handshake -> identity tokens -> keepalive).
"""

import logging
import os
import secrets
import socket
import struct
import sys
import threading
import time
from dataclasses import dataclass

from PySide6.QtCore import QObject, Signal

# Import protocol code from the sibling server-monitor tool.
_PARENT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_SM_PATH = os.path.join(_PARENT, "server-monitor")
if _SM_PATH not in sys.path:
    sys.path.append(_SM_PATH)

from protocol import frame_packet, DataInputStream
from packets import (
    PRE_LOGIN, LOGIN, KEEP_ALIVE, CUSTOM_PAYLOAD, DISCONNECT,
    build_prelogin, build_login, build_keep_alive,
    build_custom_payload, parse_prelogin_response,
    parse_login_response_stream, DISCONNECT_REASONS,
)

logger = logging.getLogger("perfmon")

# Reuse the same constants/patterns as server-monitor/connection.py
SMALLID_REJECT = 0xFF

CIPHER_KEY_PREFIX = (
    b"\xFA\x00\x07"
    b"\x00\x4D\x00\x43\x00\x7C\x00\x43\x00\x4B\x00\x65\x00\x79"
    b"\x00\x20"
)
CIPHER_KEY_TOTAL = len(CIPHER_KEY_PREFIX) + 32

CIPHER_ON_PATTERN = (
    b"\xFA\x00\x06"
    b"\x00\x4D\x00\x43\x00\x7C\x00\x43\x00\x4F\x00\x6E"
    b"\x00\x00"
)

CIPHER_ACK_CHANNEL = "MC|CAck"
IDENTITY_TOKEN_ISSUE = "MC|CTIssue"
IDENTITY_TOKEN_CHALLENGE = "MC|CTChallenge"
IDENTITY_TOKEN_RESPONSE = "MC|CTResponse"

PHASE_HANDSHAKE = 0
PHASE_CIPHER_SCAN = 1
PHASE_MONITORING = 2


def _build_channel_pattern(channel: str) -> bytes:
    result = b"\xFA"
    result += struct.pack(">h", len(channel))
    for ch in channel:
        result += struct.pack(">H", ord(ch))
    return result


class CipherState:
    def __init__(self, key: bytes, iv: bytes):
        from Crypto.Cipher import AES
        from Crypto.Util import Counter
        ctr = Counter.new(128, initial_value=int.from_bytes(iv, "big"))
        self._cipher = AES.new(key, AES.MODE_CTR, counter=ctr)

    def process(self, data: bytes) -> bytes:
        return self._cipher.encrypt(data)


@dataclass
class BotState:
    name: str
    thread: threading.Thread | None = None
    sock: socket.socket | None = None
    running: bool = False
    connected: bool = False


class BotConnection:
    """Single bot connection -- mirrors ServerConnection from server-monitor."""

    def __init__(self, bot: BotState, host: str, port: int, xuid: int,
                 log_fn, on_connected_fn):
        self._bot = bot
        self._host = host
        self._port = port
        self._xuid = xuid
        self._log = log_fn
        self._on_connected = on_connected_fn
        self._sock: socket.socket | None = None
        self._recv_buf = bytearray()
        self._scan_buf = bytearray()
        self._phase = PHASE_HANDSHAKE
        self._got_prelogin = False
        self._got_login = False
        self._recv_cipher: CipherState | None = None
        self._send_cipher: CipherState | None = None
        self._cipher_key = b""
        self._cipher_iv = b""
        self._cipher_scan_start = 0.0
        self._identity_token = b""

    def run(self):
        self._log(f"Bot '{self._bot.name}' connecting to {self._host}:{self._port}...")
        self._sock = socket.create_connection((self._host, self._port), timeout=10)
        self._bot.sock = self._sock
        self._sock.settimeout(5.0)

        # Phase 1: SmallID assignment
        first_byte = self._recv_exact(1)
        if first_byte[0] == SMALLID_REJECT:
            reject_data = self._recv_exact(5)
            reason_id = struct.unpack(">i", reject_data[1:5])[0]
            reason = DISCONNECT_REASONS.get(reason_id, f"Unknown({reason_id})")
            self._log(f"Bot '{self._bot.name}' rejected: {reason}")
            return

        self._log(f"Bot '{self._bot.name}' assigned smallId={first_byte[0]}")

        # Send PreLogin
        self._send_packet(PRE_LOGIN, build_prelogin(self._bot.name))
        self._log(f"Bot '{self._bot.name}' sent PreLogin")

        # Main recv loop
        self._sock.settimeout(1.0)
        last_keepalive = time.time()
        keepalive_counter = 0

        while self._bot.running:
            try:
                chunk = self._sock.recv(65536)
            except socket.timeout:
                # Cipher scan timeout
                if (self._phase == PHASE_CIPHER_SCAN
                        and not self._cipher_key
                        and time.time() - self._cipher_scan_start > 3.0):
                    self._handle_cipher_scan(b"")
                # Keepalive
                now = time.time()
                if now - last_keepalive >= 10.0:
                    keepalive_counter += 1
                    self._send_packet(KEEP_ALIVE, build_keep_alive(keepalive_counter))
                    last_keepalive = now
                continue
            except OSError:
                break

            if not chunk:
                self._log(f"Bot '{self._bot.name}' server closed connection")
                break

            self._recv_buf.extend(chunk)
            self._process_frames()

            now = time.time()
            if now - last_keepalive >= 10.0:
                keepalive_counter += 1
                self._send_packet(KEEP_ALIVE, build_keep_alive(keepalive_counter))
                last_keepalive = now

        self._log(f"Bot '{self._bot.name}' disconnected")

    def _process_frames(self):
        while len(self._recv_buf) >= 4:
            payload_len = struct.unpack(">I", self._recv_buf[:4])[0]
            total = 4 + payload_len
            if len(self._recv_buf) < total:
                break

            raw_payload = bytes(self._recv_buf[4:total])
            del self._recv_buf[:total]

            if self._recv_cipher:
                raw_payload = self._recv_cipher.process(raw_payload)

            if not raw_payload:
                continue

            if self._phase == PHASE_HANDSHAKE:
                self._handle_handshake(raw_payload)
            elif self._phase == PHASE_CIPHER_SCAN:
                self._handle_cipher_scan(raw_payload)
            elif self._phase == PHASE_MONITORING:
                self._handle_monitoring(raw_payload)

    def _handle_handshake(self, payload: bytes):
        if not payload:
            return
        packet_id = payload[0]
        data = payload[1:]

        if packet_id == PRE_LOGIN and not self._got_prelogin:
            try:
                parsed = parse_prelogin_response(data)
                self._got_prelogin = True
                self._log(f"Bot '{self._bot.name}' PreLogin OK: {parsed['player_count']} players online")
                self._send_packet(LOGIN, build_login(self._bot.name, self._xuid))
                self._log(f"Bot '{self._bot.name}' sent Login")
            except Exception as e:
                self._log(f"Bot '{self._bot.name}' PreLogin parse error: {e}")

        elif packet_id == LOGIN and not self._got_login:
            try:
                dis = DataInputStream(data)
                parsed = parse_login_response_stream(dis)
                self._got_login = True
                self._log(f"Bot '{self._bot.name}' Login OK: entityId={parsed['entity_id']}")
                self._phase = PHASE_CIPHER_SCAN
                self._cipher_scan_start = time.time()
                # Feed remaining bytes from this frame into cipher scan
                remaining = data[len(data) - dis.remaining:]
                if remaining:
                    self._handle_cipher_scan(remaining)
            except Exception as e:
                self._log(f"Bot '{self._bot.name}' Login parse error: {e}")

        elif packet_id == DISCONNECT:
            try:
                dis = DataInputStream(data)
                reason_id = dis.read_int()
                reason = DISCONNECT_REASONS.get(reason_id, f"Unknown({reason_id})")
                self._log(f"Bot '{self._bot.name}' kicked: {reason}")
            except Exception:
                self._log(f"Bot '{self._bot.name}' kicked (unknown reason)")
            self._bot.running = False

    def _handle_cipher_scan(self, payload: bytes):
        self._scan_buf.extend(payload)

        # Timeout: no cipher key after 3s = server has cipher disabled
        if (not self._cipher_key
                and time.time() - self._cipher_scan_start > 3.0):
            self._log(f"Bot '{self._bot.name}' no cipher, proceeding unencrypted")
            self._phase = PHASE_MONITORING
            self._bot.connected = True
            self._on_connected()
            buf = bytes(self._scan_buf)
            self._scan_buf.clear()
            if buf:
                self._handle_monitoring(buf)
            return

        # Look for MC|CKey
        if not self._cipher_key:
            idx = self._scan_buf.find(CIPHER_KEY_PREFIX)
            if idx >= 0 and idx + CIPHER_KEY_TOTAL <= len(self._scan_buf):
                key_start = idx + len(CIPHER_KEY_PREFIX)
                key_data = bytes(self._scan_buf[key_start:key_start + 32])
                self._cipher_key = key_data[:16]
                self._cipher_iv = key_data[16:32]
                self._log(f"Bot '{self._bot.name}' got cipher key")

                # Send MC|CAck and activate send cipher
                self._send_packet(CUSTOM_PAYLOAD, build_custom_payload(CIPHER_ACK_CHANNEL))
                iv_send = bytearray(self._cipher_iv)
                iv_send[0] ^= 0x80
                self._send_cipher = CipherState(self._cipher_key, bytes(iv_send))
                self._log(f"Bot '{self._bot.name}' send cipher active")

                del self._scan_buf[:idx + CIPHER_KEY_TOTAL]

        # Look for MC|COn
        if self._cipher_key and not self._recv_cipher:
            idx = self._scan_buf.find(CIPHER_ON_PATTERN)
            if idx >= 0:
                self._recv_cipher = CipherState(self._cipher_key, self._cipher_iv)
                self._log(f"Bot '{self._bot.name}' cipher handshake complete")
                self._phase = PHASE_MONITORING
                self._bot.connected = True
                self._on_connected()
                self._scan_buf.clear()

    def _handle_monitoring(self, payload: bytes):
        self._scan_buf.extend(payload)
        self._scan_identity_tokens()
        # Trim buffer
        if len(self._scan_buf) > 4096:
            del self._scan_buf[:-512]

    def _scan_identity_tokens(self):
        # MC|CTIssue: server sends a token for us to store
        issue_prefix = _build_channel_pattern(IDENTITY_TOKEN_ISSUE)
        idx = self._scan_buf.find(issue_prefix)
        if idx >= 0:
            data_start = idx + len(issue_prefix)
            if data_start + 2 + 32 <= len(self._scan_buf):
                data_len = struct.unpack(">h", self._scan_buf[data_start:data_start+2])[0]
                if data_len == 32:
                    token = bytes(self._scan_buf[data_start+2:data_start+2+32])
                    self._identity_token = token
                    self._log(f"Bot '{self._bot.name}' received identity token")
                    del self._scan_buf[:data_start + 2 + 32]
                    return

        # MC|CTChallenge: server wants us to present our token
        challenge_prefix = _build_channel_pattern(IDENTITY_TOKEN_CHALLENGE)
        idx = self._scan_buf.find(challenge_prefix)
        if idx >= 0:
            data_start = idx + len(challenge_prefix)
            if data_start + 2 <= len(self._scan_buf):
                del self._scan_buf[:data_start + 2]
                if self._identity_token and len(self._identity_token) == 32:
                    self._send_packet(CUSTOM_PAYLOAD,
                        build_custom_payload(IDENTITY_TOKEN_RESPONSE, self._identity_token))
                    self._log(f"Bot '{self._bot.name}' sent identity token response")
                else:
                    self._send_packet(CUSTOM_PAYLOAD,
                        build_custom_payload(IDENTITY_TOKEN_RESPONSE))
                    self._log(f"Bot '{self._bot.name}' sent empty identity token response")

    def _send_packet(self, packet_id: int, payload: bytes):
        raw = frame_packet(packet_id, payload)
        if self._send_cipher:
            header = raw[:4]
            encrypted = self._send_cipher.process(raw[4:])
            raw = header + encrypted
        try:
            if self._sock:
                self._sock.sendall(raw)
        except OSError:
            pass

    def _recv_exact(self, n: int) -> bytes:
        data = b""
        while len(data) < n:
            chunk = self._sock.recv(n - len(data))
            if not chunk:
                raise ConnectionError("Connection closed during recv")
            data += chunk
        return data


class BotManager(QObject):
    """Manages bot player connections for load testing."""

    bot_added = Signal(str)
    bot_removed = Signal(str)
    bot_error = Signal(str, str)
    log = Signal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self._bots: dict[str, BotState] = {}
        self._lock = threading.Lock()
        self._next_id = 1

    @property
    def bot_count(self) -> int:
        with self._lock:
            return len(self._bots)

    def add_bot(self, host: str, port: int) -> str:
        with self._lock:
            name = f"Bot{self._next_id}"
            self._next_id += 1
            xuid = secrets.randbits(62) | (1 << 32)

        bot = BotState(name=name)
        bot.running = True
        bot.thread = threading.Thread(
            target=self._run_bot, args=(bot, host, port, xuid),
            daemon=True)

        with self._lock:
            self._bots[name] = bot

        bot.thread.start()
        return name

    def remove_bot(self) -> str | None:
        with self._lock:
            if not self._bots:
                return None
            name = list(self._bots.keys())[-1]
            bot = self._bots.pop(name)

        bot.running = False
        if bot.sock:
            try:
                bot.sock.close()
            except OSError:
                pass
        self.bot_removed.emit(name)
        self.log.emit(f"Bot '{name}' disconnecting")
        return name

    def remove_all(self):
        with self._lock:
            names = list(self._bots.keys())
        for _ in names:
            self.remove_bot()

    def _run_bot(self, bot: BotState, host: str, port: int, xuid: int):
        try:
            def on_connected():
                self.bot_added.emit(bot.name)

            conn = BotConnection(bot, host, port, xuid, self.log.emit, on_connected)
            conn.run()
        except Exception as e:
            self.bot_error.emit(bot.name, str(e))
            self.log.emit(f"Bot '{bot.name}' error: {e}")
        finally:
            bot.running = False
            bot.connected = False
            if bot.sock:
                try:
                    bot.sock.close()
                except OSError:
                    pass
                bot.sock = None
            with self._lock:
                self._bots.pop(bot.name, None)