Merge branch 'main' into AIK-3613

.codecov.yml

@@ -1 +1,6 @@
 comment: false
+ignore:
+  - "*_test.py"
+  - "sample-apps/"
+  - "end2end/"
+  - "benchmarks/"

.github/workflows/benchmark.yml

@@ -42,6 +42,7 @@ jobs:
       - name: Start flask-mysql
         working-directory: ./sample-apps/flask-mysql
         run: |
+          cp .env.benchmark .env.example
           docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
       - name: Install K6
         uses: grafana/setup-k6-action@v1
@@ -60,14 +61,12 @@ jobs:
       - name: Start starlette multi-threaded
         working-directory: ./sample-apps/starlette-postgres-uvicorn
         run: |
+          cp .env.benchmark .env.example
           docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
       - name: Install wrk
         run: |
-          sudo apt-get install build-essential libssl-dev git -y
-          git clone https://github.com/wg/wrk.git wrk
-          cd wrk
-          sudo make
-          sudo cp wrk /usr/local/bin
+          sudo apt-get update
+          sudo apt-get install -y wrk
       - name: Set up Python 3.9
         uses: actions/setup-python@v2
         with:

.github/workflows/end2end.yml

@@ -8,59 +8,121 @@ jobs:
     steps:
     - name: Checkout code
       uses: actions/checkout@v2
+
     - name: Start databases
       working-directory:  ./sample-apps/databases
       run: docker compose up --build -d
+    - name: Build and start aikido mock server
+      working-directory: ./end2end/server
+      run: docker build -t mock-core . && docker run --name mock_core -d -p 5000:5000 mock-core
+
+    - 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
+
+    # django-mysql
     - 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: Run end2end tests for django-mysql
+      run: sleep 5 && poetry run pytest ./end2end/django_mysql_test.py
+
+    # django-mysql-gunicorn
+    - name: Restart mock server
+      run: docker restart mock_core
     - 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: Run end2end tests for django-mysql-gunicorn
+      run: sleep 5 && poetry run pytest ./end2end/django_mysql_gunicorn_test.py
+
+    # django-postgres-gunicorn
+    - name: Restart mock server
+      run: docker restart mock_core
     - name: Start django-postgres-gunicorn
       working-directory: ./sample-apps/django-postgres-gunicorn
       run: |
         docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Run end2end tests for django-postgres-gunicorn
+      run: sleep 5 && poetry run pytest ./end2end/django_postgres_gunicorn_test.py
+
+    # flask-mongo
+    - name: Restart mock server
+      run: docker restart mock_core
     - 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: Run end2end tests for flask-mongo
+      run: sleep 5 && poetry run pytest ./end2end/flask_mongo_test.py
+
+    # flask-mysql
+    - name: Restart mock server
+      run: docker restart mock_core
     - 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: Run end2end tests for flask-mysql
+      run: sleep 5 && poetry run pytest ./end2end/flask_mysql_test.py
+
+    # flask-mysql-uwsgi
+    - name: Restart mock server
+      run: docker restart mock_core
     - 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: Run end2end tests for flask-mysql-uwsgi
+      run: sleep 5 && poetry run pytest ./end2end/flask_mysql_uwsgi_test.py
+
+    # flask-postgres
+    - name: Restart mock server
+      run: docker restart mock_core
     - 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: Run end2end tests for flask-postgres
+      run: sleep 5 && poetry run pytest ./end2end/flask_postgres_test.py
+
+    # flask-postgres-xml
+    - name: Restart mock server
+      run: docker restart mock_core
     - name: Start flask-postgres-xml
       working-directory: ./sample-apps/flask-postgres-xml
       run: |
         docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Run end2end tests for flask-postgres-xml
+      run: | 
+        sleep 5
+        poetry run pytest ./end2end/flask_postgres_xml_test.py
+        docker restart mock_core
+        poetry run pytest ./end2end/flask_postgres_xml_lxml_test.py
+
+    # quart-postgres-uvicorn
+    - name: Restart mock server
+      run: docker restart mock_core
     - name: Start quart-postgres-uvicorn
       working-directory: ./sample-apps/quart-postgres-uvicorn
       run: |
         docker compose -f docker-compose.yml -f docker-compose.benchmark.yml up --build -d
+    - name: Run end2end tests for quart-postgres-uvicorn
+      run: sleep 5 && poetry run pytest ./end2end/quart_postgres_uvicorn_test.py
+
+    # starlette-postgres-uvicorn
+    - name: Restart mock server
+      run: docker restart mock_core
     - name: Start starlette-postgres-uvicorn
       working-directory: ./sample-apps/starlette-postgres-uvicorn
       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
+    - name: Run end2end tests for starlette-postgres-uvicorn
+      run: sleep 5 && poetry run pytest ./end2end/starlette_postgres_uvicorn_test.py

.github/workflows/publish.yml

@@ -16,10 +16,8 @@ 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: Install poetry
-      run: pip install poetry
-    - name: Install dependencies
-      run: poetry install
+    - name: Installation
+      run: make install
     - name: Run Pylint
       run: |
         poetry run pylint --fail-under=9 --rcfile=.pylintrc aikido_zen/
@@ -50,9 +48,8 @@ jobs:
     - uses: actions/checkout@v4
     - name: Install poetry
       run: pip install poetry
-    - name: Install dependencies
-      run: poetry install
-
+    - name: Download binaries
+      run: make binaries
     - name: Publish to PyPI
       env:
         POETRY_HTTP_BASIC_PYPI_USERNAME: __token__

.github/workflows/test-publish.yml

@@ -17,10 +17,8 @@ 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: Install poetry
-      run: pip install poetry
-    - name: Install dependencies
-      run: poetry install
+    - name: Installation
+      run: make install
     - name: Run Pylint
       run: |
         poetry run pylint --fail-under=9 --rcfile=.pylintrc aikido_zen/
@@ -42,45 +40,30 @@ jobs:
       run: |
         make test
   build:
-    name: Build distribution 📦
+    name: Build distribution 📦 and Publish to TestPyPI
+    needs:
+      - tests
     runs-on: ubuntu-latest
 
     steps:
     - uses: actions/checkout@v4
     - name: Install poetry
       run: pip install poetry
+    - name: Download binaries
+      run: make binaries
     - name: Install dependencies
       run: poetry install
 
-    - name: Build distribution packages
-      run: poetry build
-
+    - name: Publish to TestPyPI
+      env:
+        POETRY_HTTP_BASIC_PYPI_USERNAME: __token__
+        POETRY_HTTP_BASIC_PYPI_PASSWORD: ${{ secrets.TEST_PYPI_TOKEN }}
+      run: |
+        poetry config repositories.test-pypi https://test.pypi.org/legacy/
+        poetry config pypi-token.test-pypi ${{ secrets.TEST_PYPI_TOKEN }}
+        poetry publish -r test-pypi --build
     - name: Store the distribution packages
       uses: actions/upload-artifact@v3
       with:
         name: python-package-distributions
         path: dist/
-  publish-to-testpypi:
-    name: Publish Python 🐍 distribution 📦 to TestPyPI
-    needs:
-    - build
-    - tests
-    runs-on: ubuntu-latest
-
-    environment:
-      name: testpypi
-      url: https://test.pypi.org/p/aikido_zen
-
-    permissions:
-      id-token: write  # IMPORTANT: mandatory for trusted publishing
-
-    steps:
-    - name: Download all the dists
-      uses: actions/download-artifact@v4.1.7
-      with:
-        name: python-package-distributions
-        path: dist/
-    - name: Publish distribution 📦 to TestPyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
-      with:
-        repository-url: https://test.pypi.org/legacy/

.github/workflows/unit-test.yml

@@ -27,7 +27,6 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         make install
-
     - name: Run tests with coverage
       run: |
         make cov

Makefile

@@ -1,5 +1,5 @@
-.PHONY: build
-build:
+build: binaries _build
+_build:
 	poetry build
 
 .PHONY: clean
@@ -11,14 +11,12 @@ lint:
 	poetry run black aikido_zen/
 	poetry run pylint aikido_zen/
 
-.PHONY: install
-install:
+install: binaries
 	pip install poetry
 	poetry install
 
-.PHONY: dev_install
-dev_install:
-	pip install poetry
+dev_install: install _dev_install
+_dev_install:
 	poetry install --with=dev
 
 
@@ -37,3 +35,24 @@ cov:
 .PHONY: benchmark
 benchmark:
 	k6 run -q ./benchmarks/flask-mysql-benchmarks.js
+
+BASE_URL = https://github.com/AikidoSec/zen-internals/releases/download/v0.1.34
+FILES = \
+    libzen_internals_aarch64-apple-darwin.dylib \
+    libzen_internals_aarch64-apple-darwin.dylib.sha256sum \
+    libzen_internals_aarch64-unknown-linux-gnu.so \
+    libzen_internals_aarch64-unknown-linux-gnu.so.sha256sum \
+    libzen_internals_x86_64-apple-darwin.dylib \
+    libzen_internals_x86_64-apple-darwin.dylib.sha256sum \
+    libzen_internals_x86_64-pc-windows-gnu.dll \
+    libzen_internals_x86_64-pc-windows-gnu.dll.sha256sum \
+    libzen_internals_x86_64-unknown-linux-gnu.so \
+    libzen_internals_x86_64-unknown-linux-gnu.so.sha256sum
+
+binaries: binaries_make_dir $(addprefix aikido_zen/lib/, $(FILES))
+binaries_make_dir:
+	rm -rf aikido_zen/lib
+	mkdir -p aikido_zen/lib/
+aikido_zen/lib/%:
+	@echo "Downloading $*..."
+	curl -L -o $@ $(BASE_URL)/$*

README.md

@@ -18,7 +18,7 @@ Zen protects your Python apps by preventing user input containing dangerous stri
 Zen will autonomously protect your Python applications from the inside against:
 
 * 🛡️ [NoSQL injection attacks](https://www.aikido.dev/blog/web-application-security-vulnerabilities)
-* 🛡️ [SQL injection attacks]([https://www.aikido.dev/blog/web-application-security-vulnerabilities](https://owasp.org/www-community/attacks/SQL_Injection))
+* 🛡️ [SQL injection attacks](https://www.aikido.dev/blog/the-state-of-sql-injections)
 * 🛡️ [Command injection attacks](https://owasp.org/www-community/attacks/Command_Injection)
 * 🛡️ [Path traversal attacks](https://owasp.org/www-community/attacks/Path_Traversal)
 * 🛡️ [Server-side request forgery (SSRF)](./docs/ssrf.md)
@@ -37,6 +37,7 @@ Zen for Python 3 is compatible with:
 * ✅ [Flask](docs/flask.md)
 * ✅ [Quart](docs/quart.md)
 * ✅ [Starlette](docs/starlette.md)
+* ✅ [FastAPI](docs/starlette.md) (via Starlette)
 
 
 ### WSGI servers

aikido_zen/__init__.py

@@ -6,18 +6,19 @@
 
 from dotenv import load_dotenv
 
-# Re-export set_current_user :
+# Re-export functions :
 from aikido_zen.context.users import set_user
+from aikido_zen.middleware import should_block_request
 
 # Import logger
 from aikido_zen.helpers.logging import logger
 
 # Import background process
 from aikido_zen.background_process import start_background_process
 
-# Load environment variables and constants
 # Load environment variables and constants
 from aikido_zen.config import PKG_VERSION
+from aikido_zen.helpers.aikido_disabled_flag_active import aikido_disabled_flag_active
 
 load_dotenv()
 
@@ -30,6 +31,9 @@ def protect(mode="daemon"):
     - daemon_disabled : This will import sinks/sources but won't start a background process
     Protect user's application
     """
+    if aikido_disabled_flag_active():
+        # Do not run any aikido code when the disabled flag is on
+        return
     if mode in ("daemon", "daemon_only"):
         start_background_process()
     if mode == "daemon_only":

aikido_zen/api_discovery/README.md

@@ -1,3 +1,2 @@
 # Feature flag
-
-This feature is currently disabled by default. Enable it by setting the environment variable `AIKIDO_FEATURE_COLLECT_API_SCHEMA` to `true`.
+This feature is now on by default.

aikido_zen/api_discovery/get_api_info.py

@@ -10,9 +10,6 @@
 def get_api_info(context):
     """Generates an apispec based on the context passed along"""
     try:
-        # Check if feature flag COLLECT_API_SCHEMA is enabled
-        if not is_feature_enabled("COLLECT_API_SCHEMA"):
-            return {}
         body_info = get_body_info(context)
         query_info = get_query_info(context)
         auth_info = get_auth_types(context)