Root plugins (#35)

* root plugins

plugins can now specify if they want their methods to be ran as root. this is done via the multiprocess module. method calls are delegated to a separate process that is then down-privileged by default to user 1000, so the loader can safely be ran as root

except it isn't really safe because the plugin is imported as root anyway

* working implementation

- follows the new plugin format with the plugin.json file
- plugins are loaded in their own isolated process along with their own event loop and unix socket server for calling methods
- private methods are now prepended with _ instead of __

* converted format to f-strings

3 files changed, 105 insertions, 29 deletions

diff --git a/plugin_loader/loader.py b/plugin_loader/loader.py
index 20a05096..9523d73d 100644
--- a/plugin_loader/loader.py
+++ b/plugin_loader/loader.py
@@ -4,10 +4,10 @@ from watchdog.observers.polling import PollingObserver as Observer
 from watchdog.events import FileSystemEventHandler
 
 from os import path, listdir
-from importlib.util import spec_from_file_location, module_from_spec
 from logging import getLogger
 
 from injector import get_tabs, get_tab
+from plugin import PluginWrapper
 
 class FileChangeHandler(FileSystemEventHandler):
     def __init__(self, loader, plugin_path) -> None:
@@ -67,7 +67,6 @@ class Loader:
         server_instance.add_routes([
             web.get("/plugins/iframe", self.plugin_iframe_route),
             web.get("/plugins/reload", self.reload_plugins),
-            web.post("/plugins/method_call", self.handle_plugin_method_call),
             web.get("/plugins/load_main/{name}", self.load_plugin_main_view),
             web.get("/plugins/plugin_resource/{name}/{path:.+}", self.handle_sub_route),
             web.get("/plugins/load_tile/{name}", self.load_plugin_tile_view),
@@ -76,28 +75,16 @@ class Loader:
 
     def import_plugin(self, file, plugin_directory, refresh=False):
         try:
-            spec = spec_from_file_location("_", file)
-            module = module_from_spec(spec)
-            spec.loader.exec_module(module)
-
-            # add member for what directory the given plugin lives under
-            module.Plugin._plugin_directory = plugin_directory
-
-            if not hasattr(module.Plugin, "name"):
-                raise KeyError(f"Plugin {file} has not defined a name")
-            if module.Plugin.name in self.plugins:
-                    if hasattr(module.Plugin, "hot_reload") and not module.Plugin.hot_reload  and refresh:
-                        self.logger.info(f"Plugin {module.Plugin.name} is already loaded and has requested to not be re-loaded")
+            plugin = PluginWrapper(file, plugin_directory, self.plugin_path)
+            if plugin.name in self.plugins:
+                    if not "hot_reload" in plugin.flags and refresh:
+                        self.logger.info(f"Plugin {plugin.name} is already loaded and has requested to not be re-loaded")
                         return
                     else:
-                        if hasattr(self.plugins[module.Plugin.name], "task"):
-                            self.plugins[module.Plugin.name].task.cancel()
-                        self.plugins.pop(module.Plugin.name, None)
-            self.plugins[module.Plugin.name] = module.Plugin()
-            if hasattr(module.Plugin, "__main"):
-                setattr(self.plugins[module.Plugin.name], "task",
-                self.loop.create_task(self.plugins[module.Plugin.name].__main()))
-            self.logger.info(f"Loaded {module.Plugin.name}")
+                        self.plugins[plugin.name].stop(self.loop)
+                        self.plugins.pop(plugin.name, None)
+            self.plugins[plugin.name] = plugin.start(self.loop)
+            self.logger.info(f"Loaded {plugin.name}")
         except Exception as e:
             self.logger.error(f"Could not load {file}. {e}")
         finally:
@@ -117,9 +104,9 @@ class Loader:
         self.import_plugins()
 
     async def handle_plugin_method_call(self, plugin_name, method_name, **kwargs):
-        if method_name.startswith("__"):
+        if method_name.startswith("_"):
             raise RuntimeError("Tried to call private method")
-        return await getattr(self.plugins[plugin_name], method_name)(**kwargs)
+        return await self.plugins[plugin_name].execute_method(method_name, kwargs)
 
     async def get_steam_resource(self, request):
         tab = (await get_tabs())[0]
@@ -132,7 +119,7 @@ class Loader:
         plugin = self.plugins[request.match_info["name"]]
 
         # open up the main template
-        with open(path.join(self.plugin_path, plugin._plugin_directory, plugin.main_view_html), 'r') as template:
+        with open(path.join(self.plugin_path, plugin.plugin_directory, plugin.main_view_html), 'r') as template:
             template_data = template.read()
             # setup the main script, plugin, and pull in the template
             ret = f"""
@@ -150,7 +137,7 @@ class Loader:
 
         ret = ""
 
-        file_path = path.join(self.plugin_path, plugin._plugin_directory, route_path)
+        file_path = path.join(self.plugin_path, plugin.plugin_directory, route_path)
         with open(file_path, 'r') as resource_data:
             ret = resource_data.read()
 
@@ -162,8 +149,8 @@ class Loader:
         inner_content = ""
 
         # open up the tile template (if we have one defined)
-        if len(plugin.tile_view_html) > 0:
-            with open(path.join(self.plugin_path, plugin._plugin_directory, plugin.tile_view_html), 'r') as template:
+        if hasattr(plugin, "tile_view_html"):
+            with open(path.join(self.plugin_path, plugin.plugin_directory, plugin.tile_view_html), 'r') as template:
                 template_data = template.read()
                 inner_content = template_data
         
diff --git a/plugin_loader/plugin.py b/plugin_loader/plugin.py
new file mode 100644
index 00000000..cd737c72
--- /dev/null
+++ b/plugin_loader/plugin.py
@@ -0,0 +1,88 @@
+from importlib.util import spec_from_file_location, module_from_spec
+from asyncio import get_event_loop, start_unix_server, open_unix_connection, sleep, Lock
+from os import path, setuid
+from json import loads, dumps, load
+from concurrent.futures import ProcessPoolExecutor
+from time import time
+
+class PluginWrapper:
+    def __init__(self, file, plugin_directory, plugin_path) -> None:
+        self.file = file
+        self.plugin_directory = plugin_directory
+        self.reader = None
+        self.writer = None
+        self.socket_addr = f"/tmp/plugin_socket_{time()}"
+        self.method_call_lock = Lock()
+
+        json = load(open(path.join(plugin_path, plugin_directory, "plugin.json"), "r"))
+
+        self.name = json["name"]
+        self.author = json["author"]
+        self.main_view_html = json["main_view_html"]
+        self.tile_view_html = json["tile_view_html"] if "tile_view_html" in json else ""
+        self.flags = json["flags"]
+
+    def _init(self):
+        setuid(0 if "root" in self.flags else 1000)
+        spec = spec_from_file_location("_", self.file)
+        module = module_from_spec(spec)
+        spec.loader.exec_module(module)
+        self.Plugin = module.Plugin
+
+        if hasattr(self.Plugin, "_main"):
+            get_event_loop().create_task(self.Plugin._main(self.Plugin))
+        get_event_loop().create_task(self._setup_socket())
+        get_event_loop().run_forever()
+
+    async def _setup_socket(self):
+        self.socket = await start_unix_server(self._listen_for_method_call, path=self.socket_addr)
+
+    async def _listen_for_method_call(self, reader, writer):
+        while True:
+            data = loads((await reader.readline()).decode("utf-8"))
+            if "stop" in data:
+                return get_event_loop().stop()
+            d = {"res": None, "success": True}
+            try:
+                d["res"] = await getattr(self.Plugin, data["method"])(self.Plugin, **data["args"])
+            except Exception as e:
+                d["res"] = str(e)
+                d["success"] = False
+            finally:
+                writer.write((dumps(d)+"\n").encode("utf-8"))
+                await writer.drain()
+
+    async def _open_socket_if_not_exists(self):
+        if not self.reader:
+            while True:
+                try:
+                    self.reader, self.writer = await open_unix_connection(self.socket_addr)
+                    break
+                except:
+                    await sleep(0)
+
+    def start(self, loop):
+        executor = ProcessPoolExecutor()
+        loop.run_in_executor(
+            executor,
+            self._init
+        )
+        return self
+
+    def stop(self, loop):
+        async def _(self):
+            await self._open_socket_if_not_exists()
+            self.writer.write((dumps({"stop": True})+"\n").encode("utf-8"))
+            await self.writer.drain()
+        loop.create_task(_(self))
+
+    async def execute_method(self, method_name, kwargs):
+        async with self.method_call_lock:
+            await self._open_socket_if_not_exists()
+            self.writer.write(
+                (dumps({"method": method_name, "args": kwargs})+"\n").encode("utf-8"))
+            await self.writer.drain()
+            res = loads((await self.reader.readline()).decode("utf-8"))
+            if not res["success"]:
+                raise Exception(res["res"])
+            return res["res"]
diff --git a/requirements.txt b/requirements.txt
index c77a53ed..579ebc0b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 aiohttp==3.8.1
 aiohttp-jinja2==1.5.0
-watchdog==2.1.7
\ No newline at end of file
+watchdog==2.1.7
+multiprocess==0.70.12.2
\ No newline at end of file