diff --git a/lslsh.py b/lslsh.py index 7d49b04..c94e88b 100644 --- a/lslsh.py +++ b/lslsh.py @@ -1,6 +1,6 @@ import requests import warnings -from urllib3.connectionpool import InsecureRequestWarning # type: ignore +from urllib3.connectionpool import InsecureRequestWarning # type: ignore from json.decoder import JSONDecodeError import sys from typing import Dict @@ -8,10 +8,11 @@ import readline import cmd SECRET_KEY: str = "29731e5170353a8b235098c43cd2099a4e805c55fb4395890e81f437c17334a9" -INTRO_TEXT: str = "lslsh 0.0.1\nType \"help\" for more information." +INTRO_TEXT: str = 'lslsh 0.0.1\nType "help" for more information.' warnings.filterwarnings("ignore", category=InsecureRequestWarning) + class Shell(cmd.Cmd): prompt = "> " url = None @@ -30,22 +31,24 @@ class Shell(cmd.Cmd): return None def send_cmd(self, url: str, cmd: str) -> Dict: - data = {"secret_key": SECRET_KEY, - "command": cmd} - + data = {"secret_key": SECRET_KEY, "command": cmd} + try: response = requests.post(url, json=data, verify=False) response_data = response.json() - except (requests.ConnectionError, requests.exceptions.MissingSchema, - requests.exceptions.InvalidURL): + except ( + requests.ConnectionError, + requests.exceptions.MissingSchema, + requests.exceptions.InvalidURL, + ): raise requests.exceptions.InvalidURL except JSONDecodeError: raise Exception("Error: Response has malformed json") - + error = response_data.get("error", None) if error: raise Exception(f"Error: {error}") - + return response_data def do_connect(self, url): @@ -55,14 +58,16 @@ class Shell(cmd.Cmd): except requests.exceptions.InvalidURL: print("Error: Invalid URL") return - + uuid = result.get("uuid", None) if not uuid: print("Error: Invalid response") return - + print(f"Connected to {uuid}") - print("_______________________________________________________________________________") + print( + "_______________________________________________________________________________" + ) print("") self.url = url @@ -73,7 +78,7 @@ class Shell(cmd.Cmd): self.do_disconnect(None) return True - + def do_disconnect(self, arg): """Disconnect from remote.""" if self.url: @@ -86,18 +91,20 @@ class Shell(cmd.Cmd): url = None else: print("Error: Not connected to remote.") - + def add_cmd(self, name, help_text): """Make a new command available within the shell.""" + def do_cmd(arg): result = self.send_cmd(self.url, arg) print(result) - + do_cmd.__doc__ = help_text do_cmd.__name__ = name - + setattr(self, f"do_{name}", do_cmd) + def run(): shell = Shell() try: @@ -105,4 +112,5 @@ def run(): except KeyboardInterrupt: shell.do_exit(None) + run()