Extend console script to be able to select protocol type

rcon/battleye/__init__.py

@@ -1 +1,6 @@
 """BattlEye RCON implementation."""
+
+from rcon.battleye.client import Client
+
+
+__all__ = ['Client']

rcon/battleye/client.py

@@ -16,15 +16,18 @@ Host = Union[str, IPv4Address]
 class Client(BaseClient):
     """BattlEye RCon client."""
 
-    def communicate(self, data: bytes, *, recv: int = 4096) -> bytes:
+    def communicate(self, data: bytes, *, size: int = 4096) -> bytes:
         """Sends and receives packets."""
         self._socket.send(data)
-        return self._socket.recv(recv)
+        return self._socket.recv(size)
 
     def login(self, passwd: str) -> bytes:
         """Logs the user in."""
         return self.communicate(bytes(LoginRequest.from_passwd(passwd)))
 
-    def command(self, command: str) -> bytes:
+    def run(self, command: str, *args: str) -> str:
         """Executes a command."""
-        return self.communicate(bytes(Command.from_command(command)))
+        packet = Command.from_command(command, *args)
+        _ = self.communicate(bytes(packet))
+        # TODO: Process response
+        return ''

rcon/battleye/proto.py

@@ -99,6 +99,11 @@ class Command(NamedTuple):
         return Header.from_payload(self.payload)
 
     @classmethod
-    def from_command(cls, command: str):
-        """Creates a command packet from the given command."""
+    def from_string(cls, command: str):
+        """Creates a command packet from the given string."""
         return cls(0x01, 0x00, command)
+
+    @classmethod
+    def from_command(cls, command: str, *args: str):
+        """Creates a command packet from the command and arguments."""
+        return cls.from_string(' '.join([command, *args]))

rcon/client.py

@@ -61,3 +61,7 @@ class BaseClient:
     def login(self, passwd: str) -> bool:
         """Performs a login."""
         raise NotImplementedError()
+
+    def run(self, command: str, *args: str) -> str:
+        """Runs a command."""
+        raise NotImplementedError()

rcon/source/config.py → rcon/config.py

@@ -9,7 +9,7 @@ from os import getenv, name
 from pathlib import Path
 from typing import Iterable, NamedTuple, Optional, Union
 
-from rcon.source.exceptions import ConfigReadError, UserAbort
+from rcon.exceptions import ConfigReadError, UserAbort
 
 
 __all__ = ['CONFIG_FILES', 'LOG_FORMAT', 'SERVERS', 'Config', 'from_args']

rcon/source/console.py → rcon/console.py

@@ -1,10 +1,11 @@
 """An interactive console."""
 
 from getpass import getpass
+from typing import Type
 
-from rcon.source.client import Client
-from rcon.source.config import Config
-from rcon.source.exceptions import RequestIdMismatch, WrongPassword
+from rcon.client import BaseClient
+from rcon.config import Config
+from rcon.exceptions import SessionTimeout, WrongPassword
 
 
 __all__ = ['PROMPT', 'rconcmd']
@@ -75,7 +76,7 @@ def get_config(host: str, port: int, passwd: str) -> Config:
     return Config(host, port, passwd)
 
 
-def login(client: Client, passwd: str) -> str:
+def login(client: BaseClient, passwd: str) -> str:
     """Performs a login."""
 
     while True:
@@ -89,7 +90,7 @@ def login(client: Client, passwd: str) -> str:
         return passwd
 
 
-def process_input(client: Client, passwd: str, prompt: str) -> bool:
+def process_input(client: BaseClient, passwd: str, prompt: str) -> bool:
     """Processes the CLI input."""
 
     try:
@@ -111,7 +112,7 @@ def process_input(client: Client, passwd: str, prompt: str) -> bool:
 
     try:
         result = client.run(command, *args)
-    except RequestIdMismatch:
+    except SessionTimeout:
         print(MSG_SESSION_TIMEOUT)
 
         try:
@@ -125,7 +126,10 @@ def process_input(client: Client, passwd: str, prompt: str) -> bool:
     return True
 
 
-def rconcmd(host: str, port: int, passwd: str, *, prompt: str = PROMPT):
+def rconcmd(
+        client_cls: Type[BaseClient], host: str, port: int, passwd: str, *,
+        prompt: str = PROMPT
+):
     """Initializes the console."""
 
     try:
@@ -136,7 +140,7 @@ def rconcmd(host: str, port: int, passwd: str, *, prompt: str = PROMPT):
 
     prompt = prompt.format(host=host, port=port)
 
-    with Client(host, port) as client:
+    with client_cls(host, port) as client:
         try:
             passwd = login(client, passwd)
         except EOFError:

rcon/source/errorhandler.py → rcon/errorhandler.py

@@ -3,10 +3,10 @@
 from logging import Logger
 from socket import timeout
 
-from rcon.source.exceptions import ConfigReadError
-from rcon.source.exceptions import RequestIdMismatch
-from rcon.source.exceptions import UserAbort
-from rcon.source.exceptions import WrongPassword
+from rcon.exceptions import ConfigReadError
+from rcon.exceptions import SessionTimeout
+from rcon.exceptions import UserAbort
+from rcon.exceptions import WrongPassword
 
 
 __all__ = ['ErrorHandler']
@@ -18,7 +18,7 @@ ERRORS = {
     ConnectionRefusedError: (3, 'Connection refused.'),
     (TimeoutError, timeout): (4, 'Connection timed out.'),
     WrongPassword: (5, 'Wrong password.'),
-    RequestIdMismatch: (6, 'Session timed out.')
+    SessionTimeout: (6, 'Session timed out.')
 }
 
 

rcon/exceptions.py

@@ -0,0 +1,24 @@
+"""Common exceptions."""
+
+__all__ = [
+    'ConfigReadError',
+    'SessionTimeout',
+    'UserAbort',
+    'WrongPassword'
+]
+
+
+class ConfigReadError(Exception):
+    """Indicates an error while reading the configuration."""
+
+
+class SessionTimeout(Exception):
+    """Indicates that the session timed out."""
+
+
+class UserAbort(Exception):
+    """Indicates that a required action has been aborted by the user."""
+
+
+class WrongPassword(Exception):
+    """Indicates a wrong password."""

rcon/source/rconclt.py → rcon/rconclt.py

@@ -4,9 +4,9 @@ from argparse import ArgumentParser, Namespace
 from logging import DEBUG, INFO, basicConfig, getLogger
 from pathlib import Path
 
-from rcon.source.client import Client
-from rcon.source.config import CONFIG_FILES, LOG_FORMAT, from_args
-from rcon.source.errorhandler import ErrorHandler
+from rcon import battleye, source
+from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args
+from rcon.errorhandler import ErrorHandler
 
 
 __all__ = ['main']
@@ -20,6 +20,10 @@ def get_args() -> Namespace:
 
     parser = ArgumentParser(description='A Minecraft RCON client.')
     parser.add_argument('server', help='the server to connect to')
+    parser.add_argument(
+        '-B', '--battleye', action='store_true',
+        help='use BattlEye RCon instead of Source RCON'
+    )
     parser.add_argument(
         '-c', '--config', type=Path, metavar='file', default=CONFIG_FILES,
         help='the configuration file'
@@ -45,8 +49,9 @@ def run() -> None:
     args = get_args()
     basicConfig(format=LOG_FORMAT, level=DEBUG if args.debug else INFO)
     host, port, passwd = from_args(args)
+    client_cls = battleye.Client if args.battleye else source.Client
 
-    with Client(host, port, timeout=args.timeout) as client:
+    with client_cls(host, port, timeout=args.timeout) as client:
         client.login(passwd)
         text = client.run(args.command, *args.argument)
 

rcon/source/rconshell.py → rcon/rconshell.py

@@ -4,10 +4,11 @@ from argparse import ArgumentParser, Namespace
 from logging import INFO, basicConfig, getLogger
 from pathlib import Path
 
+from rcon import battleye, source
 from rcon.readline import CommandHistory
-from rcon.source.config import CONFIG_FILES, LOG_FORMAT, from_args
-from rcon.source.console import PROMPT, rconcmd
-from rcon.source.errorhandler import ErrorHandler
+from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args
+from rcon.console import PROMPT, rconcmd
+from rcon.errorhandler import ErrorHandler
 
 
 __all__ = ['get_args', 'main']
@@ -21,6 +22,10 @@ def get_args() -> Namespace:
 
     parser = ArgumentParser(description='An interactive RCON shell.')
     parser.add_argument('server', nargs='?', help='the server to connect to')
+    parser.add_argument(
+        '-B', '--battleye', action='store_true',
+        help='use BattlEye RCon instead of Source RCON'
+    )
     parser.add_argument(
         '-c', '--config', type=Path, metavar='file', default=CONFIG_FILES,
         help='the configuration file'
@@ -37,6 +42,7 @@ def run() -> None:
 
     args = get_args()
     basicConfig(level=INFO, format=LOG_FORMAT)
+    client_cls = battleye.Client if args.battleye else source.Client
 
     if args.server:
         host, port, passwd = from_args(args)
@@ -44,7 +50,7 @@ def run() -> None:
         host = port = passwd = None
 
     with CommandHistory(LOGGER):
-        rconcmd(host, port, passwd, prompt=args.prompt)
+        rconcmd(client_cls, host, port, passwd, prompt=args.prompt)
 
 
 def main() -> int:

rcon/source/__init__.py

@@ -2,7 +2,6 @@
 
 from rcon.source.async_rcon import rcon
 from rcon.source.client import Client
-from rcon.source.exceptions import RequestIdMismatch, WrongPassword
 
 
-__all__ = ['RequestIdMismatch', 'WrongPassword', 'Client', 'rcon']
+__all__ = ['Client', 'rcon']

rcon/source/async_rcon.py

@@ -2,7 +2,7 @@
 
 from asyncio import StreamReader, StreamWriter, open_connection
 
-from rcon.source.exceptions import RequestIdMismatch, WrongPassword
+from rcon.exceptions import SessionTimeout, WrongPassword
 from rcon.source.proto import Packet, Type
 
 
@@ -56,6 +56,6 @@ async def rcon(
     await close(writer)
 
     if response.id != request.id:
-        raise RequestIdMismatch(request.id, response.id)
+        raise SessionTimeout()
 
     return response.payload.decode(encoding)

rcon/source/client.py

@@ -1,7 +1,7 @@
 """Synchronous client."""
 
 from rcon.client import BaseClient
-from rcon.source.exceptions import RequestIdMismatch, WrongPassword
+from rcon.exceptions import SessionTimeout, WrongPassword
 from rcon.source.proto import Packet, Type
 
 
@@ -44,6 +44,6 @@ class Client(BaseClient):
         response = self.communicate(request)
 
         if response.id != request.id:
-            raise RequestIdMismatch(request.id, response.id)
+            raise SessionTimeout()
 
         return response.payload.decode(encoding)

rcon/source/exceptions.py

@@ -1,15 +1,6 @@
 """RCON exceptions."""
 
-__all__ = [
-    'ConfigReadError',
-    'RequestIdMismatch',
-    'UserAbort',
-    'WrongPassword'
-]
-
-
-class ConfigReadError(Exception):
-    """Indicates an error while reading the configuration."""
+__all__ = ['RequestIdMismatch']
 
 
 class RequestIdMismatch(Exception):
@@ -20,11 +11,3 @@ class RequestIdMismatch(Exception):
         super().__init__()
         self.sent = sent
         self.received = received
-
-
-class UserAbort(Exception):
-    """Indicates that a required action has been aborted by the user."""
-
-
-class WrongPassword(Exception):
-    """Indicates a wrong password."""

rcon/source/gui.py

@@ -13,8 +13,8 @@ require_version('Gtk', '3.0')
 from gi.repository import Gtk
 
 from rcon.source.client import Client
-from rcon.source.config import LOG_FORMAT
-from rcon.source.exceptions import RequestIdMismatch, WrongPassword
+from rcon.config import LOG_FORMAT
+from rcon.exceptions import SessionTimeout, WrongPassword
 
 
 __all__ = ['main']
@@ -193,11 +193,8 @@ class GUI(Gtk.Window):  # pylint: disable=R0902
             self.show_error('Connection timed out.')
         except WrongPassword:
             self.show_error('Wrong password.')
-        except RequestIdMismatch as mismatch:
-            self.show_error(
-                'Request ID mismatch.\n'
-                f'Expected {mismatch.sent}, but got {mismatch.received}.'
-            )
+        except SessionTimeout:
+            self.show_error('Session timed out.')
         else:
             self.result_text = result