Skip to content

Better e2e testing #333

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 32 commits into from
Closed
Changes from all commits
Commits
Show all changes
32 commits
Select commit Hold shift + click to select a range
7c2ca78
port code from java e2e PR
bitterpanda63 Mar 7, 2025
b774a50
Add django mysql test cases
bitterpanda63 Mar 7, 2025
f01e2f8
Add django mysql gunicorn and django postgres gunicorn
bitterpanda63 Mar 7, 2025
5a84bd5
Add quart_postgres_attack + starlette/quart e2e
bitterpanda63 Mar 7, 2025
5e7bd03
allow for 201 status codes
bitterpanda63 Mar 7, 2025
4890fa8
Fix end2end workflow
bitterpanda63 Mar 7, 2025
700d18f
Remove copied over space before unsafe_reuqest=
bitterpanda63 Mar 7, 2025
e6e535e
Also port the xml/lxml end2end tests
bitterpanda63 Mar 7, 2025
dba57be
Port flask_mysql_uwsgi end2end tests
bitterpanda63 Mar 7, 2025
35d9d1d
Remopve tests that are already ported
bitterpanda63 Mar 7, 2025
cb4275a
Add trailing slashes for django_mysql_gunicorn
bitterpanda63 Mar 7, 2025
8f24957
remove \\ from flask mysql (failed copy)
bitterpanda63 Mar 7, 2025
ce469c6
flask_mongo port test cases to new system
bitterpanda63 Mar 7, 2025
d9bc505
Convert flask_postgres to the new testing framework
bitterpanda63 Mar 7, 2025
9b73bb6
Update flask_postgres.py
bitterpanda63 Mar 8, 2025
0d53485
Create flask_mysql.py file that also tests IP Blocking and rate-limiting
bitterpanda63 Mar 10, 2025
8c3d095
Allow setting of user through user header
bitterpanda63 Mar 10, 2025
5cba5ec
Change the "dangerous" payload for flask mysql
bitterpanda63 Mar 10, 2025
1143a59
Only set user id when the user header is present
bitterpanda63 Mar 10, 2025
2a157ab
Change block message to the one used by python
bitterpanda63 Mar 10, 2025
681d9d7
More testing in quart e2e tests
bitterpanda63 Mar 10, 2025
3d6f6a4
Starlette expand testing e2e
bitterpanda63 Mar 10, 2025
c6dc209
remove standard user present in quart attack
bitterpanda63 Mar 10, 2025
b14a5c9
end2end delay make it longer
bitterpanda63 Mar 10, 2025
9992f81
Fix setting of user
bitterpanda63 Mar 10, 2025
1baeccc
Check headers for lowercase user header
bitterpanda63 Mar 10, 2025
de7b3f4
sleep 30 seconds
bitterpanda63 Mar 10, 2025
dc81d87
print the headers in quart app for SetUser
bitterpanda63 Mar 10, 2025
d5a71d4
set users dynamically in both quart and starlette apps
bitterpanda63 Mar 10, 2025
380b826
remove test_blocking from flask mysql
bitterpanda63 Mar 10, 2025
21c6aee
Remove test_blocking from quart and starlette apps as well
bitterpanda63 Mar 10, 2025
f07141e
Merge remote-tracking branch 'origin/main' into better-e2e-testing
bitterpanda63 Apr 10, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 13 additions & 14 deletions .github/workflows/end2end.yml
Original file line number Diff line number Diff line change
@@ -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 }}
4 changes: 4 additions & 0 deletions end2end/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
import json

with open('end2end/attack_events.json', 'r') as file:
events = json.load(file)
124 changes: 124 additions & 0 deletions end2end/attack_events.json
Original file line number Diff line number Diff line change
@@ -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"
}
}
17 changes: 17 additions & 0 deletions end2end/django_mysql.py
Original file line number Diff line number Diff line change
@@ -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()
12 changes: 12 additions & 0 deletions end2end/django_mysql_gunicorn.py
Original file line number Diff line number Diff line change
@@ -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()
54 changes: 0 additions & 54 deletions end2end/django_mysql_gunicorn_test.py

This file was deleted.

89 changes: 0 additions & 89 deletions end2end/django_mysql_test.py

This file was deleted.

12 changes: 12 additions & 0 deletions end2end/django_postgres_gunicorn.py
Original file line number Diff line number Diff line change
@@ -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()
54 changes: 0 additions & 54 deletions end2end/django_postgres_gunicorn_test.py

This file was deleted.

24 changes: 24 additions & 0 deletions end2end/flask_mongo.py
Original file line number Diff line number Diff line change
@@ -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()
128 changes: 0 additions & 128 deletions end2end/flask_mongo_test.py

This file was deleted.

18 changes: 18 additions & 0 deletions end2end/flask_mysql.py
Original file line number Diff line number Diff line change
@@ -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()
111 changes: 0 additions & 111 deletions end2end/flask_mysql_test.py

This file was deleted.

12 changes: 12 additions & 0 deletions end2end/flask_mysql_uwsgi.py
Original file line number Diff line number Diff line change
@@ -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()
54 changes: 0 additions & 54 deletions end2end/flask_mysql_uwsgi_test.py

This file was deleted.

17 changes: 17 additions & 0 deletions end2end/flask_postgres.py
Original file line number Diff line number Diff line change
@@ -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()
104 changes: 0 additions & 104 deletions end2end/flask_postgres_test.py

This file was deleted.

17 changes: 17 additions & 0 deletions end2end/flask_postgres_xml.py
Original file line number Diff line number Diff line change
@@ -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()
48 changes: 0 additions & 48 deletions end2end/flask_postgres_xml_lxml_test.py

This file was deleted.

53 changes: 0 additions & 53 deletions end2end/flask_postgres_xml_test.py

This file was deleted.

13 changes: 13 additions & 0 deletions end2end/quart_postgres_uvicorn.py
Original file line number Diff line number Diff line change
@@ -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()
53 changes: 0 additions & 53 deletions end2end/quart_postgres_uvicorn_test.py

This file was deleted.

27 changes: 0 additions & 27 deletions end2end/server/check_events_from_mock.py

This file was deleted.

6 changes: 6 additions & 0 deletions end2end/server/mock_aikido_core.py
Original file line number Diff line number Diff line change
@@ -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:
13 changes: 13 additions & 0 deletions end2end/starlette_postgres_uvicorn.py
Original file line number Diff line number Diff line change
@@ -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()
63 changes: 0 additions & 63 deletions end2end/starlette_postgres_uvicorn_test.py

This file was deleted.

20 changes: 20 additions & 0 deletions end2end/utils/EventHandler.py
Original file line number Diff line number Diff line change
@@ -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]
62 changes: 62 additions & 0 deletions end2end/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -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)")
4 changes: 4 additions & 0 deletions end2end/utils/assert_equals.py
Original file line number Diff line number Diff line change
@@ -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}"
30 changes: 30 additions & 0 deletions end2end/utils/request.py
Original file line number Diff line number Diff line change
@@ -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})"
36 changes: 36 additions & 0 deletions end2end/utils/test_bot_blocking.py
Original file line number Diff line number Diff line change
@@ -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.")
41 changes: 41 additions & 0 deletions end2end/utils/test_ip_blocking.py
Original file line number Diff line number Diff line change
@@ -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)
13 changes: 13 additions & 0 deletions end2end/utils/test_payloads_safe_vs_unsafe.py
Original file line number Diff line number Diff line change
@@ -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)
160 changes: 160 additions & 0 deletions end2end/utils/test_ratelimiting.py
Original file line number Diff line number Diff line change
@@ -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)
5 changes: 4 additions & 1 deletion sample-apps/flask-mysql/app.py
Original file line number Diff line number Diff line change
@@ -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)
8 changes: 7 additions & 1 deletion sample-apps/quart-postgres-uvicorn/app.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 4 additions & 1 deletion sample-apps/starlette-postgres-uvicorn/app.py
Original file line number Diff line number Diff line change
@@ -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: