diff --git a/rcon/config.py b/rcon/config.py index b2619ed..9932076 100644 --- a/rcon/config.py +++ b/rcon/config.py @@ -7,9 +7,10 @@ from getpass import getpass from logging import getLogger from os import getenv, name from pathlib import Path -from sys import exit # pylint: disable=W0622 from typing import Iterable, NamedTuple, Optional, Union +from rcon.exceptions import ConfigReadError + __all__ = ['CONFIG_FILES', 'LOG_FORMAT', 'SERVERS', 'Config', 'from_args'] @@ -92,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) - exit(2) + raise ConfigReadError(2) if passwd is None: try: @@ -100,6 +101,6 @@ def from_args(args: Namespace) -> Config: except (KeyboardInterrupt, EOFError): print() LOGGER.error('Aborted by user.') - exit(1) + raise ConfigReadError(1) return Config(host, port, passwd) diff --git a/rcon/exceptions.py b/rcon/exceptions.py index d8407b0..d84f317 100644 --- a/rcon/exceptions.py +++ b/rcon/exceptions.py @@ -1,6 +1,14 @@ """RCON exceptions.""" -__all__ = ['RequestIdMismatch', 'WrongPassword'] +__all__ = ['ConfigReadError', 'RequestIdMismatch', '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): diff --git a/rcon/rconclt.py b/rcon/rconclt.py index e73736a..55ef6e9 100644 --- a/rcon/rconclt.py +++ b/rcon/rconclt.py @@ -7,6 +7,7 @@ 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'] @@ -39,12 +40,16 @@ def get_args() -> Namespace: return parser.parse_args() -def main() -> None: +def main() -> int: """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 ErrorHandler(LOGGER): with Client(host, port, timeout=args.timeout) as client: @@ -52,3 +57,4 @@ def main() -> None: text = client.run(args.command, *args.argument) print(text, flush=True) + return 0 diff --git a/rcon/rconshell.py b/rcon/rconshell.py index 473fabc..fbd8c03 100644 --- a/rcon/rconshell.py +++ b/rcon/rconshell.py @@ -7,6 +7,7 @@ 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 @@ -32,17 +33,22 @@ def get_args() -> Namespace: return parser.parse_args() -def main() -> None: +def main() -> int: """Runs the RCON shell.""" args = get_args() basicConfig(level=INFO, format=LOG_FORMAT) if args.server: - host, port, passwd = from_args(args) + try: + host, port, passwd = from_args(args) + except ConfigReadError as cre: + return cre.exit_code else: host = port = passwd = None with ErrorHandler(LOGGER): with CommandHistory(LOGGER): rconcmd(host, port, passwd, prompt=args.prompt) + + return 0