ref: remove nubia

clients/windows/client/__init__.py

@@ -41,7 +41,6 @@ def __init__(self):
             self.configuration.get("username", "unknown"),
             self.configuration.get("password", "unknown"),
             self.configuration.get("host_key", "unknown"),
-            lambda: self.stop()
         )
 
     def get_server(self):

clients/windows/client/assets_loader.py

@@ -1,6 +1,7 @@
 import os
 
 import requests
+from termcolor import cprint
 
 
 def download(url: str):
@@ -13,15 +14,15 @@ def download(url: str):
 
     r = requests.get(url, stream=True)
     if r.ok:
-        print("saving to", os.path.abspath(file_path))
+        cprint(f"saving to {os.path.abspath(file_path)}", 'green')
         with open(file_path, 'wb') as f:
             for chunk in r.iter_content(chunk_size=1024 * 8):
                 if chunk:
                     f.write(chunk)
                     f.flush()
                     os.fsync(f.fileno())
     else:  # HTTP status code 4XX/5XX
-        print("Download failed: status code {}\n{}".format(r.status_code, r.text))
+        cprint("Download failed: status code {}\n{}".format(r.status_code, r.text), 'red')
 
 
 assets = [
@@ -49,16 +50,19 @@ def download(url: str):
 
 
 def validate_assets():
-
+    called_once = False
     for asset in assets:
         if not os.path.exists(
                 os.path.join(os.path.relpath("assets"), asset["name"])
         ):
-            print(f"Could not find {asset['name']} in assets directory.")
+            called_once = True
+            cprint(f"Could not find {asset['name']} in assets directory.", 'yellow')
             if 'handler' not in asset or asset.get('handler') is None:
-                print("I can't download it for you. This may break the application. Skipping.")
+                cprint("I can't download it for you. This may break the application. Skipping.")
                 continue
             else:
-                print(f"Downloading {asset['name']}")
+                cprint(f"Downloading {asset['name']}")
                 asset['handler'](asset)
 
+    if not called_once:
+        cprint("All assets are already loaded", "green")

clients/windows/client/plink.py

@@ -8,7 +8,6 @@
 
 
 class PLink:
-    max_retries = 5
 
     def __init__(
             self,
@@ -18,7 +17,6 @@ def __init__(
             username: str,
             password: str,
             host_key: str,
-            on_interrupt=None
     ):
         self.socks_port = socks_port
         self.server = server
@@ -28,57 +26,40 @@ def __init__(
         self.host_key = host_key
         self.process_thread = None
         self.subprocess = None
-        self.current_retries = 0
         self.stopped = False
-        self.on_interrupt = on_interrupt
 
     def get_process_path(self):
         return os.path.join(
             os.path.relpath("assets"),
             "plink.exe"
         )
 
-    def set_process(self, process):
-        self.subprocess = process
-
-    def _on_exit(self):
-        if self.current_retries != self.max_retries and not self.stopped:
-            cprint("Disconnected from server. Retrying ...", "red")
-            self.start()
-        elif not self.stopped:
-            self.stop()
-            self.on_interrupt()
-
     def start(self):
-        self.stopped = False
         if self.process_thread is not None:
             return
 
         si = subprocess.STARTUPINFO()
         si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
 
-        self.process_thread = utils.popen_and_call(
-            self._on_exit,
-            self.set_process,
-            [
-                f" -hostkey {self.host_key} -ssh {self.server} -D {self.socks_port}"
-                f" -l {self.username} -P {self.server_port} -no-antispoof -pw {self.password}"
-                f" -N"
-            ],
-            {
-                "executable": f"{self.get_process_path()}",
-                "stdout": subprocess.PIPE,
-                "stderr": subprocess.DEVNULL,
-                "startupinfo": si,
-            }
+        self.subprocess = subprocess.Popen(
+            f"{self.get_process_path()} -hostkey {self.host_key} -ssh {self.server} -D {self.socks_port}"
+            f" -l {self.username} -P {self.server_port} -no-antispoof -pw {self.password}"
+            f" -N",
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            # startupinfo=si
         )
+        cprint(f"Socks proxy available on {self.socks_port}", 'yellow')
+        self.stopped = False
         setup_windows_proxy(self.socks_port)
 
     def stop(self):
-        self.stopped = True
+        current_stop_status = self.stopped
         if self.subprocess is not None:
+            self.stopped = True
             self.subprocess.terminate()
             self.subprocess.wait()
             self.subprocess = None
-        drop_windows_proxy()
-        self.process_thread = None
+        if not current_stop_status:
+            drop_windows_proxy()
+            self.process_thread = None

clients/windows/client/utils.py

@@ -2,20 +2,19 @@
 import threading
 
 
-def popen_and_call(on_exit, on_proc, popen_args, popen_kwargs):
+def popen_and_call(on_proc, popen_args, popen_kwargs):
     """
     Runs the given args in a subprocess.Popen, and then calls the function
     on_exit when the subprocess completes.
     on_exit is a callable object, and popen_args is a list/tuple of args that
     would give to subprocess.Popen.
     """
-    def run_in_thread(on_exit, on_proc, popen_args, popen_kwargs):
+    def run_in_thread(on_proc, popen_args, popen_kwargs):
         proc = subprocess.Popen(*popen_args, **popen_kwargs)
         on_proc(proc)
         proc.wait()
-        on_exit()
         return
-    thread = threading.Thread(target=run_in_thread, args=(on_exit, on_proc, popen_args, popen_kwargs))
+    thread = threading.Thread(target=run_in_thread, args=(on_exit, on_proc, popen_args, popen_kwargs), daemon=True)
     thread.start()
     # returns immediately after the thread starts
     return thread

clients/windows/mooshak.py

@@ -1,9 +1,13 @@
-import sys
 import os
-from nubia import Nubia, Options
+import signal
+import sys
+import threading
+import time
+
 from termcolor import cprint
-from client import rights
-import client.commands
+
+from client import rights, MooshakClient
+from client.assets_loader import validate_assets
 
 os.system('color')
 
@@ -21,12 +25,41 @@
         cprint("Run this program as administrator", "red")
         exit(0)
 
-    cprint(BANNER, "light_cyan")
+    if len(sys.argv) < 2:
+        print("Argument missing. Choose between: load_assets, connect, disconnect")
+        exit(0)
+
+    command = sys.argv[1]
+
+    # Connect
+    if command == "connect":
+        cprint(BANNER, "light_cyan")
 
-    shell = Nubia(
-        name="Mooshak Proxy CLI",
-        command_pkgs=client,
-        options=Options(persistent_history=False, auto_execute_single_suggestions=False),
-    )
-    sys.exit(shell.run())
+        client = MooshakClient()
 
+        def signal_handler(signal, frame):
+            cprint('Disconnecting gracefully ...', 'green')
+            client.stop()
+            cprint('Disconnected', 'red')
+            sys.exit(0)
+
+        signal.signal(signal.SIGINT, signal_handler)
+        client.start()
+        cprint('Connected', 'green')
+
+        cprint('Press Ctrl+C to disconnect')
+        while True:
+            try:
+                time.sleep(3)
+            except InterruptedError:
+                signal.raise_signal(signal.SIGINT)
+
+    elif command == "disconnect":
+        MooshakClient().stop()
+
+    elif command == "load_assets":
+        validate_assets()
+
+    else:
+        cprint("Unknown command.")
+        exit(0)

clients/windows/requirements.txt

@@ -1,2 +1,2 @@
 requests==2.28.2
-python-nubia
+termcolor