Send heartbeat after one minute if "receivedAnyStats: false"

aikido_firewall/background_process/reporter/__init__.py

@@ -26,14 +26,21 @@
 
     timeout_in_sec = 5  # Timeout of API calls to Aikido Server
     heartbeat_secs = 600  # Heartbeat every 10 minutes
+    initial_stats_timeout = 60  # Wait 60 seconds after startup for initial stats
 
     def __init__(self, block, api, token, serverless):
         self.block = block
         self.api = api
         self.token = token  # Should be instance of the Token class!
         self.routes = Routes(200)
         self.hostnames = Hostnames(200)
-        self.conf = ServiceConfig([], get_unixtime_ms(), [], [])
+        self.conf = ServiceConfig(
+            endpoints=[],
+            last_updated_at=get_unixtime_ms(),
+            blocked_uids=[],
+            bypassed_ips=[],
+            received_any_stats=True,
+        )
         self.rate_limiter = RateLimiter(
             max_items=5000, time_to_live_in_ms=120 * 60 * 1000  # 120 minutes
         )
@@ -55,9 +62,21 @@
                 "Token was invalid, not starting heartbeats and realtime polling."
             )
             return
+        event_scheduler.enter(self.initial_stats_timeout, 1, self.report_initial_stats)
         send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler)
         start_polling_for_changes(self, event_scheduler)
 
+    def report_initial_stats(self):
+        """
+        This is run 1m after startup, and checks if we should send out
+        a preliminary heartbeat with some stats.
+        """
+        should_report_initial_stats = not (
+            self.statistics.is_empty() or self.conf.received_any_stats
+        )
+        if should_report_initial_stats:
+            self.send_heartbeat()
+
     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, blocked, stack)

aikido_firewall/background_process/reporter/update_service_config.py

@@ -16,12 +16,11 @@
         logger.debug("Updating blocking, setting blocking to : %s", res["block"])
         reporter.block = bool(res["block"])
 
-    if res["endpoints"]:
-        if not isinstance(res["endpoints"], list):
-            res["endpoints"] = []  # Empty list
+    if res.get("endpoints") is not None:
         reporter.conf = ServiceConfig(
-            endpoints=res["endpoints"],
-            last_updated_at=get_unixtime_ms(),
-            blocked_uids=res["blockedUserIds"],
-            bypassed_ips=res["allowedIPAddresses"],
+            endpoints=res.get("endpoints", []),
+            last_updated_at=res.get("configUpdatedAt", get_unixtime_ms()),
+            blocked_uids=res.get("blockedUserIds", []),
+            bypassed_ips=res.get("allowedIPAddresses", []),
+            received_any_stats=res.get("receivedAnyStats", True),
         )

aikido_firewall/background_process/service_config.py

@@ -1,5 +1,5 @@
 """
-?
+Exports ServiceConfig class
 """
 
 from aikido_firewall.helpers.match_endpoint import match_endpoint
@@ -8,11 +8,14 @@
 class ServiceConfig:
     """Class holding the config of the reporter"""
 
-    def __init__(self, endpoints, last_updated_at, blocked_uids, bypassed_ips):
+    def __init__(
+        self, endpoints, last_updated_at, blocked_uids, bypassed_ips, received_any_stats
+    ):
         self.endpoints = [endpoint for endpoint in endpoints if not endpoint["graphql"]]
         self.last_updated_at = last_updated_at
         self.bypassed_ips = set(bypassed_ips)
         self.blocked_uids = set(blocked_uids)
+        self.received_any_stats = bool(received_any_stats)
 
     def get_endpoint(self, route_metadata):
         """

aikido_firewall/background_process/service_config_test.py

@@ -40,7 +40,11 @@ def test_service_config_initialization():
     ]
     last_updated_at = "2023-10-01"
     service_config = ServiceConfig(
-        endpoints, last_updated_at, ["0", "0", "1", "5"], ["5", "1", "2", "1", "5"]
+        endpoints,
+        last_updated_at,
+        ["0", "0", "1", "5"],
+        ["5", "1", "2", "1", "5"],
+        True,
     )
 
     # Check that non-GraphQL endpoints are correctly filtered
@@ -67,6 +71,7 @@ def service_config():
         last_updated_at="2023-10-01T00:00:00Z",
         blocked_uids=["user1", "user2"],
         bypassed_ips=["192.168.1.1", "10.0.0.1"],
+        received_any_stats=True,
     )
 
 

aikido_firewall/background_process/statistics/__init__.py

@@ -70,3 +70,11 @@ def on_inspected_call(self, *args, **kwargs):
     def get_stats(self):
         """This will return the stats as a dict, from a Statistics class"""
         return get_stats(self)
+
+    def is_empty(self):
+        """This will return a boolean value indicating if the stats are empty"""
+        return (
+            len(self.stats) == 0
+            and self.requests["total"] == 0
+            and self.requests["attacksDetected"]["total"] == 0
+        )

aikido_firewall/background_process/statistics/init_test.py

@@ -361,3 +361,42 @@ def test_it_keeps_track_of_aborted_requests(stats):
             },
         },
     }
+
+
+def test_is_empty_when_stats_are_empty(stats):
+    assert stats.is_empty() is True
+
+
+def test_is_empty_when_requests_are_empty(stats):
+    stats.requests["total"] = 0
+    stats.requests["attacksDetected"]["total"] = 0
+    stats.stats = {}  # Assuming stats is a dictionary
+    assert stats.is_empty() is True
+
+
+def test_is_empty_when_requests_have_data(stats):
+    stats.requests["total"] = 1
+    stats.requests["attacksDetected"]["total"] = 0
+    stats.stats = {}
+    assert stats.is_empty() is False
+
+
+def test_is_empty_when_attacks_detected(stats):
+    stats.requests["total"] = 0
+    stats.requests["attacksDetected"]["total"] = 1
+    stats.stats = {}
+    assert stats.is_empty() is False
+
+
+def test_is_empty_when_stats_have_data(stats):
+    stats.requests["total"] = 0
+    stats.requests["attacksDetected"]["total"] = 0
+    stats.stats = {"some_stat": 1}  # Adding some data to stats
+    assert stats.is_empty() is False
+
+
+def test_is_empty_when_all_data_present(stats):
+    stats.requests["total"] = 1
+    stats.requests["attacksDetected"]["total"] = 1
+    stats.stats = {"some_stat": 1}
+    assert stats.is_empty() is False