Merge pull request #274 from AikidoSec/set-middleware-installed

Set middlewareInstalled and send on heartbeat

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -52,6 +52,7 @@ def __init__(self, block, api, token, serverless):
         self.statistics = Statistics(
             max_perf_samples_in_mem=5000, max_compressed_stats_in_mem=100
         )
+        self.middleware_installed = False
 
         if isinstance(serverless, str) and len(serverless) == 0:
             raise ValueError("Serverless cannot be an empty string")

aikido_zen/background_process/cloud_connection_manager/init_test.py

@@ -0,0 +1,48 @@
+import pytest
+from aikido_zen.helpers.token import Token
+from aikido_zen.background_process.api.http_api import ReportingApiHTTP
+from aikido_zen.background_process.service_config import ServiceConfig
+from aikido_zen.background_process.users import Users
+from aikido_zen.background_process.hostnames import Hostnames
+from aikido_zen.ratelimiting.rate_limiter import RateLimiter
+from aikido_zen.background_process.statistics import Statistics
+from . import CloudConnectionManager
+
+
+@pytest.fixture
+def setup_cloud_connection_manager():
+    block = None  # Replace with appropriate mock or object if needed
+    api = ReportingApiHTTP(None)  # Mock or create an instance of ReportingApiHTTP
+    token = Token("AIK_TOKEN_TEST")  # Mock or create an instance of Token
+    serverless = "some_value"  # Valid serverless value
+    return CloudConnectionManager(block, api, token, serverless)
+
+
+def test_cloud_connection_manager_initialization(setup_cloud_connection_manager):
+    manager = setup_cloud_connection_manager
+
+    # Check that instance variables are initialized correctly
+    assert manager.block is None
+    assert isinstance(manager.api, ReportingApiHTTP)
+    assert isinstance(manager.token, Token)
+    assert len(manager.routes) == 0
+    assert isinstance(manager.hostnames, Hostnames)
+    assert isinstance(manager.conf, ServiceConfig)
+    assert isinstance(manager.rate_limiter, RateLimiter)
+    assert isinstance(manager.users, Users)
+    assert isinstance(manager.statistics, Statistics)
+    assert manager.middleware_installed is False
+    assert manager.serverless == "some_value"
+
+
+def test_cloud_connection_manager_empty_serverless():
+    block = None  # Replace with appropriate mock or object if needed
+    api = ReportingApiHTTP(None)  # Mock or create an instance of ReportingApiHTTP
+    token = Token("Hellow")  # Mock or create an instance of Token
+    serverless = ""  # Invalid serverless value
+
+    with pytest.raises(ValueError, match="Serverless cannot be an empty string"):
+        CloudConnectionManager(block, api, token, serverless)
+
+
+# Additional tests can be added here for other edge cases or scenarios

aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py

@@ -30,6 +30,7 @@ def send_heartbeat(connection_manager):
             "hostnames": outgoing_domains,
             "routes": routes,
             "users": users,
+            "middlewareInstalled": connection_manager.middleware_installed,
         },
         connection_manager.timeout_in_sec,
     )

aikido_zen/background_process/commands/sync_data.py

@@ -25,7 +25,12 @@ def process_sync_data(connection_manager, data, conn, queue=None):
         # Update API Spec :
         update_route_info(route["apispec"], existing_route)
 
+    # Save request data :
     connection_manager.statistics.requests["total"] += data.get("reqs", 0)
+
+    # Save middleware installed :
+    if data.get("middleware_installed", False):
+        connection_manager.middleware_installed = True
     if connection_manager.conf.last_updated_at > 0:
         # Only report data if the config has been fetched.
         return {

aikido_zen/background_process/commands/sync_data_test.py

@@ -16,6 +16,7 @@ def setup_connection_manager():
     connection_manager.conf.blocked_uids = ["user1", "user2"]
     connection_manager.conf.last_updated_at = 200
     connection_manager.statistics.requests = {"total": 0}  # Initialize total requests
+    connection_manager.middleware_installed = False
     return connection_manager
 
 
@@ -38,6 +39,7 @@ def test_process_sync_data_initialization(setup_connection_manager):
             },
         },
         "reqs": 10,  # Total requests to be added
+        "middleware_installed": False,
     }
 
     result = process_sync_data(connection_manager, data, None)
@@ -63,6 +65,7 @@ def test_process_sync_data_initialization(setup_connection_manager):
     # Check that the return value is correct
     assert result["routes"] == dict(connection_manager.routes.routes)
     assert result["config"] == connection_manager.conf
+    assert connection_manager.middleware_installed == False
 
 
 def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_manager):
@@ -85,6 +88,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
             },
         },
         "reqs": 10,  # Total requests to be added
+        "middleware_installed": True,
     }
 
     result = process_sync_data(connection_manager, data, None)
@@ -106,7 +110,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
 
     # Check that the total requests were updated
     assert connection_manager.statistics.requests["total"] == 10
-
+    assert connection_manager.middleware_installed == True
     # Check that the return value is correct
     assert result == {}
 
@@ -154,6 +158,7 @@ def test_process_sync_data_existing_route(setup_connection_manager):
 
     # Check that the total requests were updated
     assert connection_manager.statistics.requests["total"] == 20  # 5 + 15
+    assert connection_manager.middleware_installed == False
 
     # Check that the return value is correct
     assert result["routes"] == dict(connection_manager.routes.routes)

aikido_zen/middleware/init_test.py

@@ -56,19 +56,22 @@ def test_with_context_with_cache():
 
     thread_cache.config.blocked_uids = ["123"]
     assert get_current_context().executed_middleware == False
+    assert thread_cache.middleware_installed == False
     assert should_block_request() == {
         "block": True,
         "trigger": "user",
         "type": "blocked",
     }
     assert get_current_context().executed_middleware == True
+    assert thread_cache.middleware_installed == True
 
     thread_cache.config.blocked_uids = []
     assert should_block_request() == {"block": False}
 
     thread_cache.config.blocked_uids = ["23", "234", "456"]
     assert should_block_request() == {"block": False}
     assert get_current_context().executed_middleware == True
+    assert thread_cache.middleware_installed == True
 
 
 def test_cache_comms_with_endpoints():
@@ -88,11 +91,13 @@ def test_cache_comms_with_endpoints():
         }
     ]
     assert get_current_context().executed_middleware == False
+    assert thread_cache.middleware_installed == False
 
     with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms:
         mock_get_comms.return_value = None  # Set the return value of get_comms
         assert should_block_request() == {"block": False}
     assert get_current_context().executed_middleware == True
+    assert thread_cache.middleware_installed == True
 
     with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms:
         mock_comms = MagicMock()

aikido_zen/middleware/should_block_request.py

@@ -17,11 +17,15 @@ def should_block_request():
         cache = get_cache()
         if not context or not cache:
             return {"block": False}
+
         context.executed_middleware = (
             True  # Update context with middleware execution set to true
         )
         context.set_as_current_context()
 
+        # Make sure we set middleware installed to true (reports back to core) :
+        cache.middleware_installed = True
+
         # Blocked users:
         if context.user and cache.is_user_blocked(context.user["id"]):
             return {"block": True, "type": "blocked", "trigger": "user"}

aikido_zen/thread/thread_cache.py

@@ -68,6 +68,7 @@ def reset(self):
         )
         self.reqs = 0
         self.last_renewal = 0
+        self.middleware_installed = False
 
     def renew(self):
         """
@@ -77,7 +78,11 @@ def renew(self):
             return
         res = comms.get_comms().send_data_to_bg_process(
             action="SYNC_DATA",
-            obj={"current_routes": dict(self.routes.routes), "reqs": self.reqs},
+            obj={
+                "current_routes": dict(self.routes.routes),
+                "reqs": self.reqs,
+                "middleware_installed": self.middleware_installed,
+            },
             receive=True,
         )
 

aikido_zen/vulnerabilities/init_test.py

@@ -89,7 +89,6 @@ def test_lifecycle_cache_bypassed_ip(caplog, get_context):
 
 
 def test_sql_injection(caplog, get_context, monkeypatch):
-
     get_context.set_as_current_context()
     cache = ThreadCache()
     monkeypatch.setenv("AIKIDO_BLOCKING", "1")
@@ -114,7 +113,6 @@ def test_sql_injection_with_route_params(caplog, get_context, monkeypatch):
 
 
 def test_sql_injection_with_comms(caplog, get_context, monkeypatch):
-
     get_context.set_as_current_context()
     cache = ThreadCache()
     cache.last_renewal = 9999999999999999999999
@@ -140,7 +138,6 @@ def test_sql_injection_with_comms(caplog, get_context, monkeypatch):
 
 
 def test_ssrf_with_comms_hostnames_add(caplog, get_context, monkeypatch):
-
     get_context.set_as_current_context()
     monkeypatch.setenv("AIKIDO_BLOCKING", "1")
     with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: