From e1e5ace3d7126b453dc5373f138548c08144339c Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 3 Sep 2024 15:02:33 +0200 Subject: [PATCH 1/4] Add a ReportingApiHTTPRatelimited class similar to node.js --- .../api/http_api_ratelimited.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 aikido_firewall/background_process/api/http_api_ratelimited.py diff --git a/aikido_firewall/background_process/api/http_api_ratelimited.py b/aikido_firewall/background_process/api/http_api_ratelimited.py new file mode 100644 index 00000000..fd49705a --- /dev/null +++ b/aikido_firewall/background_process/api/http_api_ratelimited.py @@ -0,0 +1,33 @@ +""" +Exports ReportingApiHTTPRatelimited +""" + +from aikido_firewall.background_process.api.http_api import ReportingApiHTTP +import aikido_firewall.helpers.get_current_unixtime_ms as t + + +class ReportingApiHTTPRatelimited(ReportingApiHTTP): + """HTTP Reporting API that has ratelimiting support""" + + def __init__(self, reporting_url, max_events_per_interval, interval_in_ms): + super().__init__(reporting_url) + self.interval_in_ms = interval_in_ms + self.max_events_per_interval = max_events_per_interval + self.events = [] + + def report(self, token, event, timeout_in_sec): + if event["type"] == "detected_attack": + # Remove all outdated events : + current_time = t.get_unixtime_ms() + + def event_in_interval_filter(e): + return e["time"] > current_time - self.interval_in_ms + + self.events = list(filter(event_in_interval_filter, self.events)) + + # Check if interval is exceeded : + if len(self.events) >= self.max_events_per_interval: + return {"success": False, "error": "max_attacks_reached"} + + self.events.append(event) + return super().report(token, event, timeout_in_sec) From c35e37a90dd086247896d764aff20c3d5afd6114 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 3 Sep 2024 15:44:05 +0200 Subject: [PATCH 2/4] Use clientside ratelimiting with same parameters as Node.js --- .../background_process/aikido_background_process.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/background_process/aikido_background_process.py b/aikido_firewall/background_process/aikido_background_process.py index 1d660ad4..284d75ee 100644 --- a/aikido_firewall/background_process/aikido_background_process.py +++ b/aikido_firewall/background_process/aikido_background_process.py @@ -15,7 +15,9 @@ ) from aikido_firewall.helpers.check_env_for_blocking import check_env_for_blocking from aikido_firewall.helpers.token import get_token_from_env -from aikido_firewall.background_process.api.http_api import ReportingApiHTTP +from aikido_firewall.background_process.api.http_api_ratelimited import ( + ReportingApiHTTPRatelimited, +) from .commands import process_incoming_command EMPTY_QUEUE_INTERVAL = 5 # 5 seconds @@ -70,7 +72,11 @@ def reporting_thread(self): ) # Create an event scheduler self.send_to_connection_manager(event_scheduler) - api = ReportingApiHTTP("https://guard.aikido.dev/") + api = ReportingApiHTTPRatelimited( + "https://guard.aikido.dev/", + max_events_per_interval=2, + interval_in_ms=15 * 1000, + ) # We need to pass along the scheduler so that the heartbeat also gets sent self.connection_manager = CloudConnectionManager( block=check_env_for_blocking(), From fd774ae80b33accb4d0695b4ed87d3d5c4355b5a Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Sep 2024 11:43:32 +0200 Subject: [PATCH 3/4] Use values from node.js for ratelimiting client-side --- .../background_process/aikido_background_process.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aikido_firewall/background_process/aikido_background_process.py b/aikido_firewall/background_process/aikido_background_process.py index 284d75ee..3d7d8afc 100644 --- a/aikido_firewall/background_process/aikido_background_process.py +++ b/aikido_firewall/background_process/aikido_background_process.py @@ -74,8 +74,8 @@ def reporting_thread(self): api = ReportingApiHTTPRatelimited( "https://guard.aikido.dev/", - max_events_per_interval=2, - interval_in_ms=15 * 1000, + max_events_per_interval=100, + interval_in_ms=60 * 60 * 1000, ) # We need to pass along the scheduler so that the heartbeat also gets sent self.connection_manager = CloudConnectionManager( From 27bced084e364be0fa27712f40d239914373ec4c Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Wed, 4 Sep 2024 12:26:03 +0200 Subject: [PATCH 4/4] Add unit tests to client-side ratelimited HTTP Api Client --- .../api/http_api_ratelimited.py | 4 +- .../api/http_api_ratelimited_test.py | 207 ++++++++++++++++++ 2 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 aikido_firewall/background_process/api/http_api_ratelimited_test.py diff --git a/aikido_firewall/background_process/api/http_api_ratelimited.py b/aikido_firewall/background_process/api/http_api_ratelimited.py index fd49705a..3a019f69 100644 --- a/aikido_firewall/background_process/api/http_api_ratelimited.py +++ b/aikido_firewall/background_process/api/http_api_ratelimited.py @@ -2,11 +2,11 @@ Exports ReportingApiHTTPRatelimited """ -from aikido_firewall.background_process.api.http_api import ReportingApiHTTP +import aikido_firewall.background_process.api.http_api as http_api import aikido_firewall.helpers.get_current_unixtime_ms as t -class ReportingApiHTTPRatelimited(ReportingApiHTTP): +class ReportingApiHTTPRatelimited(http_api.ReportingApiHTTP): """HTTP Reporting API that has ratelimiting support""" def __init__(self, reporting_url, max_events_per_interval, interval_in_ms): diff --git a/aikido_firewall/background_process/api/http_api_ratelimited_test.py b/aikido_firewall/background_process/api/http_api_ratelimited_test.py new file mode 100644 index 00000000..28925ea1 --- /dev/null +++ b/aikido_firewall/background_process/api/http_api_ratelimited_test.py @@ -0,0 +1,207 @@ +import pytest +from unittest.mock import patch +from aikido_firewall.background_process.api.http_api import ReportingApiHTTP +from aikido_firewall.helpers.get_current_unixtime_ms import get_unixtime_ms +from .http_api_ratelimited import ReportingApiHTTPRatelimited + + +@pytest.fixture +def reporting_api(): + """Fixture to create an instance of ReportingApiHTTPRatelimited.""" + return ReportingApiHTTPRatelimited( + reporting_url="http://example.com", + max_events_per_interval=3, + interval_in_ms=10000, + ) + + +def test_report_within_limit(reporting_api): + """Test reporting within the rate limit.""" + event = {"type": "detected_attack", "time": 1000} + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=2000, + ): + response = reporting_api.report( + token="token", event=event, timeout_in_sec=5 + ) + assert response == {"success": True} + assert len(reporting_api.events) == 1 + mock_report.assert_called_once_with("token", event, 5) + + +def test_report_exceeds_limit(reporting_api): + """Test reporting when the limit is exceeded.""" + event = {"type": "detected_attack", "time": 1000} + + # Simulate adding events to reach the limit + reporting_api.events = [ + {"type": "detected_attack", "time": 1000}, + {"type": "detected_attack", "time": 2000}, + {"type": "detected_attack", "time": 3000}, + ] + + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=4000, + ): + response = reporting_api.report(token="token", event=event, timeout_in_sec=5) + assert response == {"success": False, "error": "max_attacks_reached"} + assert len(reporting_api.events) == 3 # Should not add the new event + + +def test_report_within_limit_after_expiry(reporting_api): + """Test reporting after some events have expired.""" + event1 = {"type": "detected_attack", "time": 1000} + event2 = {"type": "detected_attack", "time": 2001} + + # Add events to the list + reporting_api.events = [event1, event2] + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=12000, + ): + event3 = {"type": "detected_attack", "time": 11000} + response = reporting_api.report( + token="token", event=event3, timeout_in_sec=5 + ) + assert response == {"success": True} + assert ( + len(reporting_api.events) == 2 + ) # One event should have expired, and the new one is added + + +def test_report_with_non_attack_event(reporting_api): + """Test reporting with a non-attack event.""" + event = {"type": "other_event", "time": 1000} + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + response = reporting_api.report(token="token", event=event, timeout_in_sec=5) + assert response == {"success": True} + assert len(reporting_api.events) == 0 # Non-attack events should not be stored + mock_report.assert_called_once_with("token", event, 5) + + +def test_report_multiple_events_within_limit(reporting_api): + """Test reporting multiple events within the rate limit.""" + events = [ + {"type": "detected_attack", "time": 1000}, + {"type": "detected_attack", "time": 2000}, + {"type": "detected_attack", "time": 3000}, + ] + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=4000, + ): + for event in events: + response = reporting_api.report( + token="token", event=event, timeout_in_sec=5 + ) + assert response == {"success": True} + assert len(reporting_api.events) == 3 + assert mock_report.call_count == 3 + + +def test_report_mixed_event_types(reporting_api): + """Test reporting with a mix of attack and non-attack events.""" + attack_event = {"type": "detected_attack", "time": 1000} + non_attack_event = {"type": "other_event", "time": 2000} + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + response = reporting_api.report( + token="token", event=attack_event, timeout_in_sec=5 + ) + assert response == {"success": True} + assert len(reporting_api.events) == 1 + + response = reporting_api.report( + token="token", event=non_attack_event, timeout_in_sec=5 + ) + assert response == {"success": True} + assert len(reporting_api.events) == 1 # Non-attack event should not be stored + + +def test_report_event_expiry(reporting_api): + """Test that events expire correctly based on the time interval.""" + event1 = {"type": "detected_attack", "time": 1000} + event2 = {"type": "detected_attack", "time": 2000} + reporting_api.events = [event1, event2] + + # Simulate time passing + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=12000, + ): + event3 = {"type": "detected_attack", "time": 11000} + response = reporting_api.report(token="token", event=event3, timeout_in_sec=5) + assert response == {"error": "timeout", "success": False} + assert ( + len(reporting_api.events) == 1 + ) # One event should have expired, and the new one is added + + +def test_report_event_at_boundary(reporting_api): + """Test reporting an event at the boundary of the interval.""" + event1 = {"type": "detected_attack", "time": 1000} + reporting_api.events = [event1] + + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=10000, + ): + event2 = {"type": "detected_attack", "time": 10000} # Exactly at the boundary + response = reporting_api.report(token="token", event=event2, timeout_in_sec=5) + assert response == {"error": "timeout", "success": False} + assert ( + len(reporting_api.events) == 2 + ) # Should be added since it's at the boundary + + +def test_report_invalid_event_type(reporting_api): + """Test reporting with an invalid event type.""" + event = {"type": "invalid_event", "time": 1000} + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + response = reporting_api.report(token="token", event=event, timeout_in_sec=5) + assert response == {"success": True} + assert ( + len(reporting_api.events) == 0 + ) # Invalid event types should not be stored + mock_report.assert_called_once_with("token", event, 5) + + +def test_report_no_events(reporting_api): + """Test reporting when no events have been reported yet.""" + event = {"type": "detected_attack", "time": 1000} + + with patch.object( + ReportingApiHTTP, "report", return_value={"success": True} + ) as mock_report: + with patch( + "aikido_firewall.helpers.get_current_unixtime_ms.get_unixtime_ms", + return_value=2000, + ): + response = reporting_api.report( + token="token", event=event, timeout_in_sec=5 + ) + assert response == {"success": True} + assert len(reporting_api.events) == 1 # Should add the first event + mock_report.assert_called_once_with("token", event, 5)