diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py index 9042724..6015555 100644 --- a/rcon/battleye/client.py +++ b/rcon/battleye/client.py @@ -1,14 +1,17 @@ """BattlEye RCon client.""" from ipaddress import IPv4Address +from logging import getLogger from socket import SOCK_DGRAM -from typing import Union +from typing import Callable, Union +from rcon.battleye.proto import RESPONSE_TYPES from rcon.battleye.proto import Command -from rcon.battleye.proto import CommandResponse from rcon.battleye.proto import Header from rcon.battleye.proto import LoginRequest -from rcon.battleye.proto import LoginResponse +from rcon.battleye.proto import Request +from rcon.battleye.proto import Response +from rcon.battleye.proto import ServerMessage from rcon.client import BaseClient from rcon.exceptions import WrongPassword @@ -17,41 +20,53 @@ __all__ = ['Client'] Host = Union[str, IPv4Address] +MessageHandler = Callable[[ServerMessage], None] + + +def log_message(server_message: ServerMessage) -> None: + """Default handler, logging the server message.""" + + getLogger('Server message').info(server_message.message) class Client(BaseClient, socket_type=SOCK_DGRAM): """BattlEye RCon client.""" - def send(self, data: bytes) -> None: - """Sends bytes.""" - with self._socket.makefile('wb') as file: - file.write(data) + def __init__( + self, *args, + message_handler: MessageHandler = log_message, + **kwargs + ): + super().__init__(*args, **kwargs) + self._handle_server_message = message_handler + + def _receive(self, max_length: int) -> Response: + """Receives a packet.""" + data = self._socket.recv(max_length)[:7] + header = Header.from_bytes(data) + return RESPONSE_TYPES[header.type].from_bytes(header, data[7:]) + + def receive(self, max_length: int = 4096) -> Response: + """Receives a message.""" + while isinstance(response := self._receive(max_length), ServerMessage): + self._handle_server_message(response) - def _login(self, login_request: LoginRequest) -> LoginResponse: + return response + + def communicate(self, request: Request) -> Response: """Logs the user in.""" - self.send(bytes(login_request)) + with self._socket.makefile('wb') as file: + file.write(bytes(request)) - with self._socket.makefile('rb') as file: - return LoginResponse.read(file) + return self.receive() def login(self, passwd: str) -> bool: """Logs the user in.""" - if not self._login(LoginRequest.from_passwd(passwd)).success: + if not self.communicate(LoginRequest.from_passwd(passwd)).success: raise WrongPassword() return True - def _run(self, command: Command) -> CommandResponse: - """Executes a command.""" - self.send(bytes(command)) - - with self._socket.makefile('rb') as file: - header = Header.read(file) - - # TODO: Can we determine the packet size? - remainder = self._socket.recv(4096) - return CommandResponse.from_bytes(header, remainder) - def run(self, command: str, *args: str) -> str: """Executes a command.""" - return self._run(Command.from_command(command, *args)).message + return self.communicate(Command.from_command(command, *args)).message diff --git a/rcon/battleye/proto.py b/rcon/battleye/proto.py index bff3598..90f0265 100644 --- a/rcon/battleye/proto.py +++ b/rcon/battleye/proto.py @@ -1,15 +1,19 @@ """Low-level protocol stuff.""" -from typing import IO, NamedTuple +from typing import NamedTuple, Union from zlib import crc32 __all__ = [ + 'RESPONSE_TYPES', 'Header', 'LoginRequest', 'LoginResponse', 'Command', - 'CommandResponse' + 'CommandResponse', + 'ServerMessage', + 'Request', + 'Response' ] @@ -28,7 +32,7 @@ class Header(NamedTuple): return b''.join(( self.prefix.encode('ascii'), self.crc32.to_bytes(4, 'little'), - self.suffix.to_bytes(1, 'big') + self.suffix.to_bytes(1, 'little') )) @classmethod @@ -41,19 +45,14 @@ class Header(NamedTuple): ) @classmethod - def from_bytes(cls, prefix: bytes, crc32sum: bytes, suffix: bytes): + def from_bytes(cls, payload: bytes): """Creates a header from the given bytes.""" return cls( - prefix.decode('ascii'), - int.from_bytes(crc32sum, 'little'), - int.from_bytes(suffix, 'big') + payload[:2].decode('ascii'), + int.from_bytes(payload[2:5], 'little'), + int.from_bytes(payload[5:6], 'little') ) - @classmethod - def read(cls, file: IO): - """Reads the packet from a file-like object.""" - return cls.from_bytes(file.read(2), file.read(4), file.read(1)) - class LoginRequest(NamedTuple): """Login request packet.""" @@ -88,19 +87,14 @@ class LoginResponse(NamedTuple): success: bool @classmethod - def from_bytes(cls, header: Header, typ: bytes, success: bytes): + def from_bytes(cls, header: Header, payload: bytes): """Creates a login response from the given bytes.""" return cls( header, - int.from_bytes(typ, 'little'), - bool(int.from_bytes(success, 'little')) + int.from_bytes(payload[:1], 'little'), + bool(int.from_bytes(payload[2:3], 'little')) ) - @classmethod - def read(cls, file: IO): - """Reads a login response from a file-like object.""" - return cls.from_bytes(Header.read(file), file.read(1), file.read(1)) - class Command(NamedTuple): """Command packet.""" @@ -146,16 +140,50 @@ class CommandResponse(NamedTuple): payload: bytes @classmethod - def from_bytes(cls, header: Header, data: bytes): + def from_bytes(cls, header: Header, payload: bytes): """Creates a command response from the given bytes.""" return cls( header, - int.from_bytes(data[:1], 'little'), - int.from_bytes(data[1:2], 'little'), - data[2:] + int.from_bytes(payload[:1], 'little'), + int.from_bytes(payload[1:2], 'little'), + payload[2:] + ) + + @property + def message(self) -> str: + """Returns the text message.""" + return self.payload.decode('ascii') + + +class ServerMessage(NamedTuple): + """A message from the server.""" + + header: Header + type: int + 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'), + int.from_bytes(payload[1:2], 'little'), + payload[2:] ) @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 +}