From aa1306d37394410ec9090f171a3d3b3898617c9a Mon Sep 17 00:00:00 2001 From: Richard Neumann Date: Sun, 13 Feb 2022 17:44:05 +0100 Subject: [PATCH] Extend console script to be able to select protocol type --- rcon/battleye/__init__.py | 7 ++++++- rcon/battleye/client.py | 11 +++++++---- rcon/battleye/proto.py | 9 +++++++-- rcon/client.py | 4 ++++ rcon/{source => }/config.py | 2 +- rcon/{source => }/console.py | 20 ++++++++++++-------- rcon/{source => }/errorhandler.py | 10 +++++----- rcon/exceptions.py | 24 ++++++++++++++++++++++++ rcon/{source => }/rconclt.py | 13 +++++++++---- rcon/{source => }/rconshell.py | 14 ++++++++++---- rcon/source/__init__.py | 3 +-- rcon/source/async_rcon.py | 4 ++-- rcon/source/client.py | 4 ++-- rcon/source/exceptions.py | 19 +------------------ rcon/source/gui.py | 11 ++++------- 15 files changed, 95 insertions(+), 60 deletions(-) rename rcon/{source => }/config.py (97%) rename rcon/{source => }/console.py (85%) rename rcon/{source => }/errorhandler.py (81%) create mode 100644 rcon/exceptions.py rename rcon/{source => }/rconclt.py (78%) rename rcon/{source => }/rconshell.py (72%) diff --git a/rcon/battleye/__init__.py b/rcon/battleye/__init__.py index 51c33c9..faae646 100644 --- a/rcon/battleye/__init__.py +++ b/rcon/battleye/__init__.py @@ -1 +1,6 @@ -"""BattlEye RCON implementation.""" \ No newline at end of file +"""BattlEye RCON implementation.""" + +from rcon.battleye.client import Client + + +__all__ = ['Client'] diff --git a/rcon/battleye/client.py b/rcon/battleye/client.py index 76b2e05..bd38ef4 100644 --- a/rcon/battleye/client.py +++ b/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 '' diff --git a/rcon/battleye/proto.py b/rcon/battleye/proto.py index 6019e94..8189dea 100644 --- a/rcon/battleye/proto.py +++ b/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])) diff --git a/rcon/client.py b/rcon/client.py index 1ab14db..751ab1f 100644 --- a/rcon/client.py +++ b/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() diff --git a/rcon/source/config.py b/rcon/config.py similarity index 97% rename from rcon/source/config.py rename to rcon/config.py index b59ea57..f81860d 100644 --- a/rcon/source/config.py +++ b/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'] diff --git a/rcon/source/console.py b/rcon/console.py similarity index 85% rename from rcon/source/console.py rename to rcon/console.py index 26eb6f0..ce82287 100644 --- a/rcon/source/console.py +++ b/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: diff --git a/rcon/source/errorhandler.py b/rcon/errorhandler.py similarity index 81% rename from rcon/source/errorhandler.py rename to rcon/errorhandler.py index 7892d18..92dbfe2 100644 --- a/rcon/source/errorhandler.py +++ b/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.') } diff --git a/rcon/exceptions.py b/rcon/exceptions.py new file mode 100644 index 0000000..4a04377 --- /dev/null +++ b/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.""" diff --git a/rcon/source/rconclt.py b/rcon/rconclt.py similarity index 78% rename from rcon/source/rconclt.py rename to rcon/rconclt.py index 145a7e6..253ec3e 100644 --- a/rcon/source/rconclt.py +++ b/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) diff --git a/rcon/source/rconshell.py b/rcon/rconshell.py similarity index 72% rename from rcon/source/rconshell.py rename to rcon/rconshell.py index 4a7fa7a..e8c3f58 100644 --- a/rcon/source/rconshell.py +++ b/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: diff --git a/rcon/source/__init__.py b/rcon/source/__init__.py index b58ae35..4d7e143 100644 --- a/rcon/source/__init__.py +++ b/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'] diff --git a/rcon/source/async_rcon.py b/rcon/source/async_rcon.py index 7533377..10d4c70 100644 --- a/rcon/source/async_rcon.py +++ b/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) diff --git a/rcon/source/client.py b/rcon/source/client.py index dfe9437..82ff2f0 100644 --- a/rcon/source/client.py +++ b/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) diff --git a/rcon/source/exceptions.py b/rcon/source/exceptions.py index 9af01bb..2f42588 100644 --- a/rcon/source/exceptions.py +++ b/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.""" diff --git a/rcon/source/gui.py b/rcon/source/gui.py index 3621928..396c53f 100644 --- a/rcon/source/gui.py +++ b/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