diff --git a/lib.py b/lib.py index 61e58af..b513252 100644 --- a/lib.py +++ b/lib.py @@ -12,6 +12,14 @@ class ErrorReceived(Exception): pass +class SessionLockedError(Exception): + pass + + +class UnauthorizedError(Exception): + pass + + def send_cmd(url: str, secret_key: str, cmd: str) -> Dict: """Send a command to the endpoint and return the response.""" data = {"secret_key": secret_key, "command": cmd} @@ -27,9 +35,10 @@ def send_cmd(url: str, secret_key: str, cmd: str) -> Dict: raise requests.exceptions.InvalidURL except requests.exceptions.HTTPError as e: code = e.response.status_code - if code == 401 or code == 423: - error_msg = e.response.json().get("error") - raise ErrorReceived(f"Error: {error_msg}") + if code == 423: + raise SessionLockedError("Error: " + e.response.json().get("error")) + elif code == 401: + raise UnauthorizedError("Error: " + e.response.json().get("error")) else: raise e diff --git a/lslsh.py b/lslsh.py index c44cedd..2fecc45 100755 --- a/lslsh.py +++ b/lslsh.py @@ -13,6 +13,8 @@ from urllib3.connectionpool import InsecureRequestWarning # type: ignore from lib import ( ErrorReceived, InvalidResponseError, + SessionLockedError, + UnauthorizedError, connect, disconnect, get_available_commands, @@ -52,7 +54,12 @@ class Shell(cmd.Cmd): except requests.exceptions.InvalidURL as e: print("Error: Invalid URL") return - except (InvalidResponseError, ErrorReceived) as e: + except ( + InvalidResponseError, + ErrorReceived, + UnauthorizedError, + SessionLockedError, + ) as e: print(e) return