Fix user agent not being reported correctly on attack

aikido_zen/background_process/cloud_connection_manager/on_detected_attack.py

@@ -1,10 +1,10 @@
 """Mainly exports on_detected_attack"""
 
 import json
+
 from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms
 from aikido_zen.helpers.logging import logger
 from aikido_zen.helpers.limit_length_metadata import limit_length_metadata
-from aikido_zen.helpers.get_ua_from_context import get_ua_from_context
 from aikido_zen.helpers.serialize_to_json import serialize_to_json
 
 
@@ -31,7 +31,7 @@ def on_detected_attack(connection_manager, attack, context, blocked, stack):
                 "method": context.method,
                 "url": context.url,
                 "ipAddress": context.remote_address,
-                "userAgent": get_ua_from_context(context),
+                "userAgent": context.get_user_agent(),
                 "body": context.body,
                 "headers": context.headers,
                 "source": context.source,

aikido_zen/background_process/cloud_connection_manager/on_detected_attack_test.py

@@ -1,6 +1,7 @@
 import pytest
 from unittest.mock import MagicMock, patch
 from .on_detected_attack import on_detected_attack
+from ...context import Context
 
 
 @pytest.fixture
@@ -16,16 +17,22 @@ def mock_connection_manager():
 
 @pytest.fixture
 def mock_context():
-    class Context:
-        method = "POST"
-        url = "http://example.com/api"
-        remote_address = "192.168.1.1"
-        body = "test body"
-        headers = {"Content-Type": "application/json"}
-        source = "test_source"
-        route = "test_route"
+    basic_wsgi_req = {
+        "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",
+        "HTTP_USER_AGENT": "Mozilla/5.0",
+    }
 
-    return Context()
+    return Context(req=basic_wsgi_req, body=123, source="django")
 
 
 def test_on_detected_attack_no_token(mock_context):
@@ -114,3 +121,64 @@ def test_on_detected_attack_with_blocked_and_stack(
     assert attack["blocked"] is True
     assert attack["stack"] == stack
     assert mock_connection_manager.api.report.call_count == 1
+
+
+def test_on_detected_attack_request_data_and_attack_data(
+    mock_connection_manager, mock_context
+):
+    attack = {
+        "payload": {"key": "value"},
+        "metadata": {"test": "true"},
+    }
+
+    on_detected_attack(
+        mock_connection_manager, attack, mock_context, blocked=False, stack=None
+    )
+
+    # Extract the call arguments for the report method
+    _, event, _ = mock_connection_manager.api.report.call_args[0]
+
+    # Verify the request attribute in the payload
+    request_data = event["request"]
+
+    assert request_data["method"] == "GET"
+    assert request_data["url"] == "http://localhost:8080/hello"
+    assert request_data["ipAddress"] == "198.51.100.23"
+    assert request_data["body"] == 123
+    assert request_data["headers"] == {
+        "CONTENT_TYPE": "application/json",
+        "USER_AGENT": "Mozilla/5.0",
+        "COOKIE": "sessionId=abc123xyz456;",
+        "HEADER_1": "header 1 value",
+        "HEADER_2": "Header 2 value",
+        "HOST": "localhost:8080",
+    }
+    assert request_data["source"] == "django"
+    assert request_data["route"] == "/hello"
+    assert request_data["userAgent"] == "Mozilla/5.0"
+
+    attack_data = event["attack"]
+    assert attack_data["blocked"] == False
+    assert attack_data["metadata"] == {"test": "true"}
+    assert attack_data["payload"] == '{"key": "value"}'
+    assert attack_data["stack"] is None
+    assert attack_data["user"] is None
+
+
+def test_on_detected_attack_with_user(mock_connection_manager, mock_context):
+    attack = {
+        "payload": {"key": "value"},
+        "metadata": {},
+    }
+    # Simulate a user in the context
+    mock_context.user = "test_user"
+
+    on_detected_attack(
+        mock_connection_manager, attack, mock_context, blocked=False, stack=None
+    )
+
+    # Extract the call arguments for the report method
+    _, event, _ = mock_connection_manager.api.report.call_args[0]
+
+    # Verify the user is included in the attack data
+    assert event["attack"]["user"] == "test_user"

end2end/django_mysql_gunicorn_test.py

@@ -50,6 +50,10 @@ def test_dangerous_response_with_firewall():
         'user': None
     }
 
+    assert attacks[0]["request"]["source"] == "django"
+    assert attacks[0]["request"]["route"] == "/app/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
 def test_dangerous_response_without_firewall():
     dog_name = 'Dangerous bobby", 1); -- '
     res = requests.post(post_url_nofw, data={'dog_name': dog_name})

end2end/flask_mysql_uwsgi_test.py

@@ -50,6 +50,10 @@ def test_dangerous_response_with_firewall():
         'user': None
     }
 
+    assert attacks[0]["request"]["source"] == "flask"
+    assert attacks[0]["request"]["route"] == "/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
 def test_dangerous_response_without_firewall():
     dog_name = 'Dangerous bobby", 1); -- '
     res = requests.post(post_url_nofw, data={'dog_name': dog_name})

end2end/quart_postgres_uvicorn_test.py

@@ -46,6 +46,10 @@ def test_dangerous_response_with_firewall():
     assert attacks[0]["attack"]["user"]["id"] == "user123"
     assert attacks[0]["attack"]["user"]["name"] == "John Doe"
 
+    assert attacks[0]["request"]["source"] == "quart"
+    assert attacks[0]["request"]["route"] == "/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
 
 def test_dangerous_response_without_firewall():
     dog_name = "Dangerous Bobby', TRUE); -- "

end2end/starlette_postgres_uvicorn_test.py

@@ -48,6 +48,10 @@ def test_dangerous_response_with_firewall():
     assert attacks[0]["attack"]["user"]["id"] == "user123"
     assert attacks[0]["attack"]["user"]["name"] == "John Doe"
 
+    assert attacks[0]["request"]["source"] == "starlette"
+    assert attacks[0]["request"]["route"] == "/create"
+    assert attacks[0]["request"]["userAgent"] == "python-requests/2.32.3"
+
 
 def test_dangerous_response_without_firewall():
     dog_name = "Dangerous Bobby', TRUE); -- "