Early return login response

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)

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: