Create an enable_idor_protection.py and an idor_protection_config.py

Author: bitterpanda63

aikido_zen/storage/idor_protection_config.py

@@ -0,0 +1,21 @@
+class IdorProtectionConfig:
+    def __init__(self, tenant_column_name, excluded_tables):
+        self.tenant_column_name = tenant_column_name
+        self.excluded_tables = excluded_tables
+
+
+class IdorProtectionStore:
+    def __init__(self):
+        self.config = None
+
+    def get(self):
+        return self.config
+
+    def set(self, config):
+        self.config = config
+
+    def clear(self):
+        self.config = None
+
+
+idor_protection_store = IdorProtectionStore()

aikido_zen/vulnerabilities/idor/enable_idor_protection.py

@@ -0,0 +1,40 @@
+from aikido_zen.helpers.logging import logger
+from aikido_zen.storage.idor_protection_config import (
+    IdorProtectionConfig,
+    idor_protection_store,
+)
+
+
+def enable_idor_protection(tenant_column_name: str, excluded_tables=None):
+    if not isinstance(tenant_column_name, str):
+        logger.info(
+            "enable_idor_protection(...) expects tenant_column_name to be a string, found %s instead.",
+            type(tenant_column_name),
+        )
+        return
+
+    if len(tenant_column_name) == 0:
+        logger.info(
+            "enable_idor_protection(...) expects tenant_column_name to be a non-empty string."
+        )
+        return
+
+    if excluded_tables is None:
+        excluded_tables = []
+
+    if not isinstance(excluded_tables, list):
+        logger.info(
+            "enable_idor_protection(...) expects excluded_tables to be a list, found %s instead.",
+            type(excluded_tables),
+        )
+        return
+
+    for table in excluded_tables:
+        if not isinstance(table, str):
+            logger.info(
+                "enable_idor_protection(...) expects excluded_tables to contain strings, found %s instead.",
+                type(table),
+            )
+            return
+
+    idor_protection_store.set(IdorProtectionConfig(tenant_column_name, excluded_tables))

aikido_zen/vulnerabilities/idor/enable_idor_protection_test.py

@@ -0,0 +1,55 @@
+import pytest
+from aikido_zen.storage.idor_protection_config import idor_protection_store
+from .enable_idor_protection import enable_idor_protection
+
+
+@pytest.fixture(autouse=True)
+def run_around_tests():
+    yield
+    idor_protection_store.clear()
+
+
+def test_enable_basic():
+    enable_idor_protection("tenant_id")
+    config = idor_protection_store.get()
+    assert config is not None
+    assert config.tenant_column_name == "tenant_id"
+    assert config.excluded_tables == []
+
+
+def test_enable_with_excluded_tables():
+    enable_idor_protection("org_id", excluded_tables=["migrations", "sessions"])
+    config = idor_protection_store.get()
+    assert config is not None
+    assert config.tenant_column_name == "org_id"
+    assert config.excluded_tables == ["migrations", "sessions"]
+
+
+def test_invalid_column_name_type(caplog):
+    enable_idor_protection(123)
+    assert idor_protection_store.get() is None
+    assert "expects tenant_column_name to be a string" in caplog.text
+
+
+def test_empty_column_name(caplog):
+    enable_idor_protection("")
+    assert idor_protection_store.get() is None
+    assert "non-empty string" in caplog.text
+
+
+def test_invalid_excluded_tables_type(caplog):
+    enable_idor_protection("tenant_id", excluded_tables="not_a_list")
+    assert idor_protection_store.get() is None
+    assert "expects excluded_tables to be a list" in caplog.text
+
+
+def test_invalid_excluded_table_item(caplog):
+    enable_idor_protection("tenant_id", excluded_tables=[123])
+    assert idor_protection_store.get() is None
+    assert "expects excluded_tables to contain strings" in caplog.text
+
+
+def test_none_column_name(caplog):
+    enable_idor_protection(None)
+    assert idor_protection_store.get() is None
+    assert "expects tenant_column_name to be a string" in caplog.text