From 87f3c197876fdeb8b2177a8e372ba5fbbe9ec426 Mon Sep 17 00:00:00 2001
From: Richard Neumann <r.neumann@homeinfo.de>
Date: Mon, 14 Feb 2022 10:22:43 +0100
Subject: [PATCH] Improve BattlEye RCon client

---
 rcon/battleye/client.py | 45 +++++++++++++++++++++++++--------
 rcon/battleye/proto.py  | 56 +++++++++++++++++++++++++++++++++++++++--
 2 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py
index 1dab69a..9042724 100644
--- a/rcon/battleye/client.py
+++ b/rcon/battleye/client.py
@@ -4,8 +4,13 @@ from ipaddress import IPv4Address
 from socket import SOCK_DGRAM
 from typing import Union
 
-from rcon.battleye.proto import Command, LoginRequest
+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.client import BaseClient
+from rcon.exceptions import WrongPassword
 
 
 __all__ = ['Client']
@@ -17,18 +22,36 @@ Host = Union[str, IPv4Address]
 class Client(BaseClient, socket_type=SOCK_DGRAM):
     """BattlEye RCon client."""
 
-    def communicate(self, data: bytes, *, read: int = 4096) -> bytes:
-        """Sends and receives packets."""
-        self._socket.send(data)
-        return self._socket.recv(read)
+    def send(self, data: bytes) -> None:
+        """Sends bytes."""
+        with self._socket.makefile('wb') as file:
+            file.write(data)
 
-    def login(self, passwd: str) -> bytes:
+    def _login(self, login_request: LoginRequest) -> LoginResponse:
         """Logs the user in."""
-        return self.communicate(bytes(LoginRequest.from_passwd(passwd)))
+        self.send(bytes(login_request))
+
+        with self._socket.makefile('rb') as file:
+            return LoginResponse.read(file)
+
+    def login(self, passwd: str) -> bool:
+        """Logs the user in."""
+        if not self._login(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."""
-        packet = Command.from_command(command, *args)
-        _ = self.communicate(bytes(packet))
-        # TODO: Process response
-        return ''
+        return self._run(Command.from_command(command, *args)).message
diff --git a/rcon/battleye/proto.py b/rcon/battleye/proto.py
index 8189dea..bff3598 100644
--- a/rcon/battleye/proto.py
+++ b/rcon/battleye/proto.py
@@ -4,7 +4,13 @@ from typing import IO, NamedTuple
 from zlib import crc32
 
 
-__all__ = ['LoginRequest', 'Command']
+__all__ = [
+    'Header',
+    'LoginRequest',
+    'LoginResponse',
+    'Command',
+    'CommandResponse'
+]
 
 
 PREFIX = 'BE'
@@ -45,7 +51,7 @@ class Header(NamedTuple):
 
     @classmethod
     def read(cls, file: IO):
-        """Reads the packet from a socket."""
+        """Reads the packet from a file-like object."""
         return cls.from_bytes(file.read(2), file.read(4), file.read(1))
 
 
@@ -74,6 +80,28 @@ class LoginRequest(NamedTuple):
         return cls(0x00, passwd)
 
 
+class LoginResponse(NamedTuple):
+    """A login response."""
+
+    header: Header
+    type: int
+    success: bool
+
+    @classmethod
+    def from_bytes(cls, header: Header, typ: bytes, success: bytes):
+        """Creates a login response from the given bytes."""
+        return cls(
+            header,
+            int.from_bytes(typ, 'little'),
+            bool(int.from_bytes(success, '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."""
 
@@ -107,3 +135,27 @@ class Command(NamedTuple):
     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
+    type: int
+    seq: int
+    payload: bytes
+
+    @classmethod
+    def from_bytes(cls, header: Header, data: 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:]
+        )
+
+    @property
+    def message(self) -> str:
+        """Returns the text message."""
+        return self.payload.decode('ascii')