From 97270d8a5261bbc584c860148d9b2eaf5895914e Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 09:55:47 +0200 Subject: [PATCH 01/11] Create a users file with validate_user function --- aikido_firewall/context/users.py | 46 +++++++++++++++++++ aikido_firewall/context/users_test.py | 64 +++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 aikido_firewall/context/users.py create mode 100644 aikido_firewall/context/users_test.py diff --git a/aikido_firewall/context/users.py b/aikido_firewall/context/users.py new file mode 100644 index 000000000..549ddef8d --- /dev/null +++ b/aikido_firewall/context/users.py @@ -0,0 +1,46 @@ +""" +Users file +""" + +from aikido_firewall.helpers.logging import logger + + +def set_user(user): + """ + External function for applications to set a user + """ + validated_user = validate_user(user) + if not validated_user: + return + logger.debug("Validated user : %s", validated_user) + + +def validate_user(user): + """This validates the user object""" + if not isinstance(user, dict): + logger.info( + "set_user(...) expects a dict with 'id' and 'name' properties, found %s instead.", + type(user), + ) + return + + # Validate user's id : + if not "id" in user: + logger.info("set_user(...) expects an object with 'id' property.") + return + if not isinstance(user["id"], str) and not isinstance(user["id"], int): + logger.info( + "set_user(...) expects an object with 'id' property of type string or number, found %s instead.", + type(user["id"]), + ) + return + if isinstance(user["id"], str) and len(user["id"]) is 0: + logger.info( + "set_user(...) expects an object with 'id' property non-empty string." + ) + return + valid_user = {"id": str(user["id"])} + if "name" in user and isinstance(user["name"], str) and len(user["name"]) > 0: + valid_user["name"] = str(user["name"]) + + return valid_user diff --git a/aikido_firewall/context/users_test.py b/aikido_firewall/context/users_test.py new file mode 100644 index 000000000..bc54e6fcc --- /dev/null +++ b/aikido_firewall/context/users_test.py @@ -0,0 +1,64 @@ +import pytest + +from .users import validate_user + + +def test_validate_user_valid_input(): + user = {"id": "123", "name": "Alice"} + result = validate_user(user) + assert result == {"id": "123", "name": "Alice"} + + +def test_validate_user_valid_input_with_int_id(): + user = {"id": 456, "name": "Bob"} + result = validate_user(user) + assert result == {"id": "456", "name": "Bob"} + + +def test_validate_user_missing_id(caplog): + user = {"name": "Charlie"} + result = validate_user(user) + assert result is None + assert "expects an object with 'id' property." in caplog.text + + +def test_validate_user_invalid_id_type(caplog): + user = {"id": 12.34, "name": "David"} + result = validate_user(user) + assert result is None + assert ( + "expects an object with 'id' property of type string or number" in caplog.text + ) + + +def test_validate_user_empty_string_id(caplog): + user = {"id": "", "name": "Eve"} + result = validate_user(user) + assert result is None + assert "expects an object with 'id' property non-empty string." in caplog.text + + +def test_validate_user_missing_name(caplog): + user = {"id": "789"} + result = validate_user(user) + assert result == {"id": "789"} + + +def test_validate_user_empty_name(caplog): + user = {"id": "101", "name": ""} + result = validate_user(user) + assert result == {"id": "101"} + + +def test_validate_user_invalid_user_type(caplog): + user = ["id", "name"] + result = validate_user(user) + assert result is None + assert "expects a dict with 'id' and 'name' properties" in caplog.text + + +def test_validate_user_invalid_user_type_dict_without_id(caplog): + user = {"name": "Frank"} + result = validate_user(user) + assert result is None + assert "expects an object with 'id' property." in caplog.text From c97e36ef4373d10e11c9607a28799fdd6368a26b Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 10:01:11 +0200 Subject: [PATCH 02/11] Set the user object in context and update context class to include user --- aikido_firewall/context/__init__.py | 12 ++++++++++++ aikido_firewall/context/users.py | 3 +++ 2 files changed, 15 insertions(+) diff --git a/aikido_firewall/context/__init__.py b/aikido_firewall/context/__init__.py index cc44696a0..0cb196ed9 100644 --- a/aikido_firewall/context/__init__.py +++ b/aikido_firewall/context/__init__.py @@ -7,6 +7,7 @@ from http.cookies import SimpleCookie from aikido_firewall.helpers.build_route_from_url import build_route_from_url from aikido_firewall.helpers.get_subdomains_from_url import get_subdomains_from_url +from aikido_firewall.helpers.logging import logger SUPPORTED_SOURCES = ["django", "flask", "django-gunicorn"] UINPUT_SOURCES = ["body", "cookies", "query", "headers"] @@ -14,6 +15,15 @@ local = threading.local() +def set_current_user(user): + """Sets the current user""" + if hasattr(local, "user") and local.user is not None: + logger.debug( + "Evicting a saved users, this probably means a user was set twice." + ) + local.user = user + + def get_current_context(): """Returns the current context""" try: @@ -70,6 +80,7 @@ def __init__(self, context_obj=None, req=None, source=None): self.set_django_gunicorn_attrs(req) self.route = build_route_from_url(self.url) self.subdomains = get_subdomains_from_url(self.url) + self.user = local.user if hasattr(local, "user") else None def set_django_gunicorn_attrs(self, req): """Set properties that are specific to django-gunicorn""" @@ -114,6 +125,7 @@ def __reduce__(self): "source": self.source, "route": self.route, "subdomains": self.subdomains, + "user": self.user, }, None, None, diff --git a/aikido_firewall/context/users.py b/aikido_firewall/context/users.py index 549ddef8d..160f248f5 100644 --- a/aikido_firewall/context/users.py +++ b/aikido_firewall/context/users.py @@ -3,6 +3,7 @@ """ from aikido_firewall.helpers.logging import logger +from . import set_current_user def set_user(user): @@ -14,6 +15,8 @@ def set_user(user): return logger.debug("Validated user : %s", validated_user) + set_current_user(validated_user) + def validate_user(user): """This validates the user object""" From 37328673e1a9651612b2fda0326d72fa0807dcc6 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 10:08:01 +0200 Subject: [PATCH 03/11] Also check for a remote ip address --- aikido_firewall/context/users.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/aikido_firewall/context/users.py b/aikido_firewall/context/users.py index 160f248f5..3301ed2ae 100644 --- a/aikido_firewall/context/users.py +++ b/aikido_firewall/context/users.py @@ -3,7 +3,7 @@ """ from aikido_firewall.helpers.logging import logger -from . import set_current_user +from . import set_current_user, get_current_context def set_user(user): @@ -17,6 +17,13 @@ def set_user(user): set_current_user(validated_user) + context = get_current_context() + if not context: + return + validated_user["lastIpAddress"] = context.remote_addr + + # Sent validated_user object to Agent + def validate_user(user): """This validates the user object""" From 699f194cde29fa99fa347bec0c430a409785d20d Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 10:45:36 +0200 Subject: [PATCH 04/11] Make a report to the background thread of the validated user --- .../background_process/aikido_background_process.py | 4 ++++ aikido_firewall/context/users.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/aikido_firewall/background_process/aikido_background_process.py b/aikido_firewall/background_process/aikido_background_process.py index a9a2a553b..c440158fc 100644 --- a/aikido_firewall/background_process/aikido_background_process.py +++ b/aikido_firewall/background_process/aikido_background_process.py @@ -62,6 +62,10 @@ def __init__(self, address, key): elif data[0] == "READ_PROPERTY": # meant to get config props if hasattr(self.reporter, data[1]): conn.send(self.reporter.__dict__[data[1]]) + elif data[0] == "USER": + pass + # The client reported a user + # TBI : self.reporter.add_user(data[1]) def reporting_thread(self): """Reporting thread""" diff --git a/aikido_firewall/context/users.py b/aikido_firewall/context/users.py index 3301ed2ae..c100b1610 100644 --- a/aikido_firewall/context/users.py +++ b/aikido_firewall/context/users.py @@ -4,6 +4,7 @@ from aikido_firewall.helpers.logging import logger from . import set_current_user, get_current_context +from aikido_firewall.background_process import get_comms def set_user(user): @@ -22,7 +23,8 @@ def set_user(user): return validated_user["lastIpAddress"] = context.remote_addr - # Sent validated_user object to Agent + # Send validated_user object to Agent + get_comms().send_data_to_bg_process("USER", validated_user) def validate_user(user): From d53c9d0fa58dda6a7d36d6bbe2c9d0b96bb33c87 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 21:46:14 +0200 Subject: [PATCH 05/11] Re-export set_user from main script --- aikido_firewall/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 0348d9589..47f6bf027 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -2,6 +2,9 @@ Aggregates from the different modules """ +# Re-export set_current_user : +from aikido_firewall.context.users import set_current_user as set_user + from dotenv import load_dotenv # Constants From dfc1d82164930b786594f5f55f1dff2deccb7b43 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:21:16 +0200 Subject: [PATCH 06/11] Move PKG_VERSION to config.py --- aikido_firewall/__init__.py | 3 --- aikido_firewall/background_process/reporter.py | 2 +- aikido_firewall/config.py | 3 +++ 3 files changed, 4 insertions(+), 4 deletions(-) create mode 100644 aikido_firewall/config.py diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 47f6bf027..1c241677d 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -7,9 +7,6 @@ from dotenv import load_dotenv -# Constants -PKG_VERSION = "0.0.1" - # Import logger from aikido_firewall.helpers.logging import logger diff --git a/aikido_firewall/background_process/reporter.py b/aikido_firewall/background_process/reporter.py index f11ec4bc3..9cb5433f1 100644 --- a/aikido_firewall/background_process/reporter.py +++ b/aikido_firewall/background_process/reporter.py @@ -11,7 +11,7 @@ from aikido_firewall.helpers.get_machine_ip import get_ip from aikido_firewall.helpers.get_ua_from_context import get_ua_from_context from aikido_firewall.helpers.get_current_unixtime_ms import get_unixtime_ms -from aikido_firewall import PKG_VERSION +from aikido_firewall.config import PKG_VERSION from aikido_firewall.background_process.heartbeats import send_heartbeats_every_x_secs from aikido_firewall.background_process.routes import Routes from .reporter_config import ReporterConfig diff --git a/aikido_firewall/config.py b/aikido_firewall/config.py new file mode 100644 index 000000000..a82a47a20 --- /dev/null +++ b/aikido_firewall/config.py @@ -0,0 +1,3 @@ +"""Contains package versions""" + +PKG_VERSION = "0.0.1" From c04853ba050cdbe9d7c8df2088324a87b165dbfd Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:22:49 +0200 Subject: [PATCH 07/11] Import set_user and actually set a user in flask-mysql app --- aikido_firewall/__init__.py | 2 +- sample-apps/flask-mysql/app.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 1c241677d..678bb1c73 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -3,7 +3,7 @@ """ # Re-export set_current_user : -from aikido_firewall.context.users import set_current_user as set_user +from aikido_firewall.context.users import set_user from dotenv import load_dotenv diff --git a/sample-apps/flask-mysql/app.py b/sample-apps/flask-mysql/app.py index de3cccd50..7681fe914 100644 --- a/sample-apps/flask-mysql/app.py +++ b/sample-apps/flask-mysql/app.py @@ -19,6 +19,10 @@ @app.route("/") def homepage(): + aikido_firewall.set_user({ + "id": 1, + "name": "Wout" + }) cursor = mysql.get_db().cursor() cursor.execute("SELECT * FROM db.dogs") dogs = cursor.fetchall() @@ -27,6 +31,10 @@ def homepage(): @app.route('/dogpage/') def get_dogpage(dog_id): + aikido_firewall.set_user({ + "id": 2, + "name": "Wout 2" + }) cursor = mysql.get_db().cursor() cursor.execute("SELECT * FROM db.dogs WHERE id = " + str(dog_id)) dog = cursor.fetchmany(1)[0] From 71baac3ddc88b6430c0873fbea9f84fa8b76857d Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:24:43 +0200 Subject: [PATCH 08/11] Fix bug "remote_addr" is actually "remote_address" --- aikido_firewall/context/users.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_firewall/context/users.py b/aikido_firewall/context/users.py index c100b1610..ab32ba9d5 100644 --- a/aikido_firewall/context/users.py +++ b/aikido_firewall/context/users.py @@ -21,7 +21,7 @@ def set_user(user): context = get_current_context() if not context: return - validated_user["lastIpAddress"] = context.remote_addr + validated_user["lastIpAddress"] = context.remote_address # Send validated_user object to Agent get_comms().send_data_to_bg_process("USER", validated_user) From f0dda267b87cd40ab5d3721e807bc9b29bcc740a Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:45:43 +0200 Subject: [PATCH 09/11] Add Users class and it's testfile --- aikido_firewall/background_process/users.py | 59 +++++++++++++++ .../background_process/users_test.py | 73 +++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 aikido_firewall/background_process/users.py create mode 100644 aikido_firewall/background_process/users_test.py diff --git a/aikido_firewall/background_process/users.py b/aikido_firewall/background_process/users.py new file mode 100644 index 000000000..f81748df7 --- /dev/null +++ b/aikido_firewall/background_process/users.py @@ -0,0 +1,59 @@ +""" +Export the Users class +""" + +from aikido_firewall.helpers.get_current_unixtime_ms import get_unixtime_ms + + +class Users: + """ + Class that holds users for the background process + """ + + def __init__(self, max_entries=1000): + self.max_entries = max_entries + self.users = {} + + def add_user(self, user): + """Store a user""" + user_id = user["id"] + current_time = get_unixtime_ms() + + existing = self.users.get(user_id) + if existing: + existing["name"] = user.get("name") + existing["lastIpAddress"] = user.get("lastIpAddress") + existing["lastSeenAt"] = current_time + return + + if len(self.users) >= self.max_entries: + # Remove the first added user (FIFO) + first_added_key = next(iter(self.users)) + del self.users[first_added_key] + + self.users[user_id] = { + "id": user_id, + "name": user.get("name"), + "lastIpAddress": user.get("lastIpAddress"), + "firstSeenAt": current_time, + "lastSeenAt": current_time, + } + + def as_array(self): + """ + Give all user entries back as an array + """ + return [ + { + "id": user["id"], + "name": user["name"], + "lastIpAddress": user["lastIpAddress"], + "firstSeenAt": user["firstSeenAt"], + "lastSeenAt": user["lastSeenAt"], + } + for user in self.users.values() + ] + + def clear(self): + """Clear out all users""" + self.users.clear() diff --git a/aikido_firewall/background_process/users_test.py b/aikido_firewall/background_process/users_test.py new file mode 100644 index 000000000..fd646185c --- /dev/null +++ b/aikido_firewall/background_process/users_test.py @@ -0,0 +1,73 @@ +import time +import pytest +from .users import Users # Assuming the Users class is in a file named users.py + + +@pytest.fixture +def users(): + """Fixture to create a Users instance with a max of 2 entries.""" + return Users(max_entries=2) + + +def test_users(users): + assert users.as_array() == [] + + users.add_user({"id": "1", "name": "John", "lastIpAddress": "::1"}) + user1 = users.as_array()[0] + assert user1["id"] == "1" + assert user1["name"] == "John" + assert user1["lastIpAddress"] == "::1" + assert user1["lastSeenAt"] >= 0 # lastSeenAt should be initialized + assert ( + user1["lastSeenAt"] == user1["firstSeenAt"] + ) # Initially, they should be equal + + # Simulate the passage of time + time.sleep(0.001) # Sleep for a short time to simulate ticking the clock + users.add_user({"id": "1", "name": "John Doe", "lastIpAddress": "1.2.3.4"}) + user1_updated = users.as_array()[0] + assert user1_updated["id"] == "1" + assert user1_updated["name"] == "John Doe" + assert user1_updated["lastIpAddress"] == "1.2.3.4" + assert ( + user1_updated["lastSeenAt"] >= user1_updated["firstSeenAt"] + ) # lastSeenAt should be >= firstSeenAt + assert ( + user1_updated["lastSeenAt"] == user1_updated["firstSeenAt"] + 1 + ) # lastSeenAt should be +1 + + users.add_user({"id": "2", "name": "Jane", "lastIpAddress": "1.2.3.4"}) + user2 = users.as_array()[1] + assert user2["id"] == "2" + assert user2["name"] == "Jane" + assert user2["lastIpAddress"] == "1.2.3.4" + assert ( + user2["lastSeenAt"] >= user2["firstSeenAt"] + ) # lastSeenAt should be >= firstSeenAt + assert ( + user2["lastSeenAt"] == user2["firstSeenAt"] + ) # Initially, they should be equal + + users.add_user({"id": "3", "name": "Alice", "lastIpAddress": "1.2.3.4"}) + user2_updated = users.as_array()[0] # Jane should still be the first user + user3 = users.as_array()[1] # Alice should be the second user + assert user2_updated["id"] == "2" + assert user2_updated["name"] == "Jane" + assert user2_updated["lastIpAddress"] == "1.2.3.4" + assert ( + user2_updated["lastSeenAt"] >= user2_updated["firstSeenAt"] + ) # lastSeenAt should be >= firstSeenAt + assert ( + user2_updated["lastSeenAt"] == user2_updated["firstSeenAt"] + ) # Should still be equal + + assert user3["id"] == "3" + assert user3["name"] == "Alice" + assert user3["lastIpAddress"] == "1.2.3.4" + assert ( + user3["lastSeenAt"] >= user3["firstSeenAt"] + ) # lastSeenAt should be >= firstSeenAt + assert user3["lastSeenAt"] == user3["firstSeenAt"] # Should be equal + + users.clear() + assert users.as_array() == [] From eadfe75d75b4ed74b8bf5a32cc7b92223ec0a644 Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:49:43 +0200 Subject: [PATCH 10/11] Import, create and use reporter User's object --- aikido_firewall/background_process/reporter.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/aikido_firewall/background_process/reporter.py b/aikido_firewall/background_process/reporter.py index 9cb5433f1..8c040c052 100644 --- a/aikido_firewall/background_process/reporter.py +++ b/aikido_firewall/background_process/reporter.py @@ -14,8 +14,9 @@ from aikido_firewall.config import PKG_VERSION from aikido_firewall.background_process.heartbeats import send_heartbeats_every_x_secs from aikido_firewall.background_process.routes import Routes -from .reporter_config import ReporterConfig from aikido_firewall.ratelimiting.rate_limiter import RateLimiter +from .users import Users +from .reporter_config import ReporterConfig class Reporter: @@ -33,6 +34,7 @@ def __init__(self, block, api, token, serverless, event_scheduler): self.rate_limiter = RateLimiter( max_items=5000, time_to_live_in_ms=120 * 60 * 1000 # 120 minutes ) + self.users = Users(1000) if isinstance(serverless, str) and len(serverless) == 0: raise ValueError("Serverless cannot be an empty string") @@ -88,6 +90,10 @@ def send_heartbeat(self): if not self.token: return logger.debug("Aikido Reporter : Sending out heartbeat") + users = self.users.as_array() + routes = list(self.routes) + self.users.clear() + self.routes.clear() res = self.api.report( self.token, { @@ -108,8 +114,8 @@ def send_heartbeat(self): }, }, "hostnames": [], - "routes": list(self.routes), - "users": [], + "routes": routes, + "users": users, }, self.timeout_in_sec, ) From d49a1f5bfed00b529e4a0b39f29535aa6e25f97a Mon Sep 17 00:00:00 2001 From: Wout Feys Date: Tue, 6 Aug 2024 23:49:56 +0200 Subject: [PATCH 11/11] Add the user to the reporter's object --- .../background_process/aikido_background_process.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/aikido_firewall/background_process/aikido_background_process.py b/aikido_firewall/background_process/aikido_background_process.py index c13b8a099..5f47b8b49 100644 --- a/aikido_firewall/background_process/aikido_background_process.py +++ b/aikido_firewall/background_process/aikido_background_process.py @@ -72,9 +72,7 @@ def __init__(self, address, key): ) ) elif data[0] == "USER": - pass - # The client reported a user - # TBI : self.reporter.add_user(data[1]) + self.reporter.users.add_user(data[1]) def reporting_thread(self): """Reporting thread"""