Load commands dynamically

oo_bin/commander.py

@@ -0,0 +1,73 @@
+import importlib
+import importlib.util
+from pathlib import Path
+
+
+class Commander:
+    """Not quite Cobra, but it snakes it's way through paths to find the common command format"""
+
+    def __init__(self, path):
+        self.commands = {}
+        self.load_from_path(path)
+
+    def find_in_path(self, path):
+        shim = "shims" in path
+        cmds = {}
+        p = Path(path)
+
+        def filename(file):
+            return (
+                str(file).replace(f"{path}/", "").replace(".py", "").replace("/", ".")
+            )
+
+        for cmd_file in p.glob("*/command/*.py"):
+            module_name = filename(cmd_file)
+            if "__init__" not in module_name:
+                cmds[module_name] = (
+                    module_name.split(".")[-1],
+                    str(cmd_file) if shim else None,
+                )
+
+        for cmd_file in p.glob("*/command.py"):
+            module_name = filename(cmd_file)
+            cmds[module_name] = (
+                module_name.split(".", 1)[0],
+                str(cmd_file) if shim else None,
+            )
+
+        return cmds
+
+    def load_from_path(self, path):
+        cmds = self.find_in_path(path)
+        for module_name in cmds:
+            (cmd_name, file) = cmds[module_name]
+            if file:
+                spec = importlib.util.spec_from_file_location(
+                    f"shim.{module_name}", file
+                )
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+            else:
+                module = importlib.import_module(f".{module_name}", "oo_bin")
+
+            if hasattr(module, "replaces"):
+                replaces = getattr(module, "replaces")
+                for name in replaces:
+                    if hasattr(module, replaces[name]):
+                        self.commands[name] = getattr(module, replaces[name])
+                    else:
+                        raise Exception(
+                            f"ERROR: could not load `{replaces[name]}` from `{module_name}`"
+                        )
+
+            if hasattr(module, cmd_name):
+                self.commands[module_name] = getattr(module, cmd_name)
+
+    def register(self, cli):
+        for cmd in self.commands:
+            cli.add_command(self.commands[cmd])
+
+    def load_from_paths(self, *paths):
+        # these paths are assumed to be shims as the canonical commands are registered above
+        for path in paths:
+            self.load_from_path(path)

oo_bin/main.py

@@ -10,14 +10,9 @@
 from xdg import BaseDirectory
 
 from oo_bin import __version__
+from oo_bin.commander import Commander
 from oo_bin.config import main_config
-from oo_bin.dnsme.command import dnsme
 from oo_bin.errors import OOBinError
-from oo_bin.hexme.command import hexme
-from oo_bin.macme.command import macme
-from oo_bin.ping.command import ping
-from oo_bin.ssh.command import ssh
-from oo_bin.tunnels.command import rdp, tunnels, vnc
 from oo_bin.utils import auto_update, update_package, update_tunnels_config
 
 dsn = main_config().get("sentry", {}).get("dsn", None)
@@ -45,19 +40,14 @@ def cli(ctx, update):
         return None
 
 
-cli.add_command(dnsme)
-cli.add_command(hexme)
-cli.add_command(macme)
-cli.add_command(ping)
-cli.add_command(rdp)
-cli.add_command(ssh)
-cli.add_command(tunnels)
-cli.add_command(vnc)
-
-
 def main():
     try:
         auto_update()
+        cmds = Commander(os.path.dirname(__file__))
+        cmds.load_from_path(
+            os.path.join(BaseDirectory.save_data_path("oo_bin"), "shims")
+        )
+        cmds.register(cli)
         cli()
     except OOBinError as e:
         # Handled errors, these errors should not be uploaded to sentry