From 4bb56907411c4c4172b772977463338bb6b41ea5 Mon Sep 17 00:00:00 2001 From: Richard Neumann Date: Thu, 10 Feb 2022 10:56:51 +0100 Subject: [PATCH] Rewrite ErrorHandler --- rcon/config.py | 6 +++--- rcon/errorhandler.py | 41 +++++++++++++++++++++-------------------- rcon/exceptions.py | 17 +++++++++++------ rcon/rconclt.py | 27 +++++++++++++-------------- rcon/rconshell.py | 23 +++++++++++------------ 5 files changed, 59 insertions(+), 55 deletions(-) diff --git a/rcon/config.py b/rcon/config.py index 9932076..f81860d 100644 --- a/rcon/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.exceptions import ConfigReadError +from rcon.exceptions import ConfigReadError, UserAbort __all__ = ['CONFIG_FILES', 'LOG_FORMAT', 'SERVERS', 'Config', 'from_args'] @@ -93,7 +93,7 @@ def from_args(args: Namespace) -> Config: host, port, passwd = SERVERS[args.server] except KeyError: LOGGER.error('No such server: %s.', args.server) - raise ConfigReadError(2) + raise ConfigReadError() from None if passwd is None: try: @@ -101,6 +101,6 @@ def from_args(args: Namespace) -> Config: except (KeyboardInterrupt, EOFError): print() LOGGER.error('Aborted by user.') - raise ConfigReadError(1) + raise UserAbort() from None return Config(host, port, passwd) diff --git a/rcon/errorhandler.py b/rcon/errorhandler.py index ad1af3d..6942eb9 100644 --- a/rcon/errorhandler.py +++ b/rcon/errorhandler.py @@ -3,12 +3,25 @@ from logging import Logger from socket import timeout -from rcon.exceptions import ConfigReadError, RequestIdMismatch, WrongPassword +from rcon.exceptions import ConfigReadError +from rcon.exceptions import RequestIdMismatch +from rcon.exceptions import UserAbort +from rcon.exceptions import WrongPassword __all__ = ['ErrorHandler'] +ERRORS = { + UserAbort: 1, + ConfigReadError: 2, + ConnectionRefusedError: 3, + (TimeoutError, timeout): 4, + WrongPassword: 5, + RequestIdMismatch: 6 +} + + class ErrorHandler: """Handles common errors and exits.""" @@ -22,23 +35,11 @@ class ErrorHandler: def __enter__(self): return self - def __exit__(self, _, value, __): + def __exit__(self, _, value: Exception, __): """Checks for connection errors and exits respectively.""" - if isinstance(value, ConnectionRefusedError): - self.logger.error('Connection refused.') - self.exit_code = 3 - elif isinstance(value, (TimeoutError, timeout)): - self.logger.error('Connection timed out.') - self.exit_code = 4 - elif isinstance(value, WrongPassword): - self.logger.error('Wrong password.') - self.exit_code = 5 - elif isinstance(value, RequestIdMismatch): - self.logger.error('Session timed out.') - self.exit_code = 6 - elif isinstance(value, ConfigReadError): - self.exit_code = value.exit_code - else: - return None - - return True + for typ, exit_code in ERRORS.items(): + if isinstance(value, typ): + self.exit_code = exit_code + return True + + return None diff --git a/rcon/exceptions.py b/rcon/exceptions.py index d84f317..9af01bb 100644 --- a/rcon/exceptions.py +++ b/rcon/exceptions.py @@ -1,25 +1,30 @@ """RCON exceptions.""" -__all__ = ['ConfigReadError', 'RequestIdMismatch', 'WrongPassword'] +__all__ = [ + 'ConfigReadError', + 'RequestIdMismatch', + 'UserAbort', + 'WrongPassword' +] class ConfigReadError(Exception): """Indicates an error while reading the configuration.""" - def __init__(self, exit_code: int): - super().__init__() - self.exit_code = exit_code - class RequestIdMismatch(Exception): """Indicates a mismatch in request IDs.""" def __init__(self, sent: int, received: int): - """Sets the sent and received IDs.""" + """Sets the IDs that have been sent and received.""" 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/rconclt.py b/rcon/rconclt.py index fbce649..b1285c3 100644 --- a/rcon/rconclt.py +++ b/rcon/rconclt.py @@ -7,7 +7,6 @@ from pathlib import Path from rcon.client import Client from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args from rcon.errorhandler import ErrorHandler -from rcon.exceptions import ConfigReadError __all__ = ['main'] @@ -40,24 +39,24 @@ def get_args() -> Namespace: return parser.parse_args() -def main() -> int: +def run() -> None: """Runs the RCON client.""" args = get_args() basicConfig(format=LOG_FORMAT, level=DEBUG if args.debug else INFO) + host, port, passwd = from_args(args) - try: - host, port, passwd = from_args(args) - except ConfigReadError as cre: - return cre.exit_code + with Client(host, port, timeout=args.timeout) as client: + client.login(passwd) + text = client.run(args.command, *args.argument) - with ErrorHandler(LOGGER) as handler: - with Client(host, port, timeout=args.timeout) as client: - client.login(passwd) - text = client.run(args.command, *args.argument) + print(text, flush=True) - if handler.exit_code: - return handler.exit_code - print(text, flush=True) - return 0 +def main() -> int: + """Runs the main script with exceptions handled.""" + + with ErrorHandler(LOGGER) as handler: + run() + + return handler.exit_code diff --git a/rcon/rconshell.py b/rcon/rconshell.py index acd3536..9e4fac5 100644 --- a/rcon/rconshell.py +++ b/rcon/rconshell.py @@ -7,7 +7,6 @@ from pathlib import Path from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args from rcon.console import PROMPT, rconcmd from rcon.errorhandler import ErrorHandler -from rcon.exceptions import ConfigReadError from rcon.readline import CommandHistory @@ -33,25 +32,25 @@ def get_args() -> Namespace: return parser.parse_args() -def main() -> int: +def run() -> None: """Runs the RCON shell.""" args = get_args() basicConfig(level=INFO, format=LOG_FORMAT) if args.server: - try: - host, port, passwd = from_args(args) - except ConfigReadError as cre: - return cre.exit_code + host, port, passwd = from_args(args) else: host = port = passwd = None - with ErrorHandler(LOGGER) as handler: - with CommandHistory(LOGGER): - rconcmd(host, port, passwd, prompt=args.prompt) + with CommandHistory(LOGGER): + rconcmd(host, port, passwd, prompt=args.prompt) + - if handler.exit_code: - return handler.exit_code +def main() -> int: + """Runs the main script with exceptions handled.""" + + with ErrorHandler(LOGGER) as handler: + run() - return 0 + return handler.exit_code