Skip stats when request is from a bypassed IP

Author: tomaisthorpe

aikido_zen/sources/functions/request_handler.py

@@ -19,7 +19,8 @@ def request_handler(stage, status_code=0):
     try:
         if stage == "init":
             cache = get_cache()
-            if ctx.get_current_context() and cache:
+            context = ctx.get_current_context()
+            if context and cache and not cache.is_bypassed_ip(context.remote_address):
                 cache.stats.increment_total_hits()
         if stage == "pre_response":
             return pre_response()
@@ -102,6 +103,9 @@ def post_response(status_code):
     if not cache:
         return
 
+    if cache.is_bypassed_ip(context.remote_address):
+        return
+
     attack_wave = attack_wave_detector_store.is_attack_wave(context)
     if attack_wave:
         cache.stats.on_detected_attack_wave(blocked=False)

aikido_zen/sources/functions/request_handler_test.py

@@ -131,6 +131,43 @@ def test_post_response_no_context(mock_get_comms):
     comms.send_data_to_bg_process.assert_not_called()
 
 
+def test_bypassed_ip_no_stats_in_init():
+    cache = get_cache()
+    cache.config.set_bypassed_ips(["1.2.3.4"])
+    cache.stats.clear()
+
+    context = MagicMock()
+    context.remote_address = "1.2.3.4"
+    with patch("aikido_zen.context.get_current_context", return_value=context):
+        request_handler("init")
+
+    assert cache.stats.get_record()["requests"]["total"] == 0
+
+
+def test_non_bypassed_ip_increments_stats_in_init():
+    cache = get_cache()
+    cache.config.set_bypassed_ips([])
+    cache.stats.clear()
+
+    context = MagicMock()
+    context.remote_address = "1.2.3.4"
+    with patch("aikido_zen.context.get_current_context", return_value=context):
+        request_handler("init")
+
+    assert cache.stats.get_record()["requests"]["total"] == 1
+
+
+def test_bypassed_ip_no_route_tracking_in_post_response(mock_context):
+    cache = get_cache()
+    cache.config.set_bypassed_ips(["5.6.7.8"])
+    mock_context.remote_address = "5.6.7.8"
+
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        request_handler("post_response", status_code=200)
+
+    assert cache.routes.routes == {}
+
+
 # Test firewall lists
 def set_context(remote_address, user_agent="", route="/posts/:number"):
     headers = Headers()