Merge branch 'main' into publish-to-pypi

Author: Wout Feys

.github/workflows/end2end.yml

@@ -0,0 +1,49 @@
+name: Run End-2-End Tests
+
+on: [pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v2
+
+    - name: Start django-mysql
+      working-directory: ./sample-apps/django-mysql
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Start django-mysql-gunicorn
+      working-directory: ./sample-apps/django-mysql-gunicorn
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Start flask-mongo
+      working-directory: ./sample-apps/flask-mongo
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Start flask-mysql
+      working-directory: ./sample-apps/flask-mysql
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Start flask-mysql-uwsgi
+      working-directory: ./sample-apps/flask-mysql-uwsgi
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Start flask-postgres
+      working-directory: ./sample-apps/flask-postgres
+      run: |
+        docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        make install
+
+    - name: Run end2end tests
+      run: |
+        make end2end

Makefile

@@ -18,11 +18,15 @@ install:
 
 .PHONY: test
 test:
-	poetry run pytest
+	poetry run pytest aikido_firewall/
+
+.PHONY: end2end
+end2end:
+	poetry run pytest end2end/ 	
 
 .PHONY: cov
 cov:
-	poetry run pytest --cov=aikido_firewall --cov-report=xml
+	poetry run pytest aikido_firewall/ --cov=aikido_firewall --cov-report=xml
 
 .PHONY: benchmark
 e2e:

aikido_firewall/background_process/aikido_background_process.py

@@ -85,6 +85,10 @@ def send_to_reporter(self, event_scheduler):
             EMPTY_QUEUE_INTERVAL, 1, self.send_to_reporter, (event_scheduler,)
         )
         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/heartbeats.py

@@ -3,6 +3,7 @@
 """
 
 from aikido_firewall.helpers.logging import logger
+from aikido_firewall.helpers.create_interval import create_interval
 
 
 def send_heartbeats_every_x_secs(reporter, interval_in_secs, event_scheduler):
@@ -18,22 +19,10 @@ def send_heartbeats_every_x_secs(reporter, interval_in_secs, event_scheduler):
 
     logger.debug("Starting heartbeats")
 
-    # Start the interval by booting the first settimeout
-    send_heartbeat_wrapper(reporter, interval_in_secs, event_scheduler, True)
-
-
-def send_heartbeat_wrapper(rep, interval_in_secs, event_scheduler, boot=False):
-    """
-    Wrapper function for send_heartbeat so we get an interval
-    """
-    event_scheduler.enter(
-        interval_in_secs,
-        1,
-        send_heartbeat_wrapper,
-        (rep, interval_in_secs, event_scheduler),
+    # Create an interval for "interval_in_secs" seconds :
+    create_interval(
+        event_scheduler=event_scheduler,
+        interval_in_secs=interval_in_secs,
+        function=lambda reporter: reporter.send_heartbeat(),
+        args=(reporter,),
     )
-    if not boot:
-        #  If boot is true it means it's the first time this gets executed, so we
-        #  need to wait for the interval in the event scheduler to finish
-        logger.debug("Heartbeat...")
-        rep.send_heartbeat()

aikido_firewall/background_process/heartbeats_test.py

@@ -1,9 +1,6 @@
 import pytest
 from unittest.mock import Mock, patch
-from aikido_firewall.background_process.heartbeats import (
-    send_heartbeats_every_x_secs,
-    send_heartbeat_wrapper,
-)
+from aikido_firewall.background_process.heartbeats import send_heartbeats_every_x_secs
 
 
 def test_send_heartbeats_serverless():
@@ -42,16 +39,3 @@ def test_send_heartbeats_success():
 
     with patch("aikido_firewall.helpers.logging.logger.debug") as mock_debug:
         send_heartbeats_every_x_secs(reporter, 5, event_scheduler)
-
-
-def test_send_heartbeat_wrapper():
-    reporter = Mock()
-    reporter.send_heartbeat = Mock()
-    event_scheduler = Mock()
-
-    send_heartbeat_wrapper(reporter, 5, event_scheduler)
-
-    reporter.send_heartbeat.assert_called_once()
-    event_scheduler.enter.assert_called_once_with(
-        5, 1, send_heartbeat_wrapper, (reporter, 5, event_scheduler)
-    )

aikido_firewall/background_process/reporter/__init__.py

@@ -70,9 +70,9 @@ def report_initial_stats(self):
         if should_report_initial_stats:
             self.send_heartbeat()
 
-    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 an 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/create_interval.py

@@ -0,0 +1,31 @@
+"""Exports create_interval"""
+
+
+def create_interval(event_scheduler, interval_in_secs, function, args):
+    """
+    This function creates an interval which first runs "function"
+    after waiting "interval_in_secs" seconds and after that keeps
+    executing the function every "interval_in_secs" seconds.
+    """
+    # Sleep interval_in_secs seconds before starting the loop :
+    event_scheduler.enter(
+        interval_in_secs,
+        1,
+        interval_loop,
+        (event_scheduler, interval_in_secs, function, args),
+    )
+
+
+def interval_loop(event_scheduler, interval_in_secs, function, args):
+    """
+    This is the actual interval loop which executes and schedules the function
+    """
+    # Execute function :
+    function(*args)
+    # Schedule the execution of the function in interval_in_secs seconds :
+    event_scheduler.enter(
+        interval_in_secs,
+        1,
+        interval_loop,
+        (event_scheduler, interval_in_secs, function, args),
+    )