(WIP) Improve readability of e2e test cases

.github/workflows/end2end.yml

@@ -25,16 +25,15 @@ jobs:
     strategy:
       matrix:
         app:
-          - { name: django-mysql, testfile: end2end/django_mysql_test.py }
-          - { name: django-mysql-gunicorn, testfile: end2end/django_mysql_gunicorn_test.py }
-          - { name: django-postgres-gunicorn, testfile: end2end/django_postgres_gunicorn_test.py }
-          - { name: flask-mongo, testfile: end2end/flask_mongo_test.py }
+          - { name: django-mysql, testfile: end2end/django_mysql.py }
+          - { name: django-mysql-gunicorn, testfile: end2end/django_mysql_gunicorn.py }
+          - { name: django-postgres-gunicorn, testfile: end2end/django_postgres_gunicorn.py }
+          - { name: flask-mongo, testfile: end2end/flask_mongo.py }
           - { name: flask-mysql, testfile: end2end/flask_mysql_test.py }
-          - { name: flask-mysql-uwsgi, testfile: end2end/flask_mysql_uwsgi_test.py }
-          - { name: flask-postgres, testfile: end2end/flask_postgres_test.py }
-          - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml_test.py }
-          - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml_lxml_test.py }
-          - { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn_test.py }
+          - { name: flask-mysql-uwsgi, testfile: end2end/flask_mysql_uwsgi.py }
+          - { name: flask-postgres, testfile: end2end/flask_postgres.py }
+          - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml.py }
+          - { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn.py }
           - { name: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_test.py }
         python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
@@ -65,7 +64,7 @@ jobs:
       - name: Start application
         working-directory: ./sample-apps/${{ matrix.app.name }}
         run: |
-          nohup make run > output.log & tail -f output.log & sleep 20
-          nohup make runZenDisabled & sleep 20
+          nohup make run > output.log & sleep 1
+          nohup make runZenDisabled & sleep 1
       - name: Run end2end tests for application
-        run: tail -f  ./sample-apps/${{ matrix.app.name }}/output.log & poetry run pytest ./${{ matrix.app.testfile }}
+        run: tail -f  ./sample-apps/${{ matrix.app.name }}/output.log & poetry run python ./${{ matrix.app.testfile }}

end2end/django_mysql.py

@@ -0,0 +1,80 @@
+from end2end.utils import assert_eq
+from utils import App, Request
+
+django_mysql_app = App(8080)
+
+django_mysql_app.add_payload(
+    "test_sql_injection",
+    safe_request=Request("/app/create", data_type="form", body={"dog_name": "Bobby Tables"}),
+    unsafe_request=Request("/app/create", data_type="form", body={"dog_name": 'Dangerous bobby", 1); -- '}),
+    test_event={
+        "blocked": True,
+        "kind": "sql_injection",
+        'metadata': {
+            'dialect': 'mysql',
+            'sql': 'INSERT INTO sample_app_dogs (dog_name, dog_boss) VALUES ("Dangerous bobby", 1); -- ", "N/A")'
+        },
+        'operation': 'MySQLdb.Cursor.execute',
+        'pathToPayload': '.dog_name.[0]',
+        'payload': '"Dangerous bobby\\", 1); -- "',
+        'source': "body",
+    }
+)
+
+django_mysql_app.add_payload(
+    "test_shell_injection",
+    safe_request=None,  # Don't test safeness
+    unsafe_request=Request("/app/shell/ls -la", "GET"),
+    test_event={
+        "blocked": True,
+        "kind": "shell_injection",
+        'metadata': {'command': 'ls -la'},
+        'operation': 'subprocess.Popen',
+        'pathToPayload': '.[0]',
+        'payload': '"ls -la"',
+        'source': "route_params",
+    }
+)
+
+
+def test_heartbeat(app):
+    heartbeat = app.get_heartbeat()
+    route1 = heartbeat["routes"][0]
+    stats = heartbeat["stats"]
+    packages = set(map(lambda x: x["name"], heartbeat["packages"]))
+
+    # Validate routes
+    assert_eq(route1["path"], "/app/create")
+    assert_eq(route1["method"], "POST")
+    assert_eq(route1["hits"], 1)
+
+    assert_eq(route1["apispec"]["body"]["type"], "form-urlencoded")
+    assert_eq(route1["apispec"]["body"]["schema"], {
+        'type': 'object',
+        'properties': {
+            'dog_name': {
+                'items': {'type': 'string'},
+                'type': 'array'
+            }
+        }
+    })
+    assert_eq(route1["apispec"]["query"], None)
+    assert_eq(route1["apispec"]["auth"], None)
+
+    # Validate stats
+    assert_eq(stats["requests"]["attacksDetected"]["blocked"], 2)
+    assert_eq(stats["requests"]["attacksDetected"]["total"], 2)
+    # There are 3-4 requests :
+    # 1. is website live request, first request not always counted
+    # 2. /app/create safe
+    # 3. /app/create sql inj
+    # 4. /app/shell/ls -la shell inj
+    total = stats["requests"]["total"]
+    assert 3 <= total <= 4, f"Unexpected amount of total requests {total}"
+
+    # Validate packages
+    assert_eq(packages, {'wrapt', 'asgiref', 'aikido_zen', 'django', 'sqlparse', 'regex', 'mysqlclient'})
+
+
+django_mysql_app.test_all_payloads()
+test_heartbeat(django_mysql_app)

end2end/django_mysql_gunicorn.py

@@ -0,0 +1,23 @@
+from utils import App, Request
+
+django_mysql_gunicorn_app = App(8082)
+
+django_mysql_gunicorn_app.add_payload(
+    "test_sql_injection",
+    safe_request=Request("/app/create/", data_type="form", body={"dog_name": "Bobby Tables"}),
+    unsafe_request=Request("/app/create/", data_type="form", body={"dog_name": 'Dangerous bobby", 1); -- '}),
+    test_event={
+        "blocked": True,
+        "kind": "sql_injection",
+        'metadata': {
+            'dialect': 'mysql',
+            'sql': 'INSERT INTO sample_app_dogs (dog_name, dog_boss) VALUES ("Dangerous bobby", 1); -- ", "N/A")'
+        },
+        'operation': 'MySQLdb.Cursor.execute',
+        'pathToPayload': '.dog_name.[0]',
+        'payload': '"Dangerous bobby\\", 1); -- "',
+        'source': "body",
+    }
+)
+
+django_mysql_gunicorn_app.test_all_payloads()

end2end/django_postgres_gunicorn.py

@@ -0,0 +1,45 @@
+from utils import App, Request
+
+django_postgres_gunicorn_app = App(8100)
+
+django_postgres_gunicorn_app.add_payload(
+    "test_sql_injection",
+    safe_request=Request("/app/create", data_type="form", body={"dog_name": "Bobby Tables"}),
+    unsafe_request=Request("/app/create", data_type="form", body={"dog_name": "Dangerous bobby', TRUE); -- "}),
+    test_event={
+        "blocked": True,
+        "kind": "sql_injection",
+        'metadata': {
+            'dialect': "postgres",
+            'sql': "INSERT INTO sample_app_Dogs (dog_name, is_admin) VALUES ('Dangerous bobby', TRUE); -- ', FALSE)"
+        },
+        'operation': "psycopg2.Connection.Cursor.execute",
+        'pathToPayload': '.dog_name.[0]',
+        'payload': "\"Dangerous bobby', TRUE); -- \"",
+        'source': "body",
+    }
+)
+
+django_postgres_gunicorn_app.add_payload(
+    "test_sql_injection_via_cookies",
+    safe_request=Request("/app/create/via_cookies", "GET", headers={
+        "Cookie": "dog_name=Safe Dog"
+    }),
+    unsafe_request=Request("/app/create/via_cookies", "GET", headers={
+        "Cookie": "dog_name=Dangerous bobby', TRUE) --; ,2=2"
+    }),
+    test_event={
+        "blocked": True,
+        "kind": "sql_injection",
+        'metadata': {
+            'dialect': "postgres",
+            'sql': "INSERT INTO sample_app_Dogs (dog_name, is_admin) VALUES ('Dangerous bobby', TRUE) --', FALSE)"
+        },
+        'operation': "psycopg2.Connection.Cursor.execute",
+        'pathToPayload': '.dog_name',
+        'payload': "\"Dangerous bobby', TRUE) --\"",
+        'source': "cookies",
+    }
+)
+
+django_postgres_gunicorn_app.test_all_payloads()