Better Logging, memory management | added result_url to response

Author: eshaan7

docs/source/Logging.md

@@ -12,10 +12,11 @@ Here's a snippet of code that shows how you can access this extension's logger o
 import logging
 # get the flask_shell2http logger
 logger = logging.getLogger("flask_shell2http")
-# log messages of severity DEBUG or lower to the console
+# create new handler
 handler = logging.StreamHandler(sys.stdout)
 logger.addHandler(handler)
-logger.setLevel(logging.DEBUG)
+# log messages of severity DEBUG or lower to the console
+logger.setLevel(logging.DEBUG)  # this is really important!
 ```
 
 Please consult the Flask's official docs on 

docs/source/Quickstart.md

@@ -56,15 +56,16 @@ returns JSON,
 
 ```json
 {
-   "key": "ddbe0a94847c65f9b8198424ffd07c50",
+   "key": "ddbe0a94847c",
+   "result_url": "http://localhost:4000/commands/saythis?key=ddbe0a94847c",
    "status": "running"
 }
 ```
 
-Then using this `key` you can query for the result,
+Then using this `key` you can query for the result or just by going to the `result_url`,
 
 ```bash
-$ curl http://localhost:4000/commands/saythis?key=ddbe0a94847c65f9b8198424ffd07c50
+$ curl http://localhost:4000/commands/saythis?key=ddbe0a94847c
 ```
 
 Returns result in JSON,
@@ -73,7 +74,7 @@ Returns result in JSON,
 {
   "end_time": 1593019807.782958, 
   "error": "", 
-  "md5": "ddbe0a94847c65f9b8198424ffd07c50", 
+  "md5": "ddbe0a94847c", 
   "process_time": 0.00748753547668457, 
   "report": {
     "result": "Hello World!\n"

docs/source/conf.py

@@ -26,7 +26,7 @@
 author = "Eshaan Bansal"
 
 # The full version, including alpha/beta/rc tags
-release = "1.1.3"
+release = "1.2.0"
 
 
 # -- General configuration ---------------------------------------------------

flask_shell2http/api.py

@@ -19,76 +19,81 @@
 
 
 logger = get_logger()
+store = ReportStore()
+request_parser = RequestParser()
 
 
 class shell2httpAPI(MethodView):
     """
-    Flask.MethodView that defines GET and POST methods for an URL rule.
-    This is invoked on `Shell2HTTP.register_command`, in which case
-    the URL rule is the given endpoint.
+    Flask.MethodView that registers GET and POST methods for a given endpoint.
+    This is invoked on `Shell2HTTP.register_command`.
     Internal use only.
     """
 
-    command_name: str
-    executor: JobExecutor
-    store: ReportStore
-    request_parser: RequestParser
-
     def get(self):
         try:
             key: str = request.args.get("key")
-            logger.info(f"Report requested for key:'{key}'.")
+            logger.info(
+                f"Job: '{key}' --> Report requested. "
+                f"Requester: '{request.remote_addr}'."
+            )
             if not key:
                 raise Exception("No key provided in arguments.")
             # check if job has been finished
             future = self.executor.get_job(key)
             if future:
                 if not future.done:
+                    logger.debug(f"Job: '{key}' --> still running.")
                     return make_response(jsonify(status="running", key=key), 200)
 
                 # pop future object since it has been finished
                 self.executor.pop_job(key)
 
             # if yes, get result from store
-            report = self.store.get_one(key)
+            report = store.pop_and_get_one(key)
             if not report:
-                raise Exception(f"Report does not exist for key:{key}.")
+                raise Exception(f"No report exists for key: '{key}'.")
 
             resp = report.to_dict()
-            logger.debug(f"Requested report: {resp}")
+            logger.debug(f"Job: '{key}' --> Requested report: {resp}")
             return make_response(jsonify(resp), HTTPStatus.OK)
 
         except Exception as e:
-            logger.exception(e)
+            logger.error(e)
             return make_response(jsonify(error=str(e)), HTTPStatus.NOT_FOUND)
 
     def post(self):
         try:
-            logger.info(f"Received request for endpoint: '{request.url_rule}'.")
+            logger.info(
+                f"Received request for endpoint: '{request.url_rule}'. "
+                f"Requester: '{request.remote_addr}'."
+            )
             # Check if command is correct and parse it
-            cmd, key = self.request_parser.parse_req(request)
+            cmd, key = request_parser.parse_req(request, self.command_name)
 
             # run executor job in background
-            job_key = JobExecutor.make_key(key)
             future = self.executor.new_job(
-                future_key=job_key, fn=self.executor.run_command, cmd=cmd, key=key
+                future_key=JobExecutor.make_key(key),
+                fn=self.executor.run_command,
+                cmd=cmd,
+                key=key,
             )
             # callback that adds result to store
-            future.add_done_callback(self.store.save_result)
+            future.add_done_callback(store.save_result)
             # callback that removes the temporary directory
-            future.add_done_callback(self.request_parser.cleanup_temp_dir)
+            future.add_done_callback(request_parser.cleanup_temp_dir)
 
-            logger.info(f"Job: '{job_key}' added to queue for command: {cmd}")
+            logger.info(f"Job: '{key}' --> added to queue for command: {cmd}")
+            result_url = f"{request.base_url}?key={key}"
             return make_response(
-                jsonify(status="running", key=key), HTTPStatus.ACCEPTED,
+                jsonify(status="running", key=key, result_url=result_url),
+                HTTPStatus.ACCEPTED,
             )
 
         except Exception as e:
-            logger.exception(e)
+            logger.error(e)
             return make_response(jsonify(error=str(e)), HTTPStatus.BAD_REQUEST)
 
-    def __init__(self, command_name, executor):
-        self.command_name = command_name
-        self.executor = JobExecutor(executor)
-        self.store = ReportStore()
-        self.request_parser = RequestParser(command_name)
+    def __init__(self, command_name, job_executor):
+        self.command_name: str = command_name
+        self.executor: JobExecutor = job_executor

flask_shell2http/base_entrypoint.py

@@ -2,6 +2,7 @@
 from collections import OrderedDict
 
 # lib imports
+from .classes import JobExecutor
 from .api import shell2httpAPI
 from .helpers import get_logger
 
@@ -49,7 +50,7 @@ def init_app(self, app, executor) -> None:
            https://flask.palletsprojects.com/en/1.1.x/patterns/appfactories/
         """
         self.app = app
-        self.__executor = executor
+        self.__executor = JobExecutor(executor)
         self.__init_extension()
 
     def __init_extension(self) -> None:
@@ -84,26 +85,33 @@ def register_command(self, endpoint: str, command_name: str) -> None:
                 endpoint="myawesomescript", command_name="./fuxsocy.py"
             )
         """
-        if not self.__commands.get(endpoint):
-            url = self.__construct_route(endpoint)
-            self.app.add_url_rule(
-                url,
-                view_func=shell2httpAPI.as_view(
-                    command_name, command_name=command_name, executor=self.__executor
-                ),
+        uri = self.__construct_route(endpoint)
+        # make sure the given endpoint is not already registered
+        cmd_already_exists = self.__commands.get(uri)
+        if cmd_already_exists:
+            logger.error(
+                "Failed to register since given endpoint: "
+                f"'{endpoint}' already maps to command: '{cmd_already_exists}'."
             )
-            self.__commands.update({command_name: url})
-            logger.info(
-                f"New endpoint: '{endpoint}' registered for command: '{command_name}'."
-            )
-
-    def get_registered_commands(self):
+            return None
+
+        # else, add new URL rule
+        self.app.add_url_rule(
+            uri,
+            view_func=shell2httpAPI.as_view(
+                command_name, command_name=command_name, job_executor=self.__executor,
+            ),
+        )
+        self.__commands.update({uri: command_name})
+        logger.info(f"New URI: '{uri}' registered for command: '{command_name}'.")
+
+    def get_registered_commands(self) -> "OrderedDict[str, str]":
         """
         Most of the time you won't need this since
         Flask provides a ``Flask.url_map`` attribute.
 
         Returns:
-            OrderedDict i.e. mapping of registered commands and their URLs.
+            OrderedDict[uri, command] i.e. mapping of registered commands and their URLs.
         """
         return self.__commands
 

flask_shell2http/classes.py

@@ -2,7 +2,9 @@
 import time
 import subprocess
 import tempfile
+import shutil
 from collections import OrderedDict
+from typing import List, Dict
 
 # web imports
 from flask import safe_join
@@ -74,7 +76,6 @@ def run_command(self, cmd, key) -> Report:
         :returns:
             A ConcurrentFuture object where future.result = Report()
         """
-        job_key: str = self.make_key(key)
         start_time = time.time()
         try:
             proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -88,7 +89,7 @@ def run_command(self, cmd, key) -> Report:
             else:
                 status = "failed"
 
-            logger.info(f"Job: '{job_key}' finished with status: '{status}'.")
+            logger.info(f"Job: '{key}' --> finished with status: '{status}'.")
             return Report(
                 key=key,
                 report=stdout,
@@ -100,7 +101,7 @@ def run_command(self, cmd, key) -> Report:
         except Exception as e:
             str_err = str(e)
             self.cancel_job(key)
-            logger.error(f"Job: '{job_key}' failed. Reason: {str_err}.")
+            logger.error(f"Job: '{key}' --> failed. Reason: {str_err}.")
             return Report(
                 key=key,
                 report=None,
@@ -128,8 +129,11 @@ def save_result(self, future) -> None:
     def get_all(self):
         return self.__results
 
-    def get_one(self, key) -> Report:
-        return self.__results.get(key)
+    def pop_and_get_one(self, key) -> Report:
+        try:
+            return self.__results.pop(key)
+        except KeyError:
+            return None
 
 
 class RequestParser:
@@ -138,8 +142,10 @@ class RequestParser:
     Internal use Only.
     """
 
+    __tmpdirs: Dict[str, str] = {}
+
     @staticmethod
-    def __parse_multipart_req(args: list, files):
+    def __parse_multipart_req(args: List[str], files) -> (List[str], str):
         # Check if file part exists
         fnames = []
         for arg in args:
@@ -154,7 +160,7 @@ def __parse_multipart_req(args: list, files):
             )
 
         # create a new temporary directory
-        tmpdir = tempfile.mkdtemp()
+        tmpdir: str = tempfile.mkdtemp()
         for fname in fnames:
             if fname not in files:
                 raise Exception(
@@ -173,33 +179,45 @@ def __parse_multipart_req(args: list, files):
         logger.debug(f"Request files saved under temp directory: '{tmpdir}'")
         return args, tmpdir
 
-    def parse_req(self, request) -> (str, str):
+    def parse_req(self, request, base_command: str) -> (str, str):
+        args: List[str] = []
+        tmpdir = None
         if request.is_json:
             # request does not contain a file
             args = request.json.get("args", [])
         elif request.files:
             # request contains file
             received_args = request.form.getlist("args")
-            args, self.tmpdir = RequestParser.__parse_multipart_req(
+            args, tmpdir = RequestParser.__parse_multipart_req(
                 received_args, request.files
             )
         else:
             # request is w/o any data
             # i.e. just run-script
             args = []
 
-        cmd: list = self.command_name.split(" ")
+        cmd: List[str] = base_command.split(" ")
         cmd.extend(args)
-        return cmd, calc_hash(cmd)
-
-    def cleanup_temp_dir(self, _) -> None:
-        if hasattr(self, "tmpdir"):
-            try:
-                __import__("shutil").rmtree(self.tmpdir)
-                logger.debug(f"temp directory: '{self.tmpdir}' was deleted.")
-                del self.tmpdir
-            except Exception:
-                logger.debug(f"Failed to clear temp directory: '{self.tmpdir}'.")
-
-    def __init__(self, command_name) -> None:
-        self.command_name = command_name
+        key: str = calc_hash(cmd)
+        if tmpdir:
+            self.__tmpdirs.update({key: tmpdir})
+
+        return cmd, key
+
+    def cleanup_temp_dir(self, future) -> None:
+        key: str = future.result().key
+        tmpdir: str = self.__tmpdirs.get(key, None)
+        if not tmpdir:
+            return None
+
+        try:
+            shutil.rmtree(tmpdir)
+            logger.debug(
+                f"Job: '{key}' --> Temporary directory: '{tmpdir}' "
+                "successfully deleted."
+            )
+            self.__tmpdirs.pop(key)
+        except Exception:
+            logger.debug(
+                f"Job: '{key}' --> Failed to clear Temporary directory: '{tmpdir}'."
+            )

flask_shell2http/helpers.py

@@ -1,7 +1,8 @@
 import logging
+from typing import List
 
 
-def list_replace(lst: list, old, new) -> None:
+def list_replace(lst: List, old, new) -> None:
     """
     replace list elements (inplace)
     """
@@ -14,15 +15,15 @@ def list_replace(lst: list, old, new) -> None:
         pass
 
 
-def calc_hash(lst: list) -> str:
+def calc_hash(lst: List) -> str:
     """
     Internal use only.
     Calculates sha1sum of given command with it's byte-string.
     This is for non-cryptographic purpose,
     that's why a faster and insecure hashing algorithm is chosen.
     """
     to_hash = " ".join(lst).encode("ascii")
-    return __import__("hashlib").sha1(to_hash).hexdigest()
+    return __import__("hashlib").sha1(to_hash).hexdigest()[:12]
 
 
 def get_logger() -> logging.Logger:

setup.py

@@ -17,7 +17,7 @@
 
 setup(
     name="Flask-Shell2HTTP",
-    version="1.1.3",
+    version="1.2.0",
     url=GITHUB_URL,
     license="BSD",
     author="Eshaan Bansal",