From d3c63a5b088dc897addeb6db2c414ac14820e9ee Mon Sep 17 00:00:00 2001
From: ujjwal96 <ujjwalverma1111@gmail.com>
Date: Wed, 6 May 2020 15:50:58 -0400
Subject: [PATCH] Allow relative and absolute paths for scripts

---
 README.md           |  6 +++---
 simserve/handler.py | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index f53ea47..ebe0c53 100644
--- a/README.md
+++ b/README.md
@@ -26,16 +26,16 @@ To configure the tool copy `conf.json.example` to `conf.json`.
     - `body` - The response body for given route. Used if `"response_type` is "static"
     - `methods` - HTTP methods for which this response should be sent.
     - `response_type` - How to generate the response. Can be "static" or "script"
-    - `script` - Name of python module from which to generate the response. Used if `"response_type` is "script"
+    - `script` - Name of python module from which to generate the response. Used if `"response_type` is "script". User can provide either a default plugin(`customresp`), a relative path to python script or an absolute path to the script.
     
 To run - `python3 server.py`
      
 ### Plugins
-Users can create their own plugins which will allow them to generate dynamic responses based on the requests received. All the plugins reside in the `plugins` directory.
+Users can create their own plugins which will allow them to generate dynamic responses based on the requests received.
 
 To create a plugin you need to import the base plugin class and override its abstract method.  
 ```python
-from plugins.base import Plugin
+from simserve.plugins.base import Plugin
 
 class MyResponse(Plugin):
     def response(self, request):
diff --git a/simserve/handler.py b/simserve/handler.py
index 7eeabc7..f4021b1 100644
--- a/simserve/handler.py
+++ b/simserve/handler.py
@@ -1,5 +1,8 @@
+import os
+import sys
 import inspect
 import importlib
+import importlib.util
 from typing import (
     Dict,
     Union,
@@ -22,9 +25,41 @@ def _get_resp(self, method: str) -> Union[Dict, None]:
                 return resp
         return None
 
+    @staticmethod
+    def _load_script(path: str):
+        if importlib.util.find_spec(f"simserve.plugins.{os.path.splitext(os.path.basename(path))[0]}"):
+            return importlib.import_module(f"simserve.plugins.{path}")
+        elif os.path.isfile(f"{os.getcwd()}/{path}.py") or os.path.isfile(f"{os.getcwd()}/{path}"):
+            path = f"{os.getcwd()}/{path}"
+            path = f"{path}.py" if os.path.isfile(f"{path}.py") else f"{path}"
+            module_name = f"__simserve_plugin__.{os.path.splitext(os.path.basename(path))[0]}"
+            try:
+                spec = importlib.util.spec_from_file_location(module_name, path)
+                module = importlib.util.module_from_spec(spec)
+                sys.modules[module_name] = module
+                spec.loader.exec_module(module)
+            except Exception as e:
+                access_log.error(f"Could not load plugin. {e}")
+                raise tornado.web.HTTPError(500)
+            return module
+        elif os.path.isfile(path):
+            module_name = f"__simserve_plugin__.{os.path.splitext(os.path.basename(path))[0]}"
+            try:
+                spec = importlib.util.spec_from_file_location(module_name, path)
+                module = importlib.util.module_from_spec(spec)
+                sys.modules[module_name] = module
+                spec.loader.exec_module(module)
+            except Exception as e:
+                access_log.error(f"Could not load plugin. {e}")
+                raise tornado.web.HTTPError(500)
+            return module
+        else:
+            access_log.error("Couldn't find plugin to load.")
+            raise tornado.web.HTTPError(500)
+
     def _get_resp_body(self, resp: Dict):
         if resp.get("response_type") == "script":
-            plug_module = importlib.import_module(f"simserve.plugins.{resp.get('script')}")
+            plug_module = self._load_script(resp.get('script'))
             _, plug = inspect.getmembers(
                 plug_module,
                 lambda x: inspect.isclass(x) and not inspect.isabstract(x) and issubclass(x, Plugin)
@@ -49,7 +84,6 @@ def _format_response(resp) -> str:
         log = f"\n----- Response -----\n"
         for k, v in sorted(resp.get("headers")):
             log += f"{k}: {v}\n"
-        # TODO: Log for script responses.
         log += f"\n{resp.get('body')}\n"
         log += f"----- End -----\n"
         return log