Merge branch 'main' into qa-tests

Author: bitterpanda63

aikido_zen/__init__.py

@@ -9,6 +9,7 @@
 # Re-export functions :
 from aikido_zen.context.users import set_user
 from aikido_zen.helpers.check_gevent import check_gevent
+from aikido_zen.helpers.python_version_not_supported import python_version_not_supported
 from aikido_zen.middleware import should_block_request
 from aikido_zen.middleware.set_rate_limit_group import set_rate_limit_group
 
@@ -34,6 +35,8 @@ def protect(mode="daemon", token=""):
     if aikido_disabled_flag_active():
         # Do not run any aikido code when the disabled flag is on
         return
+    if python_version_not_supported():
+        return
     if not test_uds_file_access():
         return  # Unable to start background process
     if check_gevent():
@@ -71,6 +74,7 @@ def protect(mode="daemon", token=""):
 
     import aikido_zen.sinks.builtins
     import aikido_zen.sinks.os
+    import aikido_zen.sinks.pathlib
     import aikido_zen.sinks.shutil
     import aikido_zen.sinks.io
     import aikido_zen.sinks.http_client

aikido_zen/helpers/python_version_not_supported.py

@@ -0,0 +1,14 @@
+import sys
+from aikido_zen.helpers.logging import logger
+
+
+def python_version_not_supported() -> bool:
+    major = sys.version_info.major
+    minor = sys.version_info.minor
+    if major != 3:
+        logger.error("This version of Zen only supports Python 3")
+        return True
+    if minor > 13:
+        logger.error("This version of Zen doesn't support versions above Python 3.13")
+        return True
+    return False

aikido_zen/middleware/init_test.py

@@ -1,8 +1,9 @@
 from unittest.mock import patch, MagicMock
 
 import pytest
-from aikido_zen.context import current_context, Context, get_current_context
-from aikido_zen.thread.thread_cache import ThreadCache, get_cache
+import aikido_zen.test_utils as test_utils
+from aikido_zen.context import current_context, get_current_context
+from aikido_zen.thread.thread_cache import get_cache
 from . import should_block_request
 from .. import set_rate_limit_group
 
@@ -22,35 +23,14 @@ def test_without_context():
     assert should_block_request() == {"block": False}
 
 
-def set_context(user=None, executed_middleware=False):
-    Context(
-        context_obj={
-            "remote_address": "::1",
-            "method": "POST",
-            "url": "http://localhost:4000",
-            "query": {
-                "abc": "def",
-            },
-            "headers": {},
-            "body": None,
-            "cookies": {},
-            "source": "flask",
-            "route": "/posts/:id",
-            "user": user,
-            "rate_limit_group": None,
-            "executed_middleware": executed_middleware,
-        }
-    ).set_as_current_context()
-
-
 def test_with_context_without_cache():
-    set_context()
+    test_utils.generate_and_set_context()
     get_cache().cache = None
     assert should_block_request() == {"block": False}
 
 
 def test_with_context_with_cache():
-    set_context(user={"id": "123"})
+    test_utils.generate_and_set_context(user={"id": "123"})
     thread_cache = get_cache()
 
     thread_cache.config.blocked_uids = ["123"]
@@ -76,7 +56,7 @@ def test_with_context_with_cache():
 
 
 def test_cache_comms_with_endpoints():
-    set_context(user={"id": "456"})
+    test_utils.generate_and_set_context(user={"id": "456"}, route="/posts/:id")
     set_rate_limit_group("my_group")
     thread_cache = get_cache()
     thread_cache.config.blocked_uids = ["123"]
@@ -145,11 +125,11 @@ def test_cache_comms_with_endpoints():
                 "route_metadata": {
                     "method": "POST",
                     "route": "/posts/:id",
-                    "url": "http://localhost:4000",
+                    "url": "http://localhost:8080/",
                 },
                 "user": {"id": "456"},
                 "group": "my_group",
-                "remote_address": "::1",
+                "remote_address": "1.1.1.1",
             },
             receive=True,
             timeout_in_sec=0.01,
@@ -168,7 +148,7 @@ def test_cache_comms_with_endpoints():
         assert thread_cache.stats.rate_limited_hits == 0
         assert should_block_request() == {
             "block": True,
-            "ip": "::1",
+            "ip": "1.1.1.1",
             "type": "ratelimited",
             "trigger": "my_trigger",
         }

aikido_zen/middleware/set_rate_limit_group_test.py

@@ -1,6 +1,6 @@
 import pytest
-from aikido_zen.context import get_current_context, Context
 from aikido_zen.thread.thread_cache import get_cache
+import aikido_zen.test_utils as test_utils
 from .set_rate_limit_group import set_rate_limit_group
 
 
@@ -15,31 +15,8 @@ def run_around_tests():
     get_cache().reset()
 
 
-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_set_rate_limit_group_valid_group_id(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group("group1")
     assert context1.rate_limit_group == "group1"
@@ -49,15 +26,15 @@ def test_set_rate_limit_group_valid_group_id(caplog):
 
 
 def test_set_rate_limit_group_empty_group_id(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group("")
     assert context1.rate_limit_group is None
     assert "Group ID cannot be empty." in caplog.text
 
 
 def test_set_rate_limit_group_none_group_id(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group(None)
     assert context1.rate_limit_group is None
@@ -73,30 +50,30 @@ def test_set_rate_limit_group_no_context(caplog):
 
 
 def test_set_rate_limit_group_middleware_already_executed(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     context1.executed_middleware = True
     set_rate_limit_group("group1")
     assert "must be called before the Zen middleware is executed" in caplog.text
     assert context1.rate_limit_group is "group1"
 
 
 def test_set_rate_limit_group_non_string_group_id(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group(123)
     assert context1.rate_limit_group == "123"
 
 
 def test_set_rate_limit_group_non_string_group_id_non_number(caplog):
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group({"a": "b"})
     assert context1.rate_limit_group is None
     assert "Group ID must be a string or a number" in caplog.text
 
 
 def test_set_rate_limit_group_overwrite_existing_group():
-    context1 = set_context_and_lifecycle()
+    context1 = test_utils.generate_and_set_context()
     assert context1.rate_limit_group is None
     set_rate_limit_group("group1")
     assert context1.rate_limit_group == "group1"

aikido_zen/sinks/builtins_import.py

@@ -1,8 +1,10 @@
+from aikido_zen.sinks import on_import, patch_function, after
+import contextvars
 import importlib.metadata
 from importlib.metadata import PackageNotFoundError
-
 from aikido_zen.background_process.packages import PackagesStore
-from aikido_zen.sinks import on_import, patch_function, after
+
+running_import_scan = contextvars.ContextVar("running_import_scan", default=False)
 
 
 @after
@@ -12,26 +14,34 @@ def _import(func, instance, args, kwargs, return_value):
 
     if not hasattr(return_value, "__package__"):
         return
-    name = getattr(return_value, "__package__")
-
-    if not name:
-        # Make sure the name exists
-        return
-    name = name.split(".")[0]  # Remove submodules
-    if name == "importlib" or name == "importlib_metadata":
-        # Avoid circular dependencies
-        return
-
-    if PackagesStore.get_package(name):
-        return
 
-    version = None
     try:
-        version = importlib.metadata.version(name)
-    except PackageNotFoundError:
-        pass
-    if version:
-        PackagesStore.add_package(name, version)
+        if running_import_scan.get():
+            return
+        running_import_scan.set(True)
+
+        name = getattr(return_value, "__package__")
+
+        if not name:
+            # Make sure the name exists
+            return
+        name = name.split(".")[0]  # Remove submodules
+        if name == "importlib" or name == "importlib_metadata":
+            # Avoid circular dependencies, this is a double safety-check for if contextvar check fails.
+            return
+
+        if PackagesStore.get_package(name):
+            return
+
+        version = None
+        try:
+            version = importlib.metadata.version(name)
+        except PackageNotFoundError:
+            pass
+        if version:
+            PackagesStore.add_package(name, version)
+    finally:
+        running_import_scan.set(False)
 
 
 @on_import("builtins")

aikido_zen/sinks/pathlib.py

@@ -0,0 +1,29 @@
+"""
+Sink module for python's `pathlib`
+"""
+
+import aikido_zen.vulnerabilities as vulns
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.helpers.register_call import register_call
+from aikido_zen.sinks import before, patch_function, on_import
+
+
+@before
+def _pathlib_truediv_patch(func, instance, args, kwargs):
+    path = get_argument(args, kwargs, 0, "key")
+    op = "pathlib.PurePath.__truediv__"
+    register_call(op, "fs_op")
+
+    vulns.run_vulnerability_scan(kind="path_traversal", op=op, args=(path,))
+
+
+@on_import("pathlib")
+def patch(m):
+    """
+    patching module pathlib
+    - patches PurePath.__truediv__ : Path() / Path() -> join operation
+    """
+
+    # PurePath() / "my/path/test.txt"
+    # This is accomplished by overloading the __truediv__ function on the Path class
+    patch_function(m, "PurePath.__truediv__", _pathlib_truediv_patch)

aikido_zen/sinks/tests/builtins_import_test.py

@@ -0,0 +1,56 @@
+import pytest
+
+import aikido_zen
+
+aikido_zen.protect()
+
+from aikido_zen.background_process.packages import PackagesStore
+
+
+@pytest.fixture(autouse=True)
+def run_around_tests():
+    PackagesStore.clear()
+
+
+def test_flask_import():
+    import flask
+
+    assert PackagesStore.get_package("flask")["version"] == "3.0.3"
+
+
+def test_django_import():
+    import django
+
+    assert PackagesStore.get_package("django")["version"] == "4.0"
+
+
+def test_recursive_package_store(monkeypatch):
+    """Test that recursive imports during package scanning don't cause max recursion depth errors."""
+
+    def recursive_get_package(name):
+        """Recursively add package and its dependencies to PackagesStore."""
+        import flask
+
+    PackagesStore.clear()
+    monkeypatch.setattr(PackagesStore, "get_package", recursive_get_package)
+
+    import flask
+
+    # Restore the original method after the test
+    monkeypatch.undo()
+
+
+def test_recursive_package_store_2(monkeypatch):
+    """Test that recursive imports during package scanning don't cause max recursion depth errors."""
+
+    def recursive_add_package(name, version):
+        """Recursively add package and its dependencies to PackagesStore."""
+        if name == "django":
+            import django
+
+    PackagesStore.clear()
+    monkeypatch.setattr(PackagesStore, "add_package", recursive_add_package)
+    import django
+
+    # Restore the original method after the test
+    monkeypatch.undo()