diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py index d2eb77a..dc5de47 100644 --- a/rcon/battleye/client.py +++ b/rcon/battleye/client.py @@ -5,11 +5,13 @@ from logging import getLogger from socket import SOCK_DGRAM from typing import Callable +from rcon.battleye.proto import HEADER_SIZE from rcon.battleye.proto import RESPONSE_TYPES from rcon.battleye.proto import CommandRequest 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 @@ -56,10 +58,10 @@ class Client(BaseClient, socket_type=SOCK_DGRAM): return RESPONSE_TYPES[ ( header := Header.from_bytes( - (data := self._socket.recv(self.max_length))[:8] + (data := self._socket.recv(self.max_length))[:HEADER_SIZE] ) ).type - ].from_bytes(header, data[8:]) + ].from_bytes(header, data[HEADER_SIZE:]) def communicate(self, request: Request) -> Response | str: """Send a request and receive a response.""" @@ -72,11 +74,11 @@ class Client(BaseClient, socket_type=SOCK_DGRAM): while True: response = self.receive() - try: - seq = response.seq - except AttributeError: + if isinstance(response, LoginResponse): return response + seq = response.seq + if isinstance(response, CommandResponse): # Collect fragmented command responses with the same seq command_responses.append(response) diff --git a/rcon/battleye/proto.py b/rcon/battleye/proto.py index 2f7c078..6414745 100644 --- a/rcon/battleye/proto.py +++ b/rcon/battleye/proto.py @@ -6,6 +6,7 @@ from zlib import crc32 __all__ = [ + "HEADER_SIZE", "RESPONSE_TYPES", "Header", "LoginRequest", @@ -19,6 +20,7 @@ __all__ = [ ] +HEADER_SIZE = 8 PREFIX = "BE" INFIX = 0xFF @@ -54,7 +56,7 @@ class Header(NamedTuple): @classmethod def from_bytes(cls, payload: bytes) -> Header: """Create a header from the given bytes.""" - if (size := len(payload)) != 8: + if (size := len(payload)) != HEADER_SIZE: raise ValueError("Invalid payload size", size) if (prefix := payload[:2].decode("ascii")) != PREFIX: