Skip to content

Commit d45f628

Browse files
bitterpanda63bitterpanda63
committed
Merge branch 'main' into os-sink-remove-leading-slashes
2 parents f05d93a + 890e78e commit d45f628

File tree

9 files changed

+266
-552
lines changed

9 files changed

+266
-552
lines changed

‎aikido_zen/middleware/init_test.py‎

Lines changed: 9 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
from unittest.mock import patch, MagicMock
22

33
import pytest
4-
from aikido_zen.context import current_context, Context, get_current_context
5-
from aikido_zen.thread.thread_cache import ThreadCache, get_cache
4+
import aikido_zen.test_utils as test_utils
5+
from aikido_zen.context import current_context, get_current_context
6+
from aikido_zen.thread.thread_cache import get_cache
67
from . import should_block_request
78
from .. import set_rate_limit_group
89

@@ -22,35 +23,14 @@ def test_without_context():
2223
assert should_block_request() == {"block": False}
2324

2425

25-
def set_context(user=None, executed_middleware=False):
26-
Context(
27-
context_obj={
28-
"remote_address": "::1",
29-
"method": "POST",
30-
"url": "http://localhost:4000",
31-
"query": {
32-
"abc": "def",
33-
},
34-
"headers": {},
35-
"body": None,
36-
"cookies": {},
37-
"source": "flask",
38-
"route": "/posts/:id",
39-
"user": user,
40-
"rate_limit_group": None,
41-
"executed_middleware": executed_middleware,
42-
}
43-
).set_as_current_context()
44-
45-
4626
def test_with_context_without_cache():
47-
set_context()
27+
test_utils.generate_and_set_context()
4828
get_cache().cache = None
4929
assert should_block_request() == {"block": False}
5030

5131

5232
def test_with_context_with_cache():
53-
set_context(user={"id": "123"})
33+
test_utils.generate_and_set_context(user={"id": "123"})
5434
thread_cache = get_cache()
5535

5636
thread_cache.config.blocked_uids = ["123"]
@@ -76,7 +56,7 @@ def test_with_context_with_cache():
7656

7757

7858
def test_cache_comms_with_endpoints():
79-
set_context(user={"id": "456"})
59+
test_utils.generate_and_set_context(user={"id": "456"}, route="/posts/:id")
8060
set_rate_limit_group("my_group")
8161
thread_cache = get_cache()
8262
thread_cache.config.blocked_uids = ["123"]
@@ -145,11 +125,11 @@ def test_cache_comms_with_endpoints():
145125
"route_metadata": {
146126
"method": "POST",
147127
"route": "/posts/:id",
148-
"url": "http://localhost:4000",
128+
"url": "http://localhost:8080/",
149129
},
150130
"user": {"id": "456"},
151131
"group": "my_group",
152-
"remote_address": "::1",
132+
"remote_address": "1.1.1.1",
153133
},
154134
receive=True,
155135
timeout_in_sec=0.01,
@@ -168,7 +148,7 @@ def test_cache_comms_with_endpoints():
168148
assert thread_cache.stats.rate_limited_hits == 0
169149
assert should_block_request() == {
170150
"block": True,
171-
"ip": "::1",
151+
"ip": "1.1.1.1",
172152
"type": "ratelimited",
173153
"trigger": "my_trigger",
174154
}

‎aikido_zen/middleware/set_rate_limit_group_test.py‎

Lines changed: 8 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import pytest
2-
from aikido_zen.context import get_current_context, Context
32
from aikido_zen.thread.thread_cache import get_cache
3+
import aikido_zen.test_utils as test_utils
44
from .set_rate_limit_group import set_rate_limit_group
55

66

@@ -15,31 +15,8 @@ def run_around_tests():
1515
get_cache().reset()
1616

1717

18-
def set_context_and_lifecycle():
19-
wsgi_request = {
20-
"REQUEST_METHOD": "GET",
21-
"HTTP_HEADER_1": "header 1 value",
22-
"HTTP_HEADER_2": "Header 2 value",
23-
"RANDOM_VALUE": "Random value",
24-
"HTTP_COOKIE": "sessionId=abc123xyz456;",
25-
"wsgi.url_scheme": "http",
26-
"HTTP_HOST": "localhost:8080",
27-
"PATH_INFO": "/hello",
28-
"QUERY_STRING": "user=JohnDoe&age=30&age=35",
29-
"CONTENT_TYPE": "application/json",
30-
"REMOTE_ADDR": "198.51.100.23",
31-
}
32-
context = Context(
33-
req=wsgi_request,
34-
body=None,
35-
source="flask",
36-
)
37-
context.set_as_current_context()
38-
return context
39-
40-
4118
def test_set_rate_limit_group_valid_group_id(caplog):
42-
context1 = set_context_and_lifecycle()
19+
context1 = test_utils.generate_and_set_context()
4320
assert context1.rate_limit_group is None
4421
set_rate_limit_group("group1")
4522
assert context1.rate_limit_group == "group1"
@@ -49,15 +26,15 @@ def test_set_rate_limit_group_valid_group_id(caplog):
4926

5027

5128
def test_set_rate_limit_group_empty_group_id(caplog):
52-
context1 = set_context_and_lifecycle()
29+
context1 = test_utils.generate_and_set_context()
5330
assert context1.rate_limit_group is None
5431
set_rate_limit_group("")
5532
assert context1.rate_limit_group is None
5633
assert "Group ID cannot be empty." in caplog.text
5734

5835

5936
def test_set_rate_limit_group_none_group_id(caplog):
60-
context1 = set_context_and_lifecycle()
37+
context1 = test_utils.generate_and_set_context()
6138
assert context1.rate_limit_group is None
6239
set_rate_limit_group(None)
6340
assert context1.rate_limit_group is None
@@ -73,30 +50,30 @@ def test_set_rate_limit_group_no_context(caplog):
7350

7451

7552
def test_set_rate_limit_group_middleware_already_executed(caplog):
76-
context1 = set_context_and_lifecycle()
53+
context1 = test_utils.generate_and_set_context()
7754
context1.executed_middleware = True
7855
set_rate_limit_group("group1")
7956
assert "must be called before the Zen middleware is executed" in caplog.text
8057
assert context1.rate_limit_group is "group1"
8158

8259

8360
def test_set_rate_limit_group_non_string_group_id(caplog):
84-
context1 = set_context_and_lifecycle()
61+
context1 = test_utils.generate_and_set_context()
8562
assert context1.rate_limit_group is None
8663
set_rate_limit_group(123)
8764
assert context1.rate_limit_group == "123"
8865

8966

9067
def test_set_rate_limit_group_non_string_group_id_non_number(caplog):
91-
context1 = set_context_and_lifecycle()
68+
context1 = test_utils.generate_and_set_context()
9269
assert context1.rate_limit_group is None
9370
set_rate_limit_group({"a": "b"})
9471
assert context1.rate_limit_group is None
9572
assert "Group ID must be a string or a number" in caplog.text
9673

9774

9875
def test_set_rate_limit_group_overwrite_existing_group():
99-
context1 = set_context_and_lifecycle()
76+
context1 = test_utils.generate_and_set_context()
10077
assert context1.rate_limit_group is None
10178
set_rate_limit_group("group1")
10279
assert context1.rate_limit_group == "group1"

‎aikido_zen/sinks/tests/clickhouse_driver_test.py‎

Lines changed: 11 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,13 @@
11
import aikido_zen.sinks.clickhouse_driver
22
import pytest
33
from aikido_zen.background_process import reset_comms
4-
from aikido_zen.context import Context
54
from aikido_zen.errors import AikidoSQLInjection
6-
7-
8-
class Context1(Context):
9-
def __init__(self, body):
10-
self.cookies = {}
11-
self.headers = {}
12-
self.remote_address = "1.1.1.1"
13-
self.method = "POST"
14-
self.url = "url"
15-
self.query = {}
16-
self.body = body
17-
self.source = "express"
18-
self.route = "/"
19-
self.parsed_userinput = {}
20-
self.protection_forced_off = False
5+
import aikido_zen.test_utils as test_utils
216

227

238
@pytest.fixture(autouse=True)
24-
def set_blocking_to_true(monkeypatch):
9+
def setup(monkeypatch):
10+
reset_comms()
2511
monkeypatch.setenv("AIKIDO_BLOCK", "1")
2612

2713

@@ -35,25 +21,22 @@ def client():
3521

3622

3723
def test_client_execute_without_context(client):
38-
reset_comms()
3924
dog_name = "Steve"
4025
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
4126
client.execute(sql)
4227

4328

4429
def test_client_execute_safe(client):
45-
reset_comms()
4630
dog_name = "Steve"
4731
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
48-
Context1({"dog_name": dog_name}).set_as_current_context()
32+
test_utils.generate_and_set_context(value=dog_name)
4933
client.execute(sql)
5034

5135

5236
def test_client_execute_unsafe(client, monkeypatch):
53-
reset_comms()
5437
dog_name = "Malicious dog', 1); -- "
5538
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
56-
Context1({"dog_name": dog_name}).set_as_current_context()
39+
test_utils.generate_and_set_context(value=dog_name)
5740

5841
with pytest.raises(AikidoSQLInjection):
5942
client.execute(sql)
@@ -66,21 +49,19 @@ def test_cursor_execute_safe():
6649
from clickhouse_driver import connect
6750

6851
conn = connect("clickhouse://localhost:9000")
69-
reset_comms()
7052
dog_name = "Steve"
7153
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
72-
Context1({"dog_name": dog_name}).set_as_current_context()
54+
test_utils.generate_and_set_context(value=dog_name)
7355
conn.cursor().execute(sql)
7456

7557

7658
def test_cursor_execute_unsafe(monkeypatch):
7759
from clickhouse_driver import connect
7860

7961
conn = connect("clickhouse://localhost:9000")
80-
reset_comms()
8162
dog_name = "Malicious dog', 1); -- "
8263
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
83-
Context1({"dog_name": dog_name}).set_as_current_context()
64+
test_utils.generate_and_set_context(value=dog_name)
8465

8566
with pytest.raises(AikidoSQLInjection):
8667
conn.cursor().execute(sql)
@@ -90,18 +71,16 @@ def test_cursor_execute_unsafe(monkeypatch):
9071

9172

9273
def test_client_execute_with_progress_safe(client):
93-
reset_comms()
9474
dog_name = "Steve"
9575
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
96-
Context1({"dog_name": dog_name}).set_as_current_context()
76+
test_utils.generate_and_set_context(value=dog_name)
9777
client.execute_with_progress(sql)
9878

9979

10080
def test_client_execute_with_progress_unsafe(client, monkeypatch):
101-
reset_comms()
10281
dog_name = "Malicious dog', 1); -- "
10382
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
104-
Context1({"dog_name": dog_name}).set_as_current_context()
83+
test_utils.generate_and_set_context(value=dog_name)
10584

10685
with pytest.raises(AikidoSQLInjection):
10786
client.execute_with_progress(sql)
@@ -111,18 +90,16 @@ def test_client_execute_with_progress_unsafe(client, monkeypatch):
11190

11291

11392
def test_client_execute_iter_safe(client):
114-
reset_comms()
11593
dog_name = "Steve"
11694
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
117-
Context1({"dog_name": dog_name}).set_as_current_context()
95+
test_utils.generate_and_set_context(value=dog_name)
11896
client.execute_iter(sql)
11997

12098

12199
def test_client_execute_iter_unsafe(client, monkeypatch):
122-
reset_comms()
123100
dog_name = "Malicious dog', 1); -- "
124101
sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
125-
Context1({"dog_name": dog_name}).set_as_current_context()
102+
test_utils.generate_and_set_context(value=dog_name)
126103

127104
with pytest.raises(AikidoSQLInjection):
128105
client.execute_iter(sql)

0 commit comments

Comments
 (0)