Attack wave detection: fix bugs & report samples

aikido_zen/helpers/create_attack_wave_event.py

@@ -1,9 +1,20 @@
+import json
 from aikido_zen.helpers.limit_length_metadata import limit_length_metadata
 from aikido_zen.helpers.logging import logger
+from aikido_zen.storage.attack_wave_detector_store import attack_wave_detector_store
 
 
-def create_attack_wave_event(context, metadata):
+def create_attack_wave_event(context):
     try:
+        metadata = {}
+
+        samples = attack_wave_detector_store.get_samples_for_ip(context.remote_address)
+        if samples:
+            # Convert samples to JSON string, since metadata is a key-value store of strings.
+            metadata["samples"] = json.dumps(samples)
+
+        attack_wave_detector_store.clear_samples_for_ip(context.remote_address)
+
         return {
             "type": "detected_attack_wave",
             "attack": {

aikido_zen/helpers/create_attack_wave_event_test.py

@@ -1,93 +1,181 @@
 import pytest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 from .create_attack_wave_event import (
     create_attack_wave_event,
     extract_request_if_possible,
 )
+from aikido_zen.storage.attack_wave_detector_store import attack_wave_detector_store
 import aikido_zen.test_utils as test_utils
 
 
 def test_create_attack_wave_event_success():
     """Test successful creation of attack wave event with basic data"""
-    metadata = {"test": "value"}
     context = test_utils.generate_context()
 
-    event = create_attack_wave_event(context, metadata)
+    # Mock the attack_wave_detector_store to return no samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=None
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
 
-    assert event is not None
-    assert event["type"] == "detected_attack_wave"
-    assert event["attack"]["user"] is None
-    assert event["attack"]["metadata"] == metadata
-    assert event["request"] is not None
+        event = create_attack_wave_event(context)
+
+        assert event is not None
+        assert event["type"] == "detected_attack_wave"
+        assert event["attack"]["user"] is None
+        assert event["attack"]["metadata"] == {}
+        assert event["request"] is not None
+
+
+def test_create_attack_wave_event_with_samples():
+    """Test attack wave event creation with samples from store"""
+    context = test_utils.generate_context()
+
+    # Create sample data (now only method and url)
+    samples = [
+        {
+            "method": "GET",
+            "url": "/test1",
+        },
+        {
+            "method": "POST",
+            "url": "/test2",
+        },
+    ]
+
+    # Mock the attack_wave_detector_store to return samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=samples
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
+
+        event = create_attack_wave_event(context)
+
+        assert event is not None
+        assert event["type"] == "detected_attack_wave"
+        assert event["attack"]["user"] is None
+        assert "samples" in event["attack"]["metadata"]
+
+        # Check that samples are now JSON stringified
+        import json
+
+        samples_json = event["attack"]["metadata"]["samples"]
+        samples = json.loads(samples_json)
+
+        assert len(samples) == 2
+
+        # Check that samples contain only method and url
+        sample1 = samples[0]
+        assert sample1["method"] == "GET"
+        assert sample1["url"] == "/test1"
+        assert "user_agent" not in sample1
+        assert "timestamp" not in sample1
 
 
 def test_create_attack_wave_event_with_user():
     """Test attack wave event creation with user information"""
-    metadata = {"test": "value"}
     context = test_utils.generate_context(user="test_user")
 
-    event = create_attack_wave_event(context, metadata)
+    # Mock the attack_wave_detector_store to return no samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=None
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
+
+        event = create_attack_wave_event(context)
 
-    assert event["attack"]["user"] == "test_user"
-    assert event["attack"]["metadata"] == metadata
+        assert event["attack"]["user"] == "test_user"
+        assert event["attack"]["metadata"] == {}
 
 
 def test_create_attack_wave_event_with_long_metadata():
-    """Test that metadata longer than 4096 characters is truncated"""
-    long_metadata = "x" * 5000  # Create metadata longer than 4096 characters
-    metadata = {"test": long_metadata}
+    """Test that metadata with long samples is truncated"""
     context = test_utils.generate_context()
 
-    event = create_attack_wave_event(context, metadata)
+    # Create samples with very long values
+    long_value = "x" * 5000
+    samples = [
+        {
+            "method": "GET",
+            "url": "/test" + long_value,  # Very long url
+        },
+    ]
 
-    assert len(event["attack"]["metadata"]["test"]) == 4096
-    assert event["attack"]["metadata"]["test"] == long_metadata[:4096]
+    # Mock the attack_wave_detector_store to return samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=samples
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
+
+        event = create_attack_wave_event(context)
+
+        # The metadata should be truncated to 4096 characters
+        metadata_json = event["attack"]["metadata"]
+        assert len(metadata_json) <= 4096
 
 
 def test_create_attack_wave_event_with_multiple_long_metadata_fields():
-    """Test that multiple metadata fields longer than 4096 characters are truncated"""
+    """Test that metadata with multiple long sample fields is truncated"""
+    context = test_utils.generate_context()
+
+    # Create samples with very long values in multiple fields
     long_value1 = "a" * 5000
     long_value2 = "b" * 6000
-    metadata = {
-        "field1": long_value1,
-        "field2": long_value2,
-    }
-    context = test_utils.generate_context()
+    samples = [
+        {
+            "method": "GET" + long_value1,
+            "url": "/test" + long_value2,
+        },
+    ]
+
+    # Mock the attack_wave_detector_store to return samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=samples
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
 
-    event = create_attack_wave_event(context, metadata)
+        event = create_attack_wave_event(context)
 
-    assert len(event["attack"]["metadata"]["field1"]) == 4096
-    assert len(event["attack"]["metadata"]["field2"]) == 4096
-    assert event["attack"]["metadata"]["field1"] == long_value1[:4096]
-    assert event["attack"]["metadata"]["field2"] == long_value2[:4096]
+        # The metadata should be truncated to 4096 characters
+        metadata_json = event["attack"]["metadata"]
+        assert len(metadata_json) <= 4096
 
 
 def test_create_attack_wave_event_request_data():
     """Test that request data is correctly extracted from context"""
-    metadata = {"test": "value"}
     context = test_utils.generate_context(
         ip="198.51.100.23",
         route="/test-route",
         headers={"user-agent": "Mozilla/5.0"},
     )
 
-    event = create_attack_wave_event(context, metadata)
+    # Mock the attack_wave_detector_store to return no samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=None
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
 
-    request_data = event["request"]
-    assert request_data["ipAddress"] == "198.51.100.23"
-    assert request_data["source"] == "flask"
-    assert request_data["userAgent"] == "Mozilla/5.0"
+        event = create_attack_wave_event(context)
+
+        request_data = event["request"]
+        assert request_data["ipAddress"] == "198.51.100.23"
+        assert request_data["source"] == "flask"
+        assert request_data["userAgent"] == "Mozilla/5.0"
 
 
 def test_create_attack_wave_event_no_context():
     """Test attack wave event creation with None context"""
-    metadata = {"test": "value"}
 
-    event = create_attack_wave_event(None, metadata)
+    event = create_attack_wave_event(None)
 
-    assert event["attack"]["user"] is None
-    assert event["attack"]["metadata"] == metadata
-    assert event["request"] is None
+    # Function returns None when context is None (due to exception handling)
+    assert event is None
 
 
 def test_create_attack_wave_event_exception_handling():
@@ -100,14 +188,18 @@ def test_create_attack_wave_event_exception_handling():
     # Make get_user_agent raise an exception
     context.get_user_agent.side_effect = Exception("Test exception")
 
-    metadata = {"test": "value"}
-
-    # This should not raise an exception, but return None
-    event = create_attack_wave_event(context, metadata)
+    # Mock the attack_wave_detector_store to raise an exception
+    with patch.object(
+        attack_wave_detector_store,
+        "get_samples_for_ip",
+        side_effect=Exception("Store exception"),
+    ):
+        # This should not raise an exception, but return None
+        event = create_attack_wave_event(context)
 
-    # Since we're mocking and causing an exception, the function should handle it
-    # and return None based on the exception handling in the function
-    assert event is None
+        # Since we're mocking and causing an exception, the function should handle it
+        # and return None based on the exception handling in the function
+        assert event is None
 
 
 def test_extract_request_if_possible_with_valid_context():
@@ -145,29 +237,56 @@ def test_extract_request_if_possible_with_minimal_context():
 
 
 def test_create_attack_wave_event_empty_metadata():
-    """Test attack wave event creation with empty metadata"""
-    metadata = {}
+    """Test attack wave event creation with no samples (empty metadata)"""
     context = test_utils.generate_context()
 
-    event = create_attack_wave_event(context, metadata)
+    # Mock the attack_wave_detector_store to return no samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=None
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
+
+        event = create_attack_wave_event(context)
 
-    assert event is not None
-    assert event["attack"]["metadata"] == {}
-    assert event["request"] is not None
+        assert event is not None
+        assert event["attack"]["metadata"] == {}
+        assert event["request"] is not None
 
 
 def test_create_attack_wave_event_complex_metadata():
-    """Test attack wave event creation with complex nested metadata"""
-    metadata = {
-        "nested": {"key1": "value1", "key2": "value2"},
-        "simple": "simple_value",
-        "json_string": "[1, 2, 3]",
-        "number_string": "42",
-    }
+    """Test attack wave event creation with complex nested samples"""
     context = test_utils.generate_context()
 
-    event = create_attack_wave_event(context, metadata)
-
-    assert event["attack"]["metadata"] == metadata
-    assert event["attack"]["metadata"]["nested"]["key1"] == "value1"
-    assert event["attack"]["metadata"]["json_string"] == "[1, 2, 3]"
+    # Create complex samples (now only method and url)
+    samples = [
+        {
+            "method": "GET",
+            "url": "/complex",
+        },
+        {
+            "method": "POST",
+            "url": "/nested",
+        },
+    ]
+
+    # Mock the attack_wave_detector_store to return samples
+    with patch.object(
+        attack_wave_detector_store, "get_samples_for_ip", return_value=samples
+    ), patch.object(
+        attack_wave_detector_store, "clear_samples_for_ip", return_value=None
+    ):
+
+        event = create_attack_wave_event(context)
+
+        # Check that samples are now JSON stringified
+        import json
+
+        samples_json = event["attack"]["metadata"]["samples"]
+        samples = json.loads(samples_json)
+
+        assert len(samples) == 2
+        assert samples[0]["method"] == "GET"
+        assert samples[1]["method"] == "POST"
+        assert samples[0]["url"] == "/complex"
+        assert samples[1]["url"] == "/nested"

aikido_zen/sources/functions/request_handler.py

@@ -98,11 +98,11 @@ def post_response(status_code):
     if not cache:
         return
 
-    attack_wave = attack_wave_detector_store.is_attack_wave(context.remote_address)
+    attack_wave = attack_wave_detector_store.is_attack_wave(context)
     if attack_wave:
         cache.stats.on_detected_attack_wave(blocked=False)
 
-        event = create_attack_wave_event(context, metadata={})
+        event = create_attack_wave_event(context)
         logger.debug("Attack wave: %s", serialize_to_json(event)[:5000])
 
         # Report in background to core (send event over IPC)