Added ErrorHandler.

rcon/config.py

@@ -9,11 +9,12 @@ from typing import Dict, Iterator, NamedTuple, Tuple
 from rcon.exceptions import InvalidConfig
 
 
-__all__ = ['servers']
+__all__ = ['CONFIG_FILE', 'LOG_FORMAT', 'Config', 'servers']
 
 
 CONFIG = ConfigParser()
 CONFIG_FILE = Path('/etc/rcon.conf')
+LOG_FORMAT = '[%(levelname)s] %(name)s: %(message)s'
 LOGGER = getLogger('RCON Config')
 
 
@@ -64,8 +65,8 @@ def entries(config_parser: ConfigParser) -> Iterator[Tuple[str, Config]]:
         yield (section, Config.from_config_section(config_parser[section]))
 
 
-def servers() -> Dict[str, Config]:
+def servers(config_file: Path = CONFIG_FILE) -> Dict[str, Config]:
     """Returns a dictionary of servers."""
 
-    CONFIG.read(CONFIG_FILE)
+    CONFIG.read(config_file)
     return dict(entries(CONFIG))

rcon/errorhandler.py

@@ -0,0 +1,30 @@
+"""Common errors handler."""
+
+from logging import Logger
+from sys import exit    # pylint: disable=W0622
+from typing import Iterable, Tuple
+
+
+__all__ = ['ErrorHandler']
+
+
+ErrorMap = Iterable[Tuple[Exception, str, int]]
+
+
+class ErrorHandler:
+    """Handles common errors and exits."""
+
+    def __init__(self, errors: ErrorMap, logger: Logger):
+        """Sets the logger."""
+        self.errors = errors
+        self.logger = logger
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, typ, *_):
+        """Checks for connection errors and exits respectively."""
+        for error, (message, returncode) in self.errors.items():
+            if isinstance(typ, error):
+                self.logger.error(message)
+                exit(returncode)

rcon/rconclt.py

@@ -3,12 +3,14 @@
 from argparse import ArgumentParser, Namespace
 from getpass import getpass
 from logging import DEBUG, INFO, basicConfig, getLogger
+from pathlib import Path
 from socket import timeout
 from sys import exit    # pylint: disable=W0622
 from typing import Tuple
 
+from rcon.errorhandler import ErrorHandler
 from rcon.exceptions import InvalidConfig
-from rcon.config import Config, servers
+from rcon.config import CONFIG_FILE, LOG_FORMAT, Config, servers
 from rcon.exceptions import RequestIdMismatch, WrongPassword
 from rcon.proto import Client
 
@@ -16,8 +18,13 @@ from rcon.proto import Client
 __all__ = ['get_credentials', 'main']
 
 
+ERRORS = (
+    (ConnectionRefusedError, 'Connection refused.', 3),
+    (timeout, 'Connection timeout.', 4),
+    (RequestIdMismatch, 'Unexpected request ID mismatch.', 5),
+    (WrongPassword, 'Wrong password.', 6),
+)
 LOGGER = getLogger('rconclt')
-LOG_FORMAT = '[%(levelname)s] %(name)s: %(message)s'
 
 
 def get_args() -> Namespace:
@@ -25,26 +32,28 @@ def get_args() -> Namespace:
 
     parser = ArgumentParser(description='A Minecraft RCON client.')
     parser.add_argument('server', help='the server to connect to')
-    parser.add_argument('-t', '--timeout', type=float, metavar='seconds',
-                        help='connection timeout in seconds')
+    parser.add_argument('-c', '--config', type=Path, metavar='file',
+                        default=CONFIG_FILE, help='the configuration file')
     parser.add_argument('-d', '--debug', action='store_true',
                         help='print additional debug information')
+    parser.add_argument('-t', '--timeout', type=float, metavar='seconds',
+                        help='connection timeout in seconds')
     parser.add_argument('command', help='command to execute on the server')
     parser.add_argument('argument', nargs='*', default=(),
                         help='arguments for the command')
     return parser.parse_args()
 
 
-def get_credentials(server: str) -> Tuple[str, int, str]:
+def get_credentials(args: Namespace) -> Tuple[str, int, str]:
     """Get the credentials for a server from the respective server name."""
 
     try:
-        host, port, passwd = Config.from_string(server)
+        host, port, passwd = Config.from_string(args.server)
     except InvalidConfig:
         try:
-            host, port, passwd = servers()[server]
+            host, port, passwd = servers(args.config)[args.server]
         except KeyError:
-            LOGGER.error('No such server: %s.', server)
+            LOGGER.error('No such server: %s.', args.server)
             exit(2)
 
     if passwd is None:
@@ -57,30 +66,17 @@ def get_credentials(server: str) -> Tuple[str, int, str]:
 
     return (host, port, passwd)
 
-
 def main():
     """Runs the RCON client."""
 
     args = get_args()
     log_level = DEBUG if args.debug else INFO
     basicConfig(level=log_level, format=LOG_FORMAT)
-    host, port, passwd = get_credentials(args.server)
+    host, port, passwd = get_credentials(args)
 
-    try:
+    with ErrorHandler(ERRORS, LOGGER):
         with Client(host, port, timeout=args.timeout) as client:
             client.login(passwd)
             text = client.run(args.command, *args.argument)
-    except ConnectionRefusedError:
-        LOGGER.error('Connection refused.')
-        exit(3)
-    except timeout:
-        LOGGER.error('Connection timeout.')
-        exit(4)
-    except RequestIdMismatch:
-        LOGGER.error('Unexpected request ID mismatch.')
-        exit(5)
-    except WrongPassword:
-        LOGGER.error('Wrong password.')
-        exit(6)
-    else:
-        print(text, flush=True)
+
+    print(text, flush=True)

rcon/rconshell.py

@@ -2,18 +2,23 @@
 
 from argparse import ArgumentParser, Namespace
 from logging import INFO, basicConfig, getLogger
+from pathlib import Path
 from socket import timeout
 from sys import exit    # pylint: disable=W0622
 
+from rcon.errorhandler import ErrorHandler
 from rcon.rconclt import get_credentials
+from rcon.config import CONFIG_FILE, LOG_FORMAT
 from rcon.console import rconcmd
 
 
 __all__ = ['get_args', 'main']
 
-
+ERRORS = (
+    (ConnectionRefusedError, 'Connection refused.', 3),
+    (timeout, 'Connection timeout.', 4)
+)
 LOGGER = getLogger('rconshell')
-LOG_FORMAT = '[%(levelname)s] %(name)s: %(message)s'
 
 
 def get_args() -> Namespace:
@@ -21,6 +26,8 @@ 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('-c', '--config', type=Path, metavar='file',
+                        default=CONFIG_FILE, help='the configuration file')
     parser.add_argument('-p', '--prompt', default='RCON> ', metavar='PS1',
                         help='the shell prompt')
     return parser.parse_args()
@@ -32,18 +39,12 @@ def main():
     args = get_args()
     basicConfig(level=INFO, format=LOG_FORMAT)
 
-    if server := args.server:
-        host, port, passwd = get_credentials(server)
+    if args.server:
+        host, port, passwd = get_credentials(args)
     else:
         host = port = passwd = None
 
-    try:
+    with ErrorHandler(ERRORS, LOGGER):
         exit_code = rconcmd(host, port, passwd, args.prompt)
-    except ConnectionRefusedError:
-        LOGGER.error('Connection refused.')
-        exit(3)
-    except timeout:
-        LOGGER.error('Connection timeout.')
-        exit(4)
 
     exit(exit_code)