Merge branch 'python-fix-importlib-recursion-bug' into add-aws-lambda-support-beta

Author: bitterpanda63

.github/workflows/publish.yml

@@ -27,7 +27,10 @@ jobs:
           echo "Pylint check failed. Please fix the issues."
           exit 1
         fi
-
+    - name: Build and start aikido mock server
+      run: |
+        cd end2end/server && docker build -t mock_core .
+        docker run --name mock_core -d -p 5050:5000 mock_core
     - name: Run Black Check
       run: |
         poetry run black --check --diff aikido_zen/

.github/workflows/test-publish.yml

@@ -27,7 +27,10 @@ jobs:
           echo "Pylint check failed. Please fix the issues."
           exit 1
         fi
-
+    - name: Build and start aikido mock server
+      run: |
+        cd end2end/server && docker build -t mock_core .
+        docker run --name mock_core -d -p 5050:5000 mock_core
     - name: Run Black Check
       run: |
         poetry run black --check --diff aikido_zen/

.github/workflows/unit-test.yml

@@ -22,6 +22,10 @@ jobs:
     - name: Add local.aikido.io to /etc/hosts
       run: |
           sudo echo "127.0.0.1 local.aikido.io" | sudo tee -a /etc/hosts
+    - name: Build and start aikido mock server
+      run: |
+        cd end2end/server && docker build -t mock_core .
+        docker run --name mock_core -d -p 5050:5000 mock_core
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:

Makefile

@@ -42,7 +42,7 @@ benchmark:
 
 
 # Binaries cache :
-BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.45
+BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.50
 FILES = \
     libzen_internals_aarch64-apple-darwin.dylib \
     libzen_internals_aarch64-apple-darwin.dylib.sha256sum \

README.md

@@ -60,6 +60,7 @@ Zen instruments the following AI SDKs to track which models are used and how man
 * ✅ [`anthropic`](https://pypi.org/project/anthropic/)
 * ✅ [`mistralai`](https://pypi.org/project/mistralai) ^1.0.0
 * ✅ [`boto3`](https://pypi.org/project/boto3) (AWS Bedrock)
+* ✅ [`groq`](https://pypi.org/project/groq)
 
 Zen is compatible with Python 3.8-3.13 and can run on Windows, Linux, and Mac OS X.
 

aikido_zen/__init__.py

@@ -9,7 +9,9 @@
 # Re-export functions :
 from aikido_zen.lambda_helper import protect_lambda
 from aikido_zen.context.users import set_user
+from aikido_zen.helpers.check_gevent import check_gevent
 from aikido_zen.middleware import should_block_request
+from aikido_zen.middleware.set_rate_limit_group import set_rate_limit_group
 
 # Import logger
 from aikido_zen.helpers.logging import logger
@@ -35,6 +37,8 @@ def protect(mode="daemon", token=""):
         return
     if not test_uds_file_access():
         return  # Unable to start background process
+    if check_gevent():
+        return  # gevent is not compatible with zen.
 
     if token:
         os.environ["AIKIDO_TOKEN"] = token

aikido_zen/background_process/__init__.py

@@ -49,7 +49,10 @@ def start_background_process():
 
     #  Daemon is set to True so that the process kills itself when the main process dies
     background_process = Process(
-        target=AikidoBackgroundProcess, args=(comms.address, comms.key), daemon=True
+        target=AikidoBackgroundProcess,
+        args=(comms.address, comms.key),
+        name="zen-agent-process",
+        daemon=True,
     )
     background_process.start()
 

aikido_zen/background_process/aikido_background_process.py

@@ -3,6 +3,7 @@
 """
 
 import multiprocessing.connection as con
+import signal
 import time
 import sched
 import traceback
@@ -40,12 +41,14 @@ def __init__(self, address, key):
         self.queue = Queue()
         self.connection_manager = None
         # Start reporting thread :
-        Thread(target=self.reporting_thread).start()
+        Thread(target=self.reporting_thread, daemon=True).start()
 
         logger.debug(
             "Background process started successfully, with UDS File : %s", address
         )
 
+        add_exit_handlers()
+
         while True:
             conn = listener.accept()
             while True:
@@ -106,3 +109,17 @@ def send_to_connection_manager(self, event_scheduler):
                 blocked=queue_attack_item[2],
                 stack=queue_attack_item[3],
             )
+
+
+def add_exit_handlers():
+    """
+    We add graceful exit handlers here since the process keeps hanging otherwise.
+    """
+
+    def exit_gracefully(sig, frame):
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, exit_gracefully)
+    signal.signal(signal.SIGTERM, exit_gracefully)
+    signal.signal(signal.SIGQUIT, exit_gracefully)
+    signal.signal(signal.SIGHUP, exit_gracefully)

aikido_zen/background_process/api/__init__.py

@@ -2,14 +2,15 @@
 init.py file for api/ folder. Includes abstract class ReportingApi
 """
 
-import json
+from aikido_zen.background_process.requests.make_request import Response
 from aikido_zen.helpers.logging import logger
 
 
 class ReportingApi:
     """This is the super class for the reporting API's"""
 
-    def to_api_response(self, res):
+    @staticmethod
+    def to_api_response(res: Response):
         """Converts results into an Api response obj"""
         status = res.status_code
         if status == 429:
@@ -18,10 +19,11 @@ def to_api_response(self, res):
             return {"success": False, "error": "invalid_token"}
         elif status == 200:
             try:
-                return json.loads(res.text)
+                return res.json()
             except Exception as e:
-                logger.debug(e)
-                logger.debug(res.text)
+                logger.debug(
+                    "Trying to load json, failed: %s (body=%s)", str(e), res.text
+                )
         return {"success": False, "error": "unknown_error"}
 
     def report(self, token, event, timeout_in_sec):

aikido_zen/background_process/api/http_api.py

@@ -2,8 +2,9 @@
 Exports the HTTP API class
 """
 
-import requests
+import aikido_zen.background_process.requests as requests
 from aikido_zen.background_process.api import ReportingApi
+from aikido_zen.background_process.requests.errors import TimeoutExceeded
 from aikido_zen.helpers.logging import logger
 
 
@@ -21,11 +22,10 @@ def report(self, token, event, timeout_in_sec):
                 timeout=timeout_in_sec,
                 headers=get_headers(token),
             )
-        except requests.exceptions.ConnectionError as e:
-            logger.error(e)
+        except TimeoutExceeded as e:
             return {"success": False, "error": "timeout"}
         except Exception as e:
-            logger.error(e)
+            logger.error("Failed to report event : %s(%s)", str(e.__class__), str(e))
             return {"success": False, "error": "unknown"}
         return self.to_api_response(res)
 
@@ -47,11 +47,10 @@ def fetch_firewall_lists(self, token):
                     "Authorization": str(token),
                 },
             )
-        except requests.exceptions.ConnectionError as e:
-            logger.error(e)
+        except TimeoutExceeded as e:
             return {"success": False, "error": "timeout"}
         except Exception as e:
-            logger.error(e)
+            logger.error("Failed to fetch firewall lists: %s", str(e))
             return {"success": False, "error": "unknown"}
         return self.to_api_response(res)