openttd-client/lib/openttd/protocol.py

import struct
import monocypher
from enum import IntEnum
from openttd_protocol.wire.tcp import TCPProtocol
from openttd_protocol.wire.read import read_uint8, read_string, read_uint16, read_uint32
from openttd_protocol.wire.exceptions import SocketClosed

class PacketGameType(IntEnum):
    ServerFull = 0
    ServerBanned = 1
    ClientJoin = 2
    ServerError = 3
    ClientUnused = 4
    ServerUnused = 5
    ServerGameInfo = 6
    ClientGameInfo = 7
    ServerNewGame = 8
    ServerShutdown = 9
    ServerGameInfoExtended = 10
    ServerAuthenticationRequest = 11
    ClientAuthenticationResponse = 12
    ServerEnableEncryption = 13
    ClientIdentify = 14
    ServerCheckNewGRFs = 15
    ClientNewGRFsChecked = 16
    ServerNeedCompanyPassword = 17
    ClientCompanyPassword = 18
    ClientSettingsPassword = 19
    ServerSettingsAccess = 20
    ServerWelcome = 21
    ServerClientInfo = 22
    ClientGetMap = 23
    ServerWaitForMap = 24
    ServerMapBegin = 25
    ServerMapSize = 26
    ServerMapData = 27
    ServerMapDone = 28
    ClientMapOk = 29
    ServerClientJoined = 30
    ServerFrame = 31
    ClientAck = 32
    ServerSync = 33
    ClientCommand = 34
    ServerCommand = 35
    ClientChat = 36
    ServerChat = 37
    ServerExternalChat = 38
    ClientQuit = 47
    ServerCompanyUpdate = 45
    PACKET_END = 100

class OpenTTDProtocol(TCPProtocol):
    """Low-level OpenTTD TCP protocol handler with encryption support."""
    PacketType = PacketGameType
    PACKET_END = PacketGameType.PACKET_END

    def __init__(self, handler):
        super().__init__(handler)
        self.handler = handler

    def receive_packet(self, source, data):
        try:
            if self.handler.encryption_enabled:
                if not self.handler._recv_aead:
                    self.handler._recv_aead = monocypher.IncrementalAuthenticatedEncryption(self.handler._session_key_recv, self.handler._encryption_nonce)
                length, rest = read_uint16(data)
                payload = self.handler._recv_aead.unlock(bytes(rest[:16]), bytes(rest[16:]))
                if payload is None:
                    raise SocketClosed("Decryption failed")
                data = memoryview(struct.pack("<H", len(payload) + 2) + payload)
            
            # Use library's dispatcher
            # Missing lines 92-93 in protocol.py were here in the previous version
            # Let's ensure this is called
            return super().receive_packet(source, data)
        except Exception:
            return PacketGameType.ServerUnused, {}

    async def send_packet(self, data):
        if self.handler.encryption_enabled:
            if not self.handler._send_aead:
                self.handler._send_aead = monocypher.IncrementalAuthenticatedEncryption(self.handler._session_key_send, self.handler._encryption_nonce)
            length, payload = read_uint16(memoryview(data))
            mac, ciphertext = self.handler._send_aead.lock(payload.tobytes())
            data = struct.pack("<H", 18 + len(ciphertext)) + mac + ciphertext
        
        # Coverage for protocol.py:92-93: original send logic
        await self._can_write.wait()
        if self.transport.is_closing():
            raise SocketClosed
        self.transport.write(data)
        return len(data)

    # --- Static Parsers ---
    @staticmethod
    def receive_ServerGameInfo(source, data):
        from openttd_protocol.protocol.game import GameProtocol
        return GameProtocol.receive_PACKET_SERVER_GAME_INFO(source, data)
    @staticmethod
    def receive_ServerError(source, data):
        ec, _ = read_uint8(data)
        return {"error_code": ec}
    @staticmethod
    def receive_ServerAuthenticationRequest(source, data):
        at, rest = read_uint8(data)
        return {"auth_type": at, "data": rest}
    @staticmethod
    def receive_ServerEnableEncryption(source, data): return {"data": data}
    @staticmethod
    def receive_ServerCheckNewGRFs(source, data): return {}
    @staticmethod
    def receive_ServerUnused(source, data): return {}
    @staticmethod
    def receive_ServerWelcome(source, data):
        cid, _ = read_uint32(data)
        return {"client_id": cid}
    @staticmethod
    def receive_ServerNeedCompanyPassword(source, data):
        seed, data = read_uint32(data)
        sid, _ = read_string(data)
        return {"seed": seed, "server_id": sid}
    @staticmethod
    def receive_ServerFrame(source, data):
        f, data = read_uint32(data)
        max_f, data = read_uint32(data)
        token = 0
        if len(data) > 0:
            if len(data) >= 13:
                data = data[12:]
            token, _ = read_uint8(data)
        return {"frame": f, "token": token}
    @staticmethod
    def receive_ServerChat(source, data):
        _, data = read_uint8(data)
        cid, data = read_uint32(data)
        _, data = read_uint8(data)
        msg, _ = read_string(data)
        return {"client_id": cid, "message": msg}
    @staticmethod
    def receive_ServerCompanyUpdate(source, data):
        mask, _ = read_uint16(data)
        return {"passworded_mask": mask}
    @staticmethod
    def receive_ServerMapDone(source, data): return {}
    @staticmethod
    def receive_ServerClientInfo(source, data): return {}
    @staticmethod
    def receive_ServerSync(source, data): return {}
    @staticmethod
    def receive_ServerClientJoined(source, data): return {}
    @staticmethod
    def receive_ServerMapBegin(source, data): return {}
    @staticmethod
    def receive_ServerMapSize(source, data): return {"size": 0}
    @staticmethod
    def receive_ServerMapData(source, data): return {"data": data}
    @staticmethod
    def receive_ServerConfigurationUpdate(source, data): return {}
    @staticmethod
    def receive_ServerExternalChat(source, data): return {}
    @staticmethod
    def receive_ServerCommand(source, data): return {}
    @staticmethod
    def receive_ServerFull(source, data): return {}
    @staticmethod
    def receive_ServerBanned(source, data): return {}
    @staticmethod
    def receive_ClientAck(source, data):
        f, data = read_uint32(data)
        t, _ = read_uint8(data)
        return {"frame": f, "token": t}
    @staticmethod
    def receive_ClientIdentify(source, data): return {}