Expose new public api for rate limiting and user blocking

aikido_zen/__init__.py

@@ -6,8 +6,9 @@
 
 from dotenv import load_dotenv
 
-# Re-export set_current_user :
+# Re-export functions :
 from aikido_zen.context.users import set_user
+from aikido_zen.middleware import should_block_request
 
 # Import logger
 from aikido_zen.helpers.logging import logger

aikido_zen/context/__init__.py

@@ -63,6 +63,8 @@ def __init__(self, context_obj=None, body=None, req=None, source=None):
         self.route_params = extract_route_params(self.url)
         self.subdomains = get_subdomains_from_url(self.url)
 
+        self.executed_middleware = False
+
     def __reduce__(self):
         return (
             self.__class__,
@@ -81,6 +83,7 @@ def __reduce__(self):
                     "user": self.user,
                     "xml": self.xml,
                     "outgoing_req_redirects": self.outgoing_req_redirects,
+                    "executed_middleware": self.executed_middleware,
                     "route_params": self.route_params,
                 },
                 None,

aikido_zen/context/init_test.py

@@ -68,6 +68,7 @@ def test_wsgi_context_1():
         "parsed_userinput": {},
         "xml": {},
         "outgoing_req_redirects": [],
+        "executed_middleware": False,
         "route_params": [],
     }
 
@@ -95,6 +96,7 @@ def test_wsgi_context_2():
         "parsed_userinput": {},
         "xml": {},
         "outgoing_req_redirects": [],
+        "executed_middleware": False,
         "route_params": [],
     }
 

aikido_zen/context/users.py

@@ -20,6 +20,12 @@ def set_user(user):
     if not context:
         logger.debug("No context set, returning")
         return
+    if context.executed_middleware is True:
+        # Middleware to rate-limit/check for users ran already. Could be misconfiguration.
+        logger.warning(
+            "set_user(...) must be called before the Zen middleware is executed."
+        )
+
     validated_user["lastIpAddress"] = context.remote_address
     context.user = validated_user
 

aikido_zen/context/users_test.py

@@ -1,6 +1,41 @@
+from lib2to3.fixes.fix_input import context
+
 import pytest
 
+from . import current_context, Context
 from .users import validate_user, set_user
+from .. import should_block_request
+
+
+@pytest.fixture(autouse=True)
+def run_around_tests():
+    yield
+    # Make sure to reset context and cache after every test so it does not
+    # interfere with other tests
+    current_context.set(None)
+
+
+def set_context_and_lifecycle():
+    wsgi_request = {
+        "REQUEST_METHOD": "GET",
+        "HTTP_HEADER_1": "header 1 value",
+        "HTTP_HEADER_2": "Header 2 value",
+        "RANDOM_VALUE": "Random value",
+        "HTTP_COOKIE": "sessionId=abc123xyz456;",
+        "wsgi.url_scheme": "http",
+        "HTTP_HOST": "localhost:8080",
+        "PATH_INFO": "/hello",
+        "QUERY_STRING": "user=JohnDoe&age=30&age=35",
+        "CONTENT_TYPE": "application/json",
+        "REMOTE_ADDR": "198.51.100.23",
+    }
+    context = Context(
+        req=wsgi_request,
+        body=None,
+        source="flask",
+    )
+    context.set_as_current_context()
+    return context
 
 
 def test_validate_user_valid_input():
@@ -67,3 +102,57 @@ def test_validate_user_invalid_user_type_dict_without_id(caplog):
 def test_set_user_with_none(caplog):
     result = set_user(None)
     assert "expects a dict with 'id' and 'name' properties" in caplog.text
+
+
+def test_set_valid_user():
+    context1 = set_context_and_lifecycle()
+    assert context1.user is None
+
+    user = {"id": 456, "name": "Bob"}
+    set_user(user)
+
+    assert context1.user == {
+        "id": "456",
+        "name": "Bob",
+        "lastIpAddress": "198.51.100.23",
+    }
+
+
+def test_re_set_valid_user():
+    context1 = set_context_and_lifecycle()
+    assert context1.user is None
+
+    user = {"id": 456, "name": "Bob"}
+    set_user(user)
+
+    assert context1.user == {
+        "id": "456",
+        "name": "Bob",
+        "lastIpAddress": "198.51.100.23",
+    }
+
+    user = {"id": "1000", "name": "Alice"}
+    set_user(user)
+
+    assert context1.user == {
+        "id": "1000",
+        "name": "Alice",
+        "lastIpAddress": "198.51.100.23",
+    }
+
+
+def test_after_middleware(caplog):
+    context1 = set_context_and_lifecycle()
+    assert context1.user is None
+
+    should_block_request()
+
+    user = {"id": 456, "name": "Bob"}
+    set_user(user)
+
+    assert "must be called before the Zen middleware is executed" in caplog.text
+    assert context1.user == {
+        "id": "456",
+        "name": "Bob",
+        "lastIpAddress": "198.51.100.23",
+    }

aikido_zen/middleware/__init__.py

@@ -0,0 +1,7 @@
+"""Re-exports middleware"""
+
+from .asgi import AikidoASGIMiddleware as AikidoQuartMiddleware
+from .asgi import AikidoASGIMiddleware as AikidoStarletteMiddleware
+from .flask import AikidoFlaskMiddleware
+from .django import AikidoDjangoMiddleware
+from .should_block_request import should_block_request

aikido_zen/middleware/asgi.py

@@ -0,0 +1,47 @@
+"""Exports ratelimiting and user blocking middleware for ASGI"""
+
+from aikido_zen.helpers.logging import logger
+from .should_block_request import should_block_request
+
+
+class AikidoASGIMiddleware:
+    """Ratelimiting and user blocking middleware for ASGI"""
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        result = should_block_request()
+        if result["block"] is not True:
+            return await self.app(scope, receive, send)
+
+        if result["type"] == "ratelimited":
+            message = "You are rate limited by Zen."
+            if result["trigger"] == "ip" and result["ip"]:
+                message += " (Your IP: " + result["ip"] + ")"
+            return await send_status_code_and_text(send, (message, 429))
+        elif result["type"] == "blocked":
+            return await send_status_code_and_text(
+                send, ("You are blocked by Zen.", 403)
+            )
+
+        logger.debug("Unknown type for blocking request: %s", result["type"])
+        return await self.app(scope, receive, send)
+
+
+async def send_status_code_and_text(send, pre_response):
+    """Sends a status code and text"""
+    await send(
+        {
+            "type": "http.response.start",
+            "status": pre_response[1],
+            "headers": [(b"content-type", b"text/plain")],
+        }
+    )
+    await send(
+        {
+            "type": "http.response.body",
+            "body": pre_response[0].encode("utf-8"),
+            "more_body": False,
+        }
+    )

aikido_zen/middleware/django.py

@@ -0,0 +1,37 @@
+"""Exports AikidoDjangoMiddleware"""
+
+from aikido_zen.helpers.logging import logger
+from .should_block_request import should_block_request
+
+
+class AikidoDjangoMiddleware:
+    """Middleware for rate-limiting and user blocking for django"""
+
+    def __init__(self, get_response):
+        logger.critical("Django middleware ised")
+        self.get_response = get_response
+        try:
+            from django.http import HttpResponse
+
+            self.HttpResponse = HttpResponse
+        except ImportError:
+            logger.warning(
+                "django.http import not working, aikido rate-limiting middleware not running."
+            )
+
+    def __call__(self, request):
+        result = should_block_request()
+        if result["block"] is not True or self.HttpResponse is None:
+            return self.get_response(request)
+
+        if result["type"] == "ratelimited":
+            message = "You are rate limited by Zen."
+            if result["trigger"] == "ip" and result["ip"]:
+                message += " (Your IP: " + result["ip"] + ")"
+            return self.HttpResponse(message, content_type="text/plain", status=429)
+        elif result["type"] == "blocked":
+            return self.HttpResponse(
+                "You are blocked by Zen.", content_type="text/plain", status=403
+            )
+        logger.debug("Unknown type for blocking request: %s", result["type"])
+        return self.get_response(request)

aikido_zen/middleware/flask.py

@@ -0,0 +1,38 @@
+"""Exports ratelimiting and user blocking middleware for Flask"""
+
+from aikido_zen.helpers.logging import logger
+from .should_block_request import should_block_request
+
+
+class AikidoFlaskMiddleware:
+    """Ratelimiting and user blocking middleware for Flask"""
+
+    def __init__(self, app):
+        self.app = app
+        try:
+            from werkzeug.wrappers import Response
+
+            self.Response = Response
+        except ImportError:
+            logger.warning(
+                "Something went wrong whilst importing werkzeug.wrappers, middleware does not work"
+            )
+
+    def __call__(self, environ, start_response):
+        result = should_block_request()
+        if result["block"] is not True or self.Response is None:
+            return self.app(environ, start_response)
+
+        if result["type"] == "ratelimited":
+            message = "You are rate limited by Zen."
+            if result["trigger"] == "ip" and result["ip"]:
+                message += " (Your IP: " + result["ip"] + ")"
+            res = self.Response(message, mimetype="text/plain", status=429)
+            return res(environ, start_response)
+        elif result["type"] == "blocked":
+            res = self.Response(
+                "You are blocked by Zen.", mimetype="text/plain", status=403
+            )
+            return res(environ, start_response)
+        logger.debug("Unknown type for blocking request: %s", result["type"])
+        return self.app(environ, start_response)