conqp.rcon/rcon/battleye/proto.py

"""Low-level protocol stuff."""

from typing import NamedTuple, Union
from zlib import crc32


__all__ = [
    'RESPONSE_TYPES',
    'Header',
    'LoginRequest',
    'LoginResponse',
    'Command',
    'CommandResponse',
    'ServerMessage',
    'Request',
    'Response'
]


PREFIX = 'BE'
INFIX = 0xff


class Header(NamedTuple):
    """Packet header."""

    crc32: int
    type: int

    def __bytes__(self):
        return b''.join((
            PREFIX.encode('ascii'),
            self.crc32.to_bytes(4, 'little'),
            INFIX.to_bytes(1, 'little'),
            self.type.to_bytes(1, 'little')
        ))

    @classmethod
    def create(cls, typ: int, payload: bytes):
        """Creates a header for the given payload."""
        return cls(
            crc32(b''.join((
                INFIX.to_bytes(1, 'little'),
                typ.to_bytes(1, 'little'),
                payload
            ))),
            typ
        )

    @classmethod
    def from_bytes(cls, payload: bytes):
        """Creates a header from the given bytes."""
        if (prefix := payload[:2].decode('ascii')) != PREFIX:
            raise ValueError('Invalid prefix', prefix)

        if (infix := int.from_bytes(payload[5:6], 'little')) != INFIX:
            raise ValueError('Invalid infix', infix)

        return cls(
            int.from_bytes(payload[2:5], 'little'),
            int.from_bytes(payload[6:7], 'little')
        )


class LoginRequest(str):
    """Login request packet."""

    def __bytes__(self):
        return bytes(self.header) + self.payload

    @property
    def payload(self) -> bytes:
        """Returns the payload."""
        return self.encode('ascii')

    @property
    def header(self) -> Header:
        """Returns the appropriate header."""
        return Header.create(0x00, self.payload)


class LoginResponse(NamedTuple):
    """A login response."""

    header: Header
    success: bool

    @classmethod
    def from_bytes(cls, header: Header, payload: bytes):
        """Creates a login response from the given bytes."""
        return cls(header, bool(int.from_bytes(payload[:1], 'little')))


class Command(NamedTuple):
    """Command packet."""

    seq: int
    command: str

    def __bytes__(self):
        return bytes(self.header) + self.payload

    @property
    def payload(self) -> bytes:
        """Returns the payload."""
        return b''.join((
            self.seq.to_bytes(1, 'little'),
            self.command.encode('ascii')
        ))

    @property
    def header(self) -> Header:
        """Returns the appropriate header."""
        return Header.create(0x01, self.payload)

    @classmethod
    def from_string(cls, command: str):
        """Creates a command packet from the given string."""
        return cls(0x00, command)

    @classmethod
    def from_command(cls, command: str, *args: str):
        """Creates a command packet from the command and arguments."""
        return cls.from_string(' '.join([command, *args]))


class CommandResponse(NamedTuple):
    """A command response."""

    header: Header
    seq: int
    payload: bytes

    @classmethod
    def from_bytes(cls, header: Header, payload: bytes):
        """Creates a command response from the given bytes."""
        return cls(
            header,
            int.from_bytes(payload[:1], 'little'),
            payload[1:]
        )

    @property
    def message(self) -> str:
        """Returns the text message."""
        return self.payload.decode('ascii')


class ServerMessage(NamedTuple):
    """A message from the server."""

    header: Header
    seq: int
    payload: bytes

    @classmethod
    def from_bytes(cls, header: Header, payload: bytes):
        """Creates a server message from the given bytes."""
        return cls(
            header,
            int.from_bytes(payload[:1], 'little'),
            payload[1:]
        )

    @property
    def message(self) -> str:
        """Returns the text message."""
        return self.payload.decode('ascii')


Request = Union[LoginRequest, Command]
Response = Union[LoginResponse, CommandResponse, ServerMessage]

RESPONSE_TYPES = {
    0x00: LoginResponse,
    0x01: CommandResponse,
    0x02: ServerMessage
}