Merge pull request #101 from AikidoSec/AIK-3353

Send a cleaned up version of the stacktrace to the backend

Author: willem-delbare

aikido_firewall/background_process/aikido_background_process.py

@@ -86,6 +86,10 @@ def send_to_reporter(self, event_scheduler):
         )
         logger.debug("Checking queue")
         while not self.queue.empty():
-            attack = self.queue.get()
-            logger.debug("Reporting attack : %s", attack)
-            self.reporter.on_detected_attack(attack[0], attack[1])
+            queue_attack_item = self.queue.get()
+            self.reporter.on_detected_attack(
+                attack=queue_attack_item[0],
+                context=queue_attack_item[1],
+                blocked=queue_attack_item[2],
+                stack=queue_attack_item[3],
+            )

aikido_firewall/background_process/commands/attack.py

@@ -4,7 +4,7 @@
 def process_attack(bg_process, data, conn):
     """
     Adds ATTACK data object to queue
-    Expected data object : [injection_results, context, blocked_or_not]
+    Expected data object : [injection_results, context, blocked_or_not, stacktrace]
     """
     bg_process.queue.put(data)
     if bg_process.reporter.statistics:

aikido_firewall/background_process/reporter/__init__.py

@@ -58,9 +58,9 @@ def start(self, event_scheduler):
         send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler)
         start_polling_for_changes(self, event_scheduler)
 
-    def on_detected_attack(self, attack, context):
+    def on_detected_attack(self, attack, context, blocked, stack):
         """This will send something to the API when an attack is detected"""
-        return on_detected_attack(self, attack, context)
+        return on_detected_attack(self, attack, context, blocked, stack)
 
     def on_start(self):
         """This will send out an Event signalling the start to the server"""

aikido_firewall/background_process/reporter/on_detected_attack.py

@@ -7,7 +7,7 @@
 from aikido_firewall.helpers.get_ua_from_context import get_ua_from_context
 
 
-def on_detected_attack(reporter, attack, context):
+def on_detected_attack(reporter, attack, context, blocked, stack):
     """
     This will send something to the API when an attack is detected
     """
@@ -18,7 +18,8 @@ def on_detected_attack(reporter, attack, context):
         attack["user"] = None
         attack["payload"] = json.dumps(attack["payload"])[:4096]
         attack["metadata"] = limit_length_metadata(attack["metadata"], 4096)
-        attack["blocked"] = reporter.block
+        attack["blocked"] = blocked
+        attack["stack"] = stack
 
         payload = {
             "type": "detected_attack",

aikido_firewall/background_process/reporter/on_detected_attack_test.py

@@ -31,7 +31,7 @@ class Context:
 def test_on_detected_attack_no_token(mock_context):
     reporter = MagicMock()
     reporter.token = None
-    on_detected_attack(reporter, {}, mock_context)
+    on_detected_attack(reporter, {}, mock_context, blocked=False, stack=None)
     reporter.api.report.assert_not_called()
 
 
@@ -42,7 +42,7 @@ def test_on_detected_attack_with_long_payload(mock_reporter, mock_context):
         "metadata": {"test": "1"},
     }
 
-    on_detected_attack(mock_reporter, attack, mock_context)
+    on_detected_attack(mock_reporter, attack, mock_context, blocked=False, stack=None)
     assert len(attack["payload"]) == 4096  # Ensure payload is truncated
     mock_reporter.api.report.assert_called_once()
 
@@ -54,7 +54,7 @@ def test_on_detected_attack_with_long_metadata(mock_reporter, mock_context):
         "metadata": {"test": long_metadata},
     }
 
-    on_detected_attack(mock_reporter, attack, mock_context)
+    on_detected_attack(mock_reporter, attack, mock_context, blocked=False, stack=None)
 
     assert (
         attack["metadata"]["test"] == long_metadata[:4096]
@@ -68,7 +68,7 @@ def test_on_detected_attack_success(mock_reporter, mock_context):
         "metadata": {},
     }
 
-    on_detected_attack(mock_reporter, attack, mock_context)
+    on_detected_attack(mock_reporter, attack, mock_context, blocked=False, stack=None)
     assert mock_reporter.api.report.call_count == 1
 
 
@@ -81,6 +81,24 @@ def test_on_detected_attack_exception_handling(mock_reporter, mock_context, capl
     # Simulate an exception during the API call
     mock_reporter.api.report.side_effect = Exception("API error")
 
-    on_detected_attack(mock_reporter, attack, mock_context)
+    on_detected_attack(mock_reporter, attack, mock_context, blocked=False, stack=None)
 
     assert "Failed to report attack" in caplog.text
+
+
+def test_on_detected_attack_with_blocked_and_stack(mock_reporter, mock_context):
+    attack = {
+        "payload": {"key": "value"},
+        "metadata": {},
+    }
+    blocked = True
+    stack = "sample stack trace"
+
+    on_detected_attack(
+        mock_reporter, attack, mock_context, blocked=blocked, stack=stack
+    )
+
+    # Check that the attack dictionary has the blocked and stack fields set
+    assert attack["blocked"] is True
+    assert attack["stack"] == stack
+    assert mock_reporter.api.report.call_count == 1

aikido_firewall/helpers/get_clean_stacktrace.py

@@ -0,0 +1,26 @@
+"""Exports function `get_clean_stacktrace`"""
+
+import inspect
+import sys
+
+
+def get_clean_stacktrace():
+    """Returns a cleaned up stacktrace"""
+    # Get the current stack
+    stack = inspect.stack()
+
+    # List of built-in modules to filter out
+    ignored_modules = sys.builtin_module_names
+
+    cleaned_stack = []
+
+    for frame_info in stack:
+        name = frame_info.frame.f_globals.get("__name__", "")
+
+        if name not in ignored_modules and not name.startswith("aikido_firewall"):
+            cleaned_stack.append(
+                f"File: {frame_info.filename}, L{frame_info.lineno} {frame_info.function}(...)"
+            )
+
+    cleaned_stack.reverse()
+    return "• " + "\n\n• ".join(cleaned_stack)

aikido_firewall/helpers/get_clean_stacktrace_test.py

@@ -0,0 +1,23 @@
+import pytest
+import pytest
+from .get_clean_stacktrace import get_clean_stacktrace
+
+
+def test_get_clean_stacktrace_no_aikido():
+    """Test that the stack trace does not include aikido frames."""
+
+    def dummy_function():
+        return get_clean_stacktrace()
+
+    result = dummy_function()
+
+    assert "/site-packages/aikido_firewall/" not in result
+
+
+def test_get_clean_stacktrace_with_aikido():
+    """Test that the stack trace includes non-aikido frames."""
+
+    def dummy_function():
+        return get_clean_stacktrace()
+
+    result = dummy_function()

aikido_firewall/vulnerabilities/__init__.py

@@ -14,6 +14,7 @@
 )
 from aikido_firewall.background_process import get_comms
 from aikido_firewall.helpers.logging import logger
+from aikido_firewall.helpers.get_clean_stacktrace import get_clean_stacktrace
 from aikido_firewall.helpers.blocking_enabled import is_blocking_enabled
 from .sql_injection.context_contains_sql_injection import context_contains_sql_injection
 from .nosql_injection.check_context import check_context_for_nosql_injection
@@ -87,6 +88,9 @@ def run_vulnerability_scan(kind, op, args):
     if injection_results:
         logger.debug("Injection results : %s", json.dumps(injection_results))
         blocked = is_blocking_enabled()
-        comms.send_data_to_bg_process("ATTACK", (injection_results, context, blocked))
+        stack = get_clean_stacktrace()
+        comms.send_data_to_bg_process(
+            "ATTACK", (injection_results, context, blocked, stack)
+        )
         if blocked:
             raise error_type(*error_args)

aikido_firewall/vulnerabilities/ssrf/inspect_getaddrinfo_result.py

@@ -9,6 +9,7 @@
 from aikido_firewall.background_process import get_comms
 from aikido_firewall.errors import AikidoSSRF
 from aikido_firewall.helpers.blocking_enabled import is_blocking_enabled
+from aikido_firewall.helpers.get_clean_stacktrace import get_clean_stacktrace
 from .imds import is_trusted_hostname, is_imds_ip_address
 from .is_private_ip import is_private_ip
 from .find_hostname_in_context import find_hostname_in_context
@@ -59,7 +60,10 @@ def inspect_getaddrinfo_result(dns_results, hostname, port):
     logger.debug("Attack results : %s", attack)
 
     logger.debug("Sending data to bg process :")
-    get_comms().send_data_to_bg_process("ATTACK", (attack, context, should_block))
+    stack = get_clean_stacktrace()
+    get_comms().send_data_to_bg_process(
+        "ATTACK", (attack, context, should_block, stack)
+    )
 
     if should_block:
         raise AikidoSSRF()