Rewrite ErrorHandler

rcon/errorhandler.py

@@ -2,9 +2,8 @@
 
 from logging import Logger
 from socket import timeout
-from sys import exit    # pylint: disable=W0622
 
-from rcon.exceptions import RequestIdMismatch, WrongPassword
+from rcon.exceptions import ConfigReadError, RequestIdMismatch, WrongPassword
 
 
 __all__ = ['ErrorHandler']
@@ -13,11 +12,12 @@ __all__ = ['ErrorHandler']
 class ErrorHandler:
     """Handles common errors and exits."""
 
-    __slots__ = ('logger',)
+    __slots__ = ('logger', 'exit_code')
 
     def __init__(self, logger: Logger):
         """Sets the logger."""
         self.logger = logger
+        self.exit_code = 0
 
     def __enter__(self):
         return self
@@ -26,16 +26,19 @@ class ErrorHandler:
         """Checks for connection errors and exits respectively."""
         if isinstance(value, ConnectionRefusedError):
             self.logger.error('Connection refused.')
-            exit(3)
+            self.exit_code = 3
-
+        elif isinstance(value, (TimeoutError, timeout)):
-        if isinstance(value, (TimeoutError, timeout)):
             self.logger.error('Connection timed out.')
-            exit(4)
+            self.exit_code = 4
-
+        elif isinstance(value, WrongPassword):
-        if isinstance(value, WrongPassword):
             self.logger.error('Wrong password.')
-            exit(5)
+            self.exit_code = 5
-
+        elif isinstance(value, RequestIdMismatch):
-        if isinstance(value, RequestIdMismatch):
             self.logger.error('Session timed out.')
-            exit(5)
+            self.exit_code = 6
+        elif isinstance(value, ConfigReadError):
+            self.exit_code = value.exit_code
+        else:
+            return None
+
+        return True

rcon/rconclt.py

@@ -51,10 +51,13 @@ def main() -> int:
     except ConfigReadError as cre:
         return cre.exit_code
 
-    with ErrorHandler(LOGGER):
+    with ErrorHandler(LOGGER) as handler:
         with Client(host, port, timeout=args.timeout) as client:
             client.login(passwd)
             text = client.run(args.command, *args.argument)
 
+    if handler.exit_code:
+        return handler.exit_code
+
     print(text, flush=True)
     return 0

rcon/rconshell.py

@@ -47,8 +47,11 @@ def main() -> int:
     else:
         host = port = passwd = None
 
-    with ErrorHandler(LOGGER):
+    with ErrorHandler(LOGGER) as handler:
         with CommandHistory(LOGGER):
             rconcmd(host, port, passwd, prompt=args.prompt)
 
+    if handler.exit_code:
+        return handler.exit_code
+
     return 0