Key value storage logic and auth

.gitignore

@@ -6,6 +6,9 @@ __pycache__/
 # C extensions
 *.so
 
+# Local k-v store
+*.db
+
 # Distribution / packaging
 .Python
 build/
@@ -153,7 +156,7 @@ cython_debug/
 .LSOverride
 
 # Icon must end with two \r
-Icon
+Icon
 
 # Thumbnails
 ._*

api_reflector/actions.py

@@ -3,11 +3,13 @@
 """
 import json
 import time
+from datetime import datetime, timedelta
 from enum import Enum
 
 import requests
 
 from api_reflector.reporting import get_logger
+from api_reflector.storage import set_value_in_storage
 from settings import settings
 
 log = get_logger(__name__)
@@ -20,6 +22,7 @@ class Action(Enum):
 
     DELAY = "DELAY"
     CALLBACK = "CALLBACK"
+    STORE = "STORE"
 
     def __str__(self) -> str:
         return self.value
@@ -64,7 +67,29 @@ def process_callback(*args, **kwargs):
         log.warning(f"Check all Action arguments have been provided and that the callback service is running`{ex}`")
 
 
+def set_value(*args, **kwargs):
+    """
+    Sets value for key specified in args into the storage object
+    """
+    key = args[0]
+    value = args[1]
+    endpoint = kwargs["endpoint"].rsplit("/", 1)[0]
+
+    if settings.session_timeout:
+        # FIXME: Could also allow expiry via action args. i.e delta = args[2]
+        delta = settings.session_timeout
+        expiry = datetime.now() + timedelta(minutes=int(delta))
+        storage_value = {key: value, "expiry": expiry.timestamp()}
+    else:
+        storage_value = {
+            key: value,
+        }
+
+    set_value_in_storage(key=endpoint, value=storage_value)
+
+
 action_executors = {
     Action.DELAY: delay,
     Action.CALLBACK: process_callback,
+    Action.STORE: set_value,
 }

api_reflector/migrations/script.py.mako

@@ -16,9 +16,9 @@ branch_labels = ${repr(branch_labels)}
 depends_on = ${repr(depends_on)}
 
 
-def upgrade():
+def upgrade() -> None:
     ${upgrades if upgrades else "pass"}
 
 
-def downgrade():
+def downgrade() -> None:
     ${downgrades if downgrades else "pass"}

api_reflector/migrations/versions/f0be9a5e8cec_storage.py

@@ -0,0 +1,26 @@
+"""storage
+
+Revision ID: f0be9a5e8cec
+Revises: d42bddfac4f1
+Create Date: 2022-11-04 20:22:43.252600
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f0be9a5e8cec"
+down_revision = "d42bddfac4f1"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.get_context().autocommit_block():
+        op.execute("ALTER TYPE action ADD VALUE 'STORE'")
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    pass
+    # ### end Alembic commands ###

api_reflector/models.py

@@ -100,7 +100,9 @@ def execute_actions(self, req_json: Mapping[str, Any], content: str):
 
         for action in self.actions:
             log.debug(f"Executing action: {action}")
-            actions.action_executors[action.action](*action.arguments, request=req_json, response=content)
+            actions.action_executors[action.action](
+                *action.arguments, request=req_json, response=content, endpoint=self.endpoint.path
+            )
 
 
 class Rule(Model):
@@ -155,6 +157,7 @@ def __str__(self) -> str:
         action_str = {
             actions.Action.DELAY: "Delay for {} second(s)",
             actions.Action.CALLBACK: "Do the callback",
+            actions.Action.STORE: "Store key-value pair in storage",
         }[self.action]
         return action_str.format(*self.arguments)
 

api_reflector/rules_engine.py

@@ -7,6 +7,8 @@
 from enum import Enum
 from typing import Any, Callable, Mapping, NamedTuple, TypeVar, Union
 
+from werkzeug import exceptions
+
 from api_reflector.templating import default_context, template_env
 
 
@@ -71,6 +73,7 @@ class TemplatableRequest(NamedTuple):
     json: Mapping[str, Any]
     query: Mapping[str, Any]
     headers: Mapping[str, Any]
+    storage: Mapping[str, bool] | None = None
 
 
 class ScoringRule(NamedTuple):
@@ -101,7 +104,10 @@ def score_response(request: TemplatableRequest, rules: list[ScoringRule]) -> flo
         args = [template_env.from_string(arg).render(**template_context) for arg in rule.arguments]
         evaluator = evaluators[rule.operator]
         if not evaluator(*args):
-            return -1
+            for argument in rule.arguments:
+                if "authorized" in argument:
+                    raise exceptions.Unauthorized
+                return -1
     return len(rules)
 
 
@@ -122,6 +128,7 @@ def find_best_response(
         )
         for response, rules in response_rules
     ]
+
     scores = [(score, response) for score, response in scores if score >= 0]
     # sort by score
     scores = sorted(scores, key=lambda score: score[0], reverse=True)

api_reflector/storage.py

@@ -0,0 +1,16 @@
+"""
+Handle connections to supported k-v storage databases
+"""
+import pickledb
+
+pickled = pickledb.load("pickle.db", auto_dump=True, sig=True)
+
+
+def set_value_in_storage(key, value):
+    """Sets key value in storage"""
+    pickled.set(key, value)
+
+
+def get_value_from_storage(key):
+    """Gets value for a key from storage"""
+    return pickled.get(key)

api_reflector/views.py

@@ -1,6 +1,8 @@
+# views.py
 """
 Defines the project's API endpoints.
 """
+from datetime import datetime
 from typing import Any, Mapping
 
 import psycopg2
@@ -13,12 +15,23 @@
 from api_reflector.auth import requires_auth
 from api_reflector.endpoint import ensure_leading_slash
 from api_reflector.reporting import get_logger
+from api_reflector.storage import get_value_from_storage
 from api_reflector.templating import default_context, template_env
 
 api = Blueprint("api", __name__)
 log = get_logger(__name__)
 
 
+def _get_storage_obj(req_endpoint_path: str) -> dict:
+    storage_val = get_value_from_storage(req_endpoint_path)
+    if storage_val:
+        expiry = storage_val.get("expiry")
+        if expiry and expiry < datetime.now().timestamp():
+            return {}
+        return storage_val
+    return {}
+
+
 def match_endpoint(path: str) -> tuple[models.Endpoint, Mapping[str, Any]]:
     """
     Uses werkzeug routing to match the given path to an endpoint.
@@ -86,6 +99,7 @@ def home() -> tuple[Any, int]:
     return render_template("home.html", endpoints=endpoints, tags=tags), 200
 
 
+@api.route("/mock/auth", methods=["POST"])
 @api.route("/mock/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
 def mock(path: str) -> Response:
     """
@@ -102,6 +116,7 @@ def mock(path: str) -> Response:
     if not active_responses:
         return Response("No Mock Responses configured or active for this endpoint", status=501)
 
+    storage_obj = _get_storage_obj(endpoint.path)
     response_rules = [
         (
             response,
@@ -122,7 +137,7 @@ def mock(path: str) -> Response:
         req_json = {}
 
     templateable_request = rules_engine.TemplatableRequest(
-        params=params, json=req_json, query=request.args, headers=request.headers
+        params=params, json=req_json, query=request.args, headers=request.headers, storage=storage_obj
     )
 
     try: