Better e2e testing

.github/workflows/end2end.yml

@@ -24,17 +24,16 @@ 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: 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: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_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.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.py }
         python-version: ["3.10", "3.11", "3.12"]
     steps:
       - name: Install packages
@@ -64,7 +63,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 & tail -f output.log & sleep 30
+          nohup make runZenDisabled & sleep 30
       - 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 & python ./${{ matrix.app.testfile }}

end2end/__init__.py

@@ -0,0 +1,4 @@
+import json
+
+with open('end2end/attack_events.json', 'r') as file:
+    events = json.load(file)

end2end/attack_events.json

@@ -0,0 +1,124 @@
+{
+  "django_mysql_attack_sql": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO sample_app_dogs (dog_name, dog_boss) VALUES (\"Dangerous bobby\", 1); -- \", \"N/A\")"
+    },
+    "operation": "MySQLdb.Cursor.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Dangerous bobby\\\", 1); -- \"",
+    "source": "body"
+  },
+  "django_mysql_attack_shell": {
+    "blocked": true,
+    "kind": "shell_injection",
+    "metadata": {
+      "command": "ls -la"
+    },
+    "operation": "subprocess.Popen",
+    "pathToPayload": ".[0]",
+    "payload": "\"ls -la\"",
+    "source": "route_params"
+  },
+  "django_postgres_attack": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "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": "body"
+  },
+  "quart_postgres_attack": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Dangerous Bobby', TRUE); -- ', FALSE)"
+    },
+    "operation": "asyncpg.connection.Connection.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Dangerous Bobby', TRUE); -- \"",
+    "source": "body"
+  },
+  "flask_xml_attack": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Malicious dog', TRUE); -- ', FALSE)"
+    },
+    "operation": "psycopg2.Connection.Cursor.execute",
+    "pathToPayload": ".dog_name.[0]",
+    "payload": "\"Malicious dog', TRUE); -- \"",
+    "source": "xml"
+  },
+  "flask_mysql_attack": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES (\"Dangerous bobby\", 1); -- \", 0)"
+    },
+    "operation": "pymysql.Cursor.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Dangerous bobby\\\", 1); -- \"",
+    "source": "body"
+  },
+  "flask_mongo_attack": {
+    "blocked": true,
+    "kind": "nosql_injection",
+    "metadata": {
+      "filter": "{\"dog_name\": \"bobby_tables\", \"pswd\": {\"$ne\": \"\"}}"
+    },
+    "operation": "pymongo.collection.Collection.find",
+    "pathToPayload": ".pswd",
+    "payload": "{\"$ne\": \"\"}",
+    "source": "body"
+  },
+  "flask_postgres_attack_body": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Dangerous Bobby', TRUE); -- ', FALSE)"
+    },
+    "operation": "psycopg2.Connection.Cursor.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Dangerous Bobby', TRUE); -- \"",
+    "source": "body"
+  },
+    "flask_postgres_attack_cookie": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Bobby', TRUE) --', FALSE)"
+    },
+    "operation": "psycopg2.Connection.Cursor.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Bobby', TRUE) --\"",
+    "source": "cookies"
+  },
+  "flask_mysql_attack_sql": {
+    "blocked": true,
+    "kind": "sql_injection",
+    "metadata": {
+      "sql": "INSERT INTO dogs (dog_name, isAdmin) VALUES (\"Dangerous bobby\", 1); -- \", 0)"
+    },
+    "operation": "pymysql.Cursor.execute",
+    "pathToPayload": ".dog_name",
+    "payload": "\"Dangerous bobby\\\", 1); -- \"",
+    "source": "body"
+  },
+  "flask_mysql_attack_shell": {
+    "blocked": true,
+    "kind": "shell_injection",
+    "metadata": {
+      "command": "ls -la"
+    },
+    "operation": "subprocess.Popen",
+    "pathToPayload": ".command",
+    "payload": "\"ls -la\"",
+    "source": "route_params",
+    "user_id": "456"
+  }
+}

end2end/django_mysql.py

@@ -0,0 +1,17 @@
+from __init__ import events
+from utils import App, Request
+
+django_mysql_app = App(8080)
+
+django_mysql_app.add_payload(
+    "sql", test_event=events["django_mysql_attack_sql"],
+    safe_request=Request(route="/app/create", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/app/create", body={'dog_name': 'Dangerous bobby", 1); -- '}, data_type="form")
+)
+django_mysql_app.add_payload(
+    "shell", test_event=events["django_mysql_attack_shell"],
+    safe_request=Request(route="/app/shell/bobby", method="GET"),
+    unsafe_request=Request(route="/app/shell/ls -la", method="GET")
+)
+
+django_mysql_app.test_all_payloads()

end2end/django_mysql_gunicorn.py

@@ -0,0 +1,12 @@
+from __init__ import events
+from utils import App, Request
+
+django_mysql_gunicorn_app = App(8082)
+
+django_mysql_gunicorn_app.add_payload(
+    "sql", test_event=events["django_mysql_attack_sql"],
+    safe_request=Request(route="/app/create/", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/app/create/", body={'dog_name': 'Dangerous bobby", 1); -- '}, data_type="form")
+)
+
+django_mysql_gunicorn_app.test_all_payloads()

end2end/django_postgres_gunicorn.py

@@ -0,0 +1,12 @@
+from __init__ import events
+from utils import App, Request
+
+django_postgres_gunicorn_app = App(8100)
+
+django_postgres_gunicorn_app.add_payload(
+    "sql", test_event=events["django_postgres_attack"],
+    safe_request=Request(route="/app/create", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/app/create", body={'dog_name': "Dangerous bobby', TRUE); -- "}, data_type="form")
+)
+
+django_postgres_gunicorn_app.test_all_payloads()

end2end/flask_mongo.py

@@ -0,0 +1,24 @@
+from __init__ import events
+from utils import App, Request
+
+flask_mongo_app = App(8094)
+
+# Create dog :
+status_code_create = Request(route="/create", body={
+    "dog_name": "bobby_tables", "pswd": "bobby123"
+}, data_type="form").execute(flask_mongo_app.urls["enabled"])
+assert status_code_create == 200
+
+# Payloads :
+flask_mongo_app.add_payload(
+    "nosql", test_event=events["flask_mongo_attack"],
+    safe_request=Request("/auth", body={"dog_name": "bobby_tables", "pswd": "bobby123"}),
+    unsafe_request=Request("/auth", body={"dog_name": "bobby_tables", "pswd": { "$ne": ""}})
+)
+flask_mongo_app.add_payload(
+    "nosql_force", test_event=events["flask_mongo_attack"],
+    safe_request=Request("/auth_force", body={"dog_name": "bobby_tables", "pswd": "bobby123"}),
+    unsafe_request=Request("/auth_force", body={"dog_name": "bobby_tables", "pswd": { "$ne": ""}})
+)
+
+flask_mongo_app.test_all_payloads()

end2end/flask_mysql.py

@@ -0,0 +1,18 @@
+from __init__ import events
+from utils import App, Request
+
+flask_mysql_app = App(8086)
+
+flask_mysql_app.add_payload(
+    "sql", test_event=events["flask_mysql_attack_sql"],
+    safe_request=Request("/create", body={"dog_name": "Bobby"}, data_type="form"),
+    unsafe_request=Request("/create", body={"dog_name": "Dangerous bobby\", 1); -- "}, data_type="form")
+)
+flask_mysql_app.add_payload(
+    "shell", test_event=events["flask_mysql_attack_shell"],
+    safe_request=Request(route="/shell/bobby", method="GET"),
+    unsafe_request=Request(route="/shell/ls -la", method="GET", headers={"user": "456"})
+)
+
+flask_mysql_app.test_all_payloads()
+flask_mysql_app.test_rate_limiting()

end2end/flask_mysql_uwsgi.py

@@ -0,0 +1,12 @@
+from __init__ import events
+from utils import App, Request
+
+flask_mysql_uwsgi_app = App(8088)
+
+flask_mysql_uwsgi_app.add_payload(
+    "sql", test_event=events["flask_mysql_attack"],
+    safe_request=Request(route="/create", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/create", body={'dog_name': 'Dangerous bobby", 1); -- '}, data_type="form")
+)
+
+flask_mysql_uwsgi_app.test_all_payloads()

end2end/flask_postgres.py

@@ -0,0 +1,17 @@
+from __init__ import events
+from utils import App, Request
+
+flask_postgres_app = App(8090)
+
+flask_postgres_app.add_payload(
+    "sql", test_event=events["flask_postgres_attack_body"],
+    safe_request=Request("/create", body={"dog_name": "Bobby Tables"}, data_type="form"),
+    unsafe_request=Request("/create", body={"dog_name": "Dangerous Bobby', TRUE); -- "}, data_type="form")
+)
+flask_postgres_app.add_payload(
+    "sql_cookie", test_event=events["flask_postgres_attack_cookie"],
+    safe_request=Request("/create_with_cookie", method="GET", cookies={"dog_name": "Bobby Tables"}),
+    unsafe_request=Request("/create_with_cookie", method="GET", cookies={"dog_name": "Bobby', TRUE) -- "})
+)
+
+flask_postgres_app.test_all_payloads()

end2end/flask_postgres_xml.py

@@ -0,0 +1,17 @@
+from __init__ import events
+from utils import App, Request
+
+flask_xml_app = App(8092)
+
+flask_xml_app.add_payload(
+    "sql_with_xml", test_event=events["flask_xml_attack"],
+    safe_request=Request(route="/xml_post", body='<dogs><dog dog_name="Bobby" /></dogs>', data_type="form"),
+    unsafe_request=Request(route="/xml_post", body='<dogs><dog dog_name="Malicious dog\', TRUE); -- " /></dogs>', data_type="form")
+)
+flask_xml_app.add_payload(
+    "sql_with_lxml", test_event=events["flask_xml_attack"],
+    safe_request=Request(route="/xml_post_lxml", body='<dogs><dog dog_name="Bobby" /></dogs>', data_type="form"),
+    unsafe_request=Request(route="/xml_post_lxml", body='<dogs><dog dog_name="Malicious dog\', TRUE); -- " /></dogs>', data_type="form")
+)
+
+flask_xml_app.test_all_payloads()

end2end/quart_postgres_uvicorn.py

@@ -0,0 +1,13 @@
+from __init__ import events
+from utils import App, Request
+
+quart_postgres_app = App(8096, status_code_valid=201)
+
+quart_postgres_app.add_payload(
+    "sql", test_event=events["quart_postgres_attack"],
+    safe_request=Request(route="/create", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/create", body={'dog_name': "Dangerous Bobby', TRUE); -- "}, data_type="form")
+)
+
+quart_postgres_app.test_all_payloads()
+quart_postgres_app.test_rate_limiting()

end2end/server/mock_aikido_core.py

@@ -84,6 +84,12 @@ def mock_set_config():
 def mock_get_events():
     return jsonify(events)
 
+@app.route('/mock/reset', methods=['GET'])
+def mock_reset():
+    global events
+    events = [] # Reset events
+    return jsonify({})
+
 
 if __name__ == '__main__':
     if len(sys.argv) < 2 or len(sys.argv) > 3:

end2end/starlette_postgres_uvicorn.py

@@ -0,0 +1,13 @@
+from __init__ import events
+from utils import App, Request
+
+starlette_postgres_app = App(8102, status_code_valid=201)
+
+starlette_postgres_app.add_payload(
+    "sql", test_event=events["quart_postgres_attack"],
+    safe_request=Request(route="/create", body={'dog_name': 'Bobby'}, data_type="form"),
+    unsafe_request=Request(route="/create", body={'dog_name': "Dangerous Bobby', TRUE); -- "}, data_type="form")
+)
+
+starlette_postgres_app.test_all_payloads()
+starlette_postgres_app.test_rate_limiting()

end2end/utils/EventHandler.py

@@ -0,0 +1,20 @@
+import time
+import requests
+import json
+
+class EventHandler:
+    def __init__(self, url="http://localhost:5000"):
+        self.url = url
+    def reset(self):
+        print("Resetting stored events on mock server")
+        res = requests.get(self.url + "/mock/reset", timeout=5)
+        time.sleep(1)
+    def fetch_events_from_mock(self):
+        res = requests.get(self.url + "/mock/events", timeout=5)
+        json_events = json.loads(res.content.decode("utf-8"))
+        return json_events
+    def fetch_attacks(self):
+        return filter_on_event_type(self.fetch_events_from_mock(), "detected_attack")
+
+def filter_on_event_type(events, type):
+    return [event for event in events if event["type"] == type]

end2end/utils/__init__.py

@@ -0,0 +1,62 @@
+import time
+
+from .EventHandler import EventHandler
+from .assert_equals import assert_eq
+from .request import Request
+from .test_bot_blocking import test_bot_blocking
+from .test_ip_blocking import test_ip_blocking
+from .test_ratelimiting import test_ratelimiting, test_ratelimiting_per_user
+from .test_payloads_safe_vs_unsafe import test_payloads_safe_vs_unsafe
+
+class App:
+    def __init__(self, port, status_code_valid=200):
+        self.urls = {
+            "enabled": f"http://localhost:{port}",
+            "disabled": f"http://localhost:{port + 1}"
+        }
+        self.payloads = {}
+        self.event_handler = EventHandler()
+        self.status_code_valid = status_code_valid
+
+    def add_payload(self,key, safe_request, unsafe_request, test_event=None):
+        self.payloads[key] = {
+            "safe": safe_request,
+            "unsafe": unsafe_request,
+            "test_event": test_event
+        }
+
+    def test_payload(self, key):
+        if key not in self.payloads:
+            raise Exception("Payload " + key + " not found.")
+        payload = self.payloads.get(key)
+
+        self.event_handler.reset()
+        test_payloads_safe_vs_unsafe(payload, self.urls, status_code1=self.status_code_valid)
+        print("✅ Tested payload: " + key)
+
+        if payload["test_event"]:
+            time.sleep(5)
+            attacks = self.event_handler.fetch_attacks()
+            assert_eq(len(attacks), equals=1)
+            for k, v in payload["test_event"].items():
+                if k == "user_id":  # exemption rule for user ids
+                    assert_eq(attacks[0]["attack"]["user"]["id"], v)
+                else:
+                    assert_eq(attacks[0]["attack"][k], equals=v)
+            print("✅ Tested accurate event reporting for: " + key)
+
+    def test_all_payloads(self):
+        for key in self.payloads.keys():
+            self.test_payload(key)
+
+    def test_blocking(self):
+        #test_bot_blocking(self.urls["enabled"])
+        #print("✅ Tested bot blocking")
+        test_ip_blocking(self.urls["enabled"])
+        print("✅ Tested IP Blocking")
+
+    def test_rate_limiting(self, route="/test_ratelimiting_1"):
+        test_ratelimiting(self.urls["enabled"] + route)
+        print("✅ Tested rate-limiting")
+        test_ratelimiting_per_user(self.urls["enabled"] + route)
+        print("✅ Tested rate-limiting (User)")

end2end/utils/assert_equals.py

@@ -0,0 +1,4 @@
+def assert_eq(val1, equals, val2=None):
+    assert val1 == equals, f"Assertion failed: Expected {equals} != {val1}"
+    if val2 is not None:
+        assert val2 == equals, f"Assertion failed: Expected {equals} != {val2}"

end2end/utils/request.py

@@ -0,0 +1,30 @@
+import requests
+
+class Request:
+    def __init__(self, route, method='POST', headers=None, data_type='json', body=None, cookies={}):
+        self.method = method
+        self.route = route
+        self.headers = headers if headers is not None else {}
+        self.data_type = data_type  # 'json' or 'form'
+        self.body = body
+        self.cookies = cookies
+
+    def execute(self, base_url):
+        """Execute the request and return the status code."""
+        url = f"{base_url}{self.route}"
+
+        if self.method.upper() == 'POST':
+            if self.data_type == 'json':
+                response = requests.post(url, json=self.body, headers=self.headers, cookies=self.cookies)
+            elif self.data_type == 'form':
+                response = requests.post(url, data=self.body, headers=self.headers, cookies=self.cookies)
+            else:
+                raise ValueError("Unsupported data type. Use 'json' or 'form'.")
+        elif self.method.upper() == 'GET':
+            response = requests.get(url, headers=self.headers, cookies=self.cookies)
+        else:
+            raise ValueError("Unsupported HTTP method. Use 'GET' or 'POST'.")
+
+        return response.status_code
+    def __str__(self):
+        return f"Request(method={self.method}, route={self.route}, headers={self.headers}, data_type={self.data_type}, body={self.body})"

end2end/utils/test_bot_blocking.py

@@ -0,0 +1,36 @@
+import requests
+from utils.assert_equals import assert_eq
+
+def test_bot_blocking(url):
+    # Allowed User-Agents :
+    res = requests.get(url, headers={
+        'User-Agent': "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'User-Agent': "Mozilla/5.0 (iPhone; CPU iPhone OS 14_4 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0 Mobile/15E148 Safari/604.1"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'User-Agent': "Mozilla/5.0 (Linux; Android 10; SM-G973F) AppleWebKit/537.36 (KHTML, like Gecko) SamsungBrowser/14.0 Chrome/91.0.4472.114 Mobile Safari/537.36"
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # Blocked User-Agent :
+    res = requests.get(url, headers={
+        'User-Agent': "BYTESPIDER"
+    })
+    assert_eq(res.status_code, equals=403)
+    assert_eq(res.text, equals="You are not allowed to access this resource because you have been identified as a bot.")
+
+    # More complex blocked User-Agent :
+    res = requests.get(url, headers={
+        'User-Agent': "Mozilla/5.0 (compatible; Bytespider/1.0; +http://bytespider.com/bot.html)"
+    })
+    assert_eq(res.status_code, equals=403)
+    assert_eq(res.text, equals="You are not allowed to access this resource because you have been identified as a bot.")
+    res = requests.get(url, headers={
+        'User-Agent': "Mozilla/5.0 (compatible; AI2Bot/1.0; +http://www.aaai.org/Press/Reports/AI2Bot)"
+    })
+    assert_eq(res.status_code, equals=403)
+    assert_eq(res.text, equals="You are not allowed to access this resource because you have been identified as a bot.")

end2end/utils/test_ip_blocking.py

@@ -0,0 +1,41 @@
+import requests
+from utils.assert_equals import assert_eq
+
+def test_ip_blocking(url):
+    # Allowed IP :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # Blocked IP :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "1.2.3.4"
+    })
+    assert_eq(res.status_code, equals=403)
+    assert_eq(res.text, equals="Your IP address is blocked due to geo restrictions (Your IP: 1.2.3.4)")
+
+    # More complex X-Forwarded-For :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "invalid.ip.here.now, 1.2.3.4 "
+    })
+    assert_eq(res.status_code, equals=403)
+    assert_eq(res.text, equals="Your IP address is blocked due to geo restrictions (Your IP: 1.2.3.4)")
+
+    # More complex but safe X-Forwarded-For :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "invalid.ip.here.now, 192.168.1.1 "
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # It should only use the first valid ip
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1, 1.2.3.4"
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # It should work with an empty X-Forwarded-For
+    res = requests.get(url, headers={
+        'X-Forwarded-For': ""
+    })
+    assert_eq(res.status_code, equals=200)

end2end/utils/test_payloads_safe_vs_unsafe.py

@@ -0,0 +1,13 @@
+from . import assert_eq
+
+def test_payloads_safe_vs_unsafe(payloads, urls, status_code1=200):
+    print("Safe req to : (1) " + urls["enabled"] + payloads["safe"].route)
+    assert_eq(val1=payloads["safe"].execute(urls["enabled"]), equals=status_code1)
+
+    print("Safe req to : (0) " + urls["disabled"] + payloads["safe"].route)
+    assert_eq(val1=payloads["safe"].execute(urls["disabled"]), equals=status_code1)
+
+    print("Unsafe req to : (1) " + urls["enabled"] + payloads["unsafe"].route)
+    assert_eq(val1=payloads["unsafe"].execute(urls["enabled"]), equals=500)
+    print("Unsafe req to : (0) " + urls["disabled"] + payloads["unsafe"].route)
+    assert_eq(val1=payloads["unsafe"].execute(urls["disabled"]), equals=status_code1)

end2end/utils/test_ratelimiting.py

@@ -0,0 +1,160 @@
+import time
+import requests
+from utils.assert_equals import assert_eq
+
+def test_ratelimiting(url):
+    # Test route (First req & 2nd Req) :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # 3rd & 4th (blocked) requests :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen. (Your IP: 192.168.1.1)")
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen. (Your IP: 192.168.1.1)")
+
+    # Now do the same but with a different IP in the same time block :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen. (Your IP: 192.168.1.2)")
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen. (Your IP: 192.168.1.2)")
+
+    # Now wait 5 seconds so your window is over and re-request :
+    time.sleep(5)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.1"
+    })
+    assert_eq(res.status_code, equals=200)
+
+
+def test_ratelimiting_per_user(url):
+    # Test route (First req & 2nd Req) :
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # 3rd & 4th (blocked) requests :
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen.")
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen.")
+
+    # Now do the same but with a different User in the same time block :
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen.")
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=429)
+    assert_eq(res.text, equals="You are rate limited by Zen.")
+
+    # Now wait 5 seconds so your window is over and re-request :
+    time.sleep(5)
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id2'
+    })
+    assert_eq(res.status_code, equals=200)
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'user': 'id1'
+    })
+    assert_eq(res.status_code, equals=200)
+
+    # Test it prefers user over IP :
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2"
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2",
+        'user': 'id3'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2",
+        'user': 'id4'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2",
+        'user': 'id5'
+    })
+    assert_eq(res.status_code, equals=200)
+    res = requests.get(url, headers={
+        'X-Forwarded-For': "192.168.1.2",
+        'user': 'id6'
+    })
+    assert_eq(res.status_code, equals=200)

sample-apps/flask-mysql/app.py

@@ -13,6 +13,7 @@
 import subprocess
 from flask import Flask, render_template, request
 from flaskext.mysql import MySQL
+from werkzeug.wrappers import Request
 import requests
 import subprocess
 
@@ -25,7 +26,9 @@ class SetUserMiddleware:
         def __init__(self, app):
             self.app = app
         def __call__(self, environ, start_response):
-            aikido_zen.set_user({"id": "123", "name": "John Doe"})
+            req = Request(environ, shallow=True)
+            if req.headers.get("USER"):
+                aikido_zen.set_user({"id": req.headers.get("USER"), "name": "John Doe"})
             return self.app(environ, start_response)
     app.wsgi_app = AikidoFlaskMiddleware(app.wsgi_app)
     app.wsgi_app = SetUserMiddleware(app.wsgi_app)

sample-apps/quart-postgres-uvicorn/app.py

@@ -11,7 +11,9 @@ def __init__(self, app):
         self.app = app
 
     async def __call__(self, scope, receive, send):
-        aikido_zen.set_user({"id": "user123", "name": "John Doe"})
+        for header, value in scope['headers']:
+            if header.decode("utf-8").upper() == 'USER':
+                aikido_zen.set_user({"id": value.decode("utf-8"), "name": "John Doe"})
         return await self.app(scope, receive, send)
 
 app.asgi_app = AikidoQuartMiddleware(app.asgi_app)
@@ -64,6 +66,10 @@ async def create_dog():
 
     return jsonify({"message": f'Dog {dog_name} created successfully'}), 201
 
+@app.route("/test_ratelimiting_1", methods=['GET'])
+async def test_ratelimiting_1():
+    return jsonify({"message": "OK"})
+
 @app.route("/create_many", methods=['POST'])
 async def create_dog_many():
     data = await request.form

sample-apps/starlette-postgres-uvicorn/app.py

@@ -73,7 +73,9 @@ def __init__(self, app):
             self.app = app
 
         async def __call__(self, scope, receive, send):
-            aikido_zen.set_user({"id": "user123", "name": "John Doe"})
+            for header, value in scope['headers']:
+                if header.decode("utf-8").upper() == 'USER':
+                    aikido_zen.set_user({"id": value.decode("utf-8"), "name": "John Doe"})
             return await self.app(scope, receive, send)
     middleware.append(Middleware(SetUserMiddleware))
     middleware.append(Middleware(AikidoStarletteMiddleware))
@@ -87,6 +89,7 @@ async def __call__(self, scope, receive, send):
     Route("/create", create_dog, methods=["POST"]),
     Route("/sync_route", sync_route),
     Route("/just", just,  methods=["GET"]),
+    Route("/test_ratelimiting_1", just,  methods=["GET"]),
     Route("/delayed_route", delayed_route, methods=["GET"])
 ]
 if len(middleware) != 0: