Add server message handling

While at it, rewrite response reading.

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

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
+}