diff --git a/aikido_zen/background_process/cloud_connection_manager/__init__.py b/aikido_zen/background_process/cloud_connection_manager/__init__.py index 44c62878e..35b02afdf 100644 --- a/aikido_zen/background_process/cloud_connection_manager/__init__.py +++ b/aikido_zen/background_process/cloud_connection_manager/__init__.py @@ -1,7 +1,7 @@ """ This file simply exports the CloudConnectionManager class""" from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs -from aikido_zen.background_process.routes import Routes +from aikido_zen.storage.routes import Routes from aikido_zen.ratelimiting.rate_limiter import RateLimiter from aikido_zen.helpers.logging import logger from .update_firewall_lists import update_firewall_lists @@ -32,7 +32,7 @@ def __init__(self, block, api, token, serverless): self.block = block self.api: ReportingApiHTTP = api self.token = token # Should be instance of the Token class! - self.routes = Routes(200) + self.routes = Routes(max_size=1000) self.hostnames = Hostnames(200) self.conf = ServiceConfig( endpoints=[], @@ -73,7 +73,7 @@ def report_initial_stats(self): """ data_present = ( not self.statistics.empty() - or len(self.routes.routes) > 0 + or not self.routes.empty() or not self.ai_stats.empty() ) should_report_initial_stats = data_present and not self.conf.received_any_stats diff --git a/aikido_zen/background_process/cloud_connection_manager/init_test.py b/aikido_zen/background_process/cloud_connection_manager/init_test.py index 3c58a85af..de4ddd2d5 100644 --- a/aikido_zen/background_process/cloud_connection_manager/init_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/init_test.py @@ -25,7 +25,7 @@ def test_cloud_connection_manager_initialization(setup_cloud_connection_manager) assert manager.block is None assert isinstance(manager.api, ReportingApiHTTP) assert isinstance(manager.token, Token) - assert len(manager.routes) == 0 + assert len(manager.routes.export()) == 0 assert isinstance(manager.hostnames, Hostnames) assert isinstance(manager.conf, ServiceConfig) assert isinstance(manager.rate_limiter, RateLimiter) diff --git a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py index 9a1ff9885..1f134a54f 100644 --- a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py +++ b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py @@ -14,7 +14,7 @@ def send_heartbeat(connection_manager): logger.debug("Aikido CloudConnectionManager : Sending out heartbeat") stats = connection_manager.statistics.get_record() users = connection_manager.users.as_array() - routes = list(connection_manager.routes) + routes = list(connection_manager.routes.export().values()) outgoing_domains = connection_manager.hostnames.as_array() ai_stats = connection_manager.ai_stats.get_stats() packages = PackagesStore.export() diff --git a/aikido_zen/background_process/commands/sync_data.py b/aikido_zen/background_process/commands/sync_data.py index bdaaada7a..7a3ad58dd 100644 --- a/aikido_zen/background_process/commands/sync_data.py +++ b/aikido_zen/background_process/commands/sync_data.py @@ -13,25 +13,13 @@ def process_sync_data(connection_manager, data, conn, queue=None): BG Process -> Thread : Routes and config """ - # Sync routes - routes = connection_manager.routes - for route in data.get("current_routes", {}).values(): - route_metadata = {"method": route["method"], "route": route["path"]} - if not routes.get(route_metadata): - routes.initialize_route(route_metadata) - existing_route = routes.get(route_metadata) - - # Update hit count : - hits_delta_since_sync = int(route.get("hits_delta_since_sync", 0)) - existing_route["hits"] += hits_delta_since_sync - - # Update API Spec : - update_route_info(route["apispec"], existing_route) - # Save middleware installed : if data.get("middleware_installed", False): connection_manager.middleware_installed = True + # Sync routes + connection_manager.routes.import_from_record(data.get("routes", {})) + # Sync hostnames for hostnames_entry in data.get("hostnames", list()): connection_manager.hostnames.add( @@ -56,7 +44,7 @@ def process_sync_data(connection_manager, data, conn, queue=None): if connection_manager.conf.last_updated_at > 0: # Only report data if the config has been fetched. return { - "routes": dict(connection_manager.routes.routes), + "routes": connection_manager.routes.export(include_apispecs=False), "config": connection_manager.conf, } return {} diff --git a/aikido_zen/background_process/commands/sync_data_test.py b/aikido_zen/background_process/commands/sync_data_test.py index b12451a02..4cb661188 100644 --- a/aikido_zen/background_process/commands/sync_data_test.py +++ b/aikido_zen/background_process/commands/sync_data_test.py @@ -1,9 +1,7 @@ -from multiprocessing.forkserver import connect_to_new_process - import pytest from unittest.mock import MagicMock from .sync_data import process_sync_data -from aikido_zen.background_process.routes import Routes +from aikido_zen.storage.routes import Routes from aikido_zen.helpers.iplist import IPList from ..packages import PackagesStore from ...storage.hostnames import Hostnames @@ -35,17 +33,18 @@ def test_process_sync_data_initialization(setup_connection_manager): test_hostnames.add("bumblebee.com", 8080) data = { - "current_routes": { - "route1": { + "routes": { + "GET:/api/v1/resource": { "method": "GET", "path": "/api/v1/resource", - "hits_delta_since_sync": 5, + "hits": 5, + "hits_delta_since_sync": 2000, "apispec": {"info": "API spec for resource"}, }, - "route2": { + "POST:/api/v1/resource": { "method": "POST", "path": "/api/v1/resource", - "hits_delta_since_sync": 3, + "hits": 3, "apispec": {"info": "API spec for resource"}, }, }, @@ -75,19 +74,9 @@ def test_process_sync_data_initialization(setup_connection_manager): result = process_sync_data(connection_manager, data, None) # Check that routes were initialized correctly - assert len(connection_manager.routes) == 2 - assert ( - connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[ - "hits" - ] - == 5 - ) - assert ( - connection_manager.routes.get({"method": "POST", "route": "/api/v1/resource"})[ - "hits" - ] - == 3 - ) + assert len(connection_manager.routes.export()) == 2 + assert connection_manager.routes.get("GET", "/api/v1/resource")["hits"] == 5 + assert connection_manager.routes.get("POST", "/api/v1/resource")["hits"] == 3 # Check that the total requests were updated assert connection_manager.statistics.get_record()["requests"] == { @@ -116,18 +105,18 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana connection_manager = setup_connection_manager connection_manager.conf.last_updated_at = -1 data = { - "current_routes": { - "route1": { + "routes": { + "GET:/api/v1/resource": { "method": "GET", "path": "/api/v1/resource", - "hits_delta_since_sync": 5, + "hits": 5, "apispec": {"info": "API spec for resource"}, }, - "route2": { + "POST:/api/v1/resource": { "method": "POST", "path": "/api/v1/resource", - "hits_delta_since_sync": 3, - "apispec": {"info": "API spec for resource"}, + "hits": 3, + "apispec": {"info": "API spec for another resource"}, }, }, "stats": { @@ -148,19 +137,19 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana result = process_sync_data(connection_manager, data, None) # Check that routes were initialized correctly - assert len(connection_manager.routes) == 2 - assert ( - connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[ - "hits" - ] - == 5 - ) - assert ( - connection_manager.routes.get({"method": "POST", "route": "/api/v1/resource"})[ - "hits" - ] - == 3 - ) + assert len(connection_manager.routes.export()) == 2 + assert connection_manager.routes.get("GET", "/api/v1/resource") == { + "hits": 5, + "method": "GET", + "path": "/api/v1/resource", + "apispec": {"info": "API spec for resource"}, + } + assert connection_manager.routes.get("POST", "/api/v1/resource") == { + "hits": 3, + "method": "POST", + "path": "/api/v1/resource", + "apispec": {"info": "API spec for another resource"}, + } # Check that the total requests were updated assert connection_manager.statistics.get_record()["requests"] == { @@ -186,11 +175,11 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager hostnames_sync.add("c.com", 443) data = { - "current_routes": { + "routes": { "route1": { "method": "GET", "path": "/api/v1/resource", - "hits_delta_since_sync": 5, + "hits": 5, "apispec": {"info": "API spec for resource"}, } }, @@ -214,11 +203,11 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager # Second call to update the existing route data_update = { - "current_routes": { + "routes": { "route1": { "method": "GET", "path": "/api/v1/resource", - "hits_delta_since_sync": 10, + "hits": 10, "apispec": {"info": "Updated API spec for resource"}, } }, @@ -239,12 +228,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager result = process_sync_data(connection_manager, data_update, None) # Check that the hit count was updated correctly - assert ( - connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[ - "hits" - ] - == 15 - ) + assert connection_manager.routes.get("GET", "/api/v1/resource")["hits"] == 15 # Check that the total requests were updated assert connection_manager.statistics.get_record()["requests"] == { @@ -267,11 +251,24 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager def test_process_sync_data_no_routes(setup_connection_manager): """Test behavior when no routes are provided.""" connection_manager = setup_connection_manager - data = {"current_routes": {}, "reqs": 0} # No requests to add + data = {"routes": {}, "reqs": 0} # No requests to add + + result = process_sync_data(connection_manager, data, None) + + # Check that no routes were initialized + assert len(connection_manager.routes.export()) == 0 + + # Check that the total requests remain unchanged + + +def test_process_sync_data_routes_undefined(setup_connection_manager): + """Test behavior when no routes are provided.""" + connection_manager = setup_connection_manager + data = {"sdfsdsdfsdfds": {}, "reqs": 0} # No requests to add result = process_sync_data(connection_manager, data, None) # Check that no routes were initialized - assert len(connection_manager.routes) == 0 + assert len(connection_manager.routes.export()) == 0 # Check that the total requests remain unchanged diff --git a/aikido_zen/background_process/routes/__init__.py b/aikido_zen/background_process/routes/__init__.py deleted file mode 100644 index 1a05ce5e0..000000000 --- a/aikido_zen/background_process/routes/__init__.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -Exports class Routes -""" - -from aikido_zen.helpers.logging import logger -from aikido_zen.api_discovery.update_route_info import update_route_info -from aikido_zen.api_discovery.get_api_info import get_api_info -from .route_to_key import route_to_key - - -class Routes: - """ - Stores all routes - """ - - def __init__(self, max_size=1000): - self.max_size = max_size - self.routes = {} - - def initialize_route(self, route_metadata): - """ - Initializes a route for the first time. `hits_delta_since_sync` counts delta between syncs. - """ - self.manage_routes_size() - key = route_to_key(route_metadata) - self.routes[key] = { - "method": route_metadata.get("method"), - "path": route_metadata.get("route"), - "hits": 0, - "hits_delta_since_sync": 0, - "apispec": {}, - } - - def increment_route(self, route_metadata): - """ - Adds a hit to the route (if it exists) specified in route_metadata. - route_metadata object includes route, url and method - """ - key = route_to_key(route_metadata) - if not self.routes.get(key): - self.initialize_route(route_metadata) - # Add hits to - route = self.routes.get(key) - route["hits"] += 1 - route["hits_delta_since_sync"] += 1 - - def update_route_with_apispec(self, route_metadata, apispec): - """ - Updates apispec of a given route (or creates it). - route_metadata object includes route, url and method - """ - key = route_to_key(route_metadata) - if not self.routes.get(key): - return - update_route_info(apispec, self.routes[key]) - - def get(self, route_metadata): - """Gets you the route entry if it exists using route metadata""" - key = route_to_key(route_metadata) - return self.routes.get(key) - - def get_routes_with_hits(self): - """Gets you all routes with a positive hits delta""" - result = dict() - for key, route in self.routes.items(): - if route["hits_delta_since_sync"] <= 0: - continue # do not add routes without a hit delta - result[key] = route - return result - - def clear(self): - """Deletes all routes""" - self.routes = {} - - def manage_routes_size(self): - """ - Evicts LRU routes if the size is too large - """ - if len(self.routes) >= self.max_size: - least_used = [None, float("inf")] - for key, route in self.routes.items(): - if route.get("hits") < least_used[1]: - least_used = [key, route.get("hits")] - if least_used[0]: - del self.routes[least_used[0]] - - def __iter__(self): - return iter(self.routes.values()) - - def __len__(self): - return len(self.routes) diff --git a/aikido_zen/background_process/routes/route_to_key.py b/aikido_zen/background_process/routes/route_to_key.py deleted file mode 100644 index 8a156cf9b..000000000 --- a/aikido_zen/background_process/routes/route_to_key.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Exports route_to_key function""" - - -def route_to_key(route_metadata): - """Creates a key from the method and path""" - method = route_metadata.get("method") - route = route_metadata.get("route") - return f"{method}:{route}" diff --git a/aikido_zen/background_process/routes/route_to_key_test.py b/aikido_zen/background_process/routes/route_to_key_test.py deleted file mode 100644 index 39aefedf5..000000000 --- a/aikido_zen/background_process/routes/route_to_key_test.py +++ /dev/null @@ -1,52 +0,0 @@ -import pytest -from .route_to_key import route_to_key - - -def gen_route_metadata(method="GET", route="/test", url="http://localhost:5000"): - return {"method": method, "route": route, "url": url} - - -def test_route_to_key_get(): - assert ( - route_to_key(gen_route_metadata(method="GET", route="/api/resource")) - == "GET:/api/resource" - ) - - -def test_route_to_key_post(): - assert ( - route_to_key(gen_route_metadata(method="POST", route="/api/resource")) - == "POST:/api/resource" - ) - - -def test_route_to_key_put(): - assert ( - route_to_key(gen_route_metadata(method="PUT", route="/api/resource")) - == "PUT:/api/resource" - ) - - -def test_route_to_key_delete(): - assert ( - route_to_key(gen_route_metadata(method="DELETE", route="/api/resource")) - == "DELETE:/api/resource" - ) - - -def test_route_to_key_with_query_params(): - assert ( - route_to_key(gen_route_metadata(method="GET", route="/api/resource?query=1")) - == "GET:/api/resource?query=1" - ) - - -def test_route_to_key_with_trailing_slash(): - assert ( - route_to_key(gen_route_metadata(method="GET", route="/api/resource/")) - == "GET:/api/resource/" - ) - - -def test_route_to_key_empty_path(): - assert route_to_key(gen_route_metadata(method="GET", route="")) == "GET:" diff --git a/aikido_zen/sources/functions/request_handler.py b/aikido_zen/sources/functions/request_handler.py index 3250ca99d..372bde95e 100644 --- a/aikido_zen/sources/functions/request_handler.py +++ b/aikido_zen/sources/functions/request_handler.py @@ -73,7 +73,6 @@ def post_response(status_code): context = ctx.get_current_context() if not context: return - route_metadata = context.get_route_metadata() is_curr_route_useful = is_useful_route( status_code, @@ -85,8 +84,8 @@ def post_response(status_code): cache = get_cache() if cache: - cache.routes.increment_route(route_metadata) + cache.routes.increment_route(context.method, context.route) # api spec generation - route = cache.routes.get(route_metadata) + route = cache.routes.get(context.method, context.route) update_route_info_from_context(context, route) diff --git a/aikido_zen/sources/functions/request_handler_test.py b/aikido_zen/sources/functions/request_handler_test.py index e13f6bb74..e99486eb6 100644 --- a/aikido_zen/sources/functions/request_handler_test.py +++ b/aikido_zen/sources/functions/request_handler_test.py @@ -42,7 +42,7 @@ def test_post_response_useful_route(mock_context): # Check that the route was initialized and updated route = cache.routes.routes.get("GET:/test/route") - assert route["hits"] == route["hits_delta_since_sync"] == 25 + assert route["hits"] == 25 assert route["method"] == "GET" assert route["path"] == "/test/route" diff --git a/aikido_zen/storage/routes/__init__.py b/aikido_zen/storage/routes/__init__.py new file mode 100644 index 000000000..225d80548 --- /dev/null +++ b/aikido_zen/storage/routes/__init__.py @@ -0,0 +1,77 @@ +from aikido_zen.api_discovery.update_route_info import update_route_info +from aikido_zen.helpers.logging import logger + + +class Routes: + """ + Stores: method & path for a given route, hit counts, the generated api spec. + """ + + def __init__(self, max_size=1000): + self.max_size = max_size + self.routes = {} + + def get(self, method, route): + key = route_to_key(method, route) + return self.routes.get(key) + + def ensure_route(self, method, route): + if self.get(method, route): + return # A route already exists + key = route_to_key(method, route) + self.routes[key] = { + "method": method, + "path": route, + "hits": 0, + "apispec": {}, + } + logger.error(self.routes) + self.manage_routes_size() + + def increment_route(self, method, route): + self.ensure_route(method, route) + self.get(method, route)["hits"] += 1 + + def update_route_with_apispec(self, method, route, apispec): + self.ensure_route(method, route) + update_route_info(apispec, self.get(method, route)) + + def export(self, include_apispecs=True): + result = dict(self.routes) + if not include_apispecs: + for route in result.values(): + route["apispec"] = {} + return result + + def import_from_record(self, new_routes): + if not isinstance(new_routes, dict): + return + for route in new_routes.values(): + self.ensure_route(method=route["method"], route=route["path"]) + existing_route = self.get(method=route["method"], route=route["path"]) + # merge + existing_route["hits"] += route.get("hits", 0) + update_route_info(route.get("apispec", {}), existing_route) + self.manage_routes_size() + + def clear(self): + self.routes = {} + + def empty(self): + return len(self.routes) == 0 + + def manage_routes_size(self): + """ + Evicts LRU routes if the size is too large + """ + if len(self.routes) > self.max_size: + least_used = [None, float("inf")] + for key, route in self.routes.items(): + if route.get("hits") < least_used[1]: + least_used = [key, route.get("hits")] + if least_used[0]: + del self.routes[least_used[0]] + + +def route_to_key(method, route): + return f"{method}:{route}" diff --git a/aikido_zen/background_process/routes/init_test.py b/aikido_zen/storage/routes/init_test.py similarity index 51% rename from aikido_zen/background_process/routes/init_test.py rename to aikido_zen/storage/routes/init_test.py index 44471f623..9e31299d6 100644 --- a/aikido_zen/background_process/routes/init_test.py +++ b/aikido_zen/storage/routes/init_test.py @@ -1,4 +1,6 @@ -from aikido_zen.background_process.routes import Routes +import pytest + +from aikido_zen.storage.routes import Routes from aikido_zen.api_discovery.get_api_info import get_api_info @@ -24,152 +26,246 @@ def __init__( self.cookies = cookies -def gen_route_metadata(method="GET", route="/test", url="http://localhost:5000"): - return {"method": method, "route": route, "url": url} +def test_init_default_max_size(): + routes = Routes() + assert routes.max_size == 1000 + assert routes.routes == {} -def test_initialization(): - routes = Routes(max_size=3) - assert routes.max_size == 3 +def test_init_custom_max_size(): + routes = Routes(max_size=500) + assert routes.max_size == 500 assert routes.routes == {} +def test_get_existing_route(): + routes = Routes() + routes.routes = {"GET:/test": {"method": "GET", "path": "/test"}} + + result = routes.get("GET", "/test") + assert result == {"method": "GET", "path": "/test"} + + +def test_get_nonexistent_route(): + routes = Routes() + + assert routes.get("GET", "/nonexistent") is None + + +def test_ensure_route_creates_new_route(): + routes = Routes() + + routes.ensure_route("POST", "/api/test") + + assert routes.get("POST", "/api/test") == { + "method": "POST", + "path": "/api/test", + "hits": 0, + "apispec": {}, + } + + +def test_ensure_route_existing_route_returns_early(): + routes = Routes() + routes.routes = {"GET:/test": {"method": "GET", "path": "/test", "hits": 5}} + + routes.ensure_route("GET", "/test") + + assert routes.get("GET", "/test")["hits"] == 5 + assert routes.export() == { + "GET:/test": {"hits": 5, "method": "GET", "path": "/test"} + } + + +def test_increment_route_new_route(): + routes = Routes() + + routes.increment_route("GET", "/test") + + assert routes.get("GET", "/test")["hits"] == 1 + + +def test_increment_route_existing_route(): + routes = Routes() + + routes.increment_route("GET", "/test") + routes.increment_route("GET", "/test") + routes.increment_route("GET", "/test") + routes.increment_route("GET", "/test") + routes.increment_route("GET", "/test") + routes.increment_route("GET", "/test") + + assert routes.get("GET", "/test")["hits"] == 6 + assert routes.export()["GET:/test"] == { + "apispec": {}, + "hits": 6, + "method": "GET", + "path": "/test", + } + + +def test_increment_route_multiple_times(): + routes = Routes() + + routes.increment_route("POST", "/api/data") + routes.increment_route("POST", "/api/data") + routes.increment_route("POST", "/api/data") + routes.increment_route("GET", "/api/data") + + assert routes.get("POST", "/api/data")["hits"] == 3 + assert routes.get("GET", "/api/data")["hits"] == 1 + + def test_get_route(): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(method="GET", route="/api/resource1")) - routes.increment_route(gen_route_metadata(method="GET", route="/api/resource1")) + # Use ensure_route instead of initialize_route + routes.ensure_route(method="GET", route="/api/resource1") + routes.increment_route(method="GET", route="/api/resource1") - routes.initialize_route(gen_route_metadata(method="POST", route="/api/resource2")) - routes.increment_route(gen_route_metadata(method="POST", route="/api/resource2")) + routes.ensure_route(method="POST", route="/api/resource2") + routes.increment_route(method="POST", route="/api/resource2") - routes.initialize_route(gen_route_metadata(method="PUT", route="/api/resource3")) - routes.increment_route(gen_route_metadata(method="PUT", route="/api/resource3")) + routes.ensure_route(method="PUT", route="/api/resource3") + routes.increment_route(method="PUT", route="/api/resource3") - assert routes.get(gen_route_metadata(method="GET", route="/api/resource1")) == { + # Update assertions to match the current structure + assert routes.get(method="GET", route="/api/resource1") == { "method": "GET", "path": "/api/resource1", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, } - assert routes.get(gen_route_metadata(method="POST", route="/api/resource2")) == { + assert routes.get(method="POST", route="/api/resource2") == { "method": "POST", "path": "/api/resource2", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, } - assert routes.get(gen_route_metadata(method="PUT", route="/api/resource3")) == { + assert routes.get(method="PUT", route="/api/resource3") == { "method": "PUT", "path": "/api/resource3", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, } - assert routes.get(gen_route_metadata(method="GE", route="/api/resource1")) == None - assert routes.get(gen_route_metadata(method="GET", route="/api/resource")) == None + assert routes.get(method="GE", route="/api/resource1") is None + assert routes.get(method="GET", route="/api/resource") is None -def test_initialize_route(): +def test_ensure_route(): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(route="/api/resource")) - assert len(routes.routes) == 1 - assert routes.routes["GET:/api/resource"]["hits"] == 0 + routes.ensure_route(method="GET", route="/api/resource") + assert len(routes.export()) == 1 + assert routes.get(method="GET", route="/api/resource")["hits"] == 0 - # Make sure does not do anything if route does not exist - routes.initialize_route(gen_route_metadata(route="/api/resource")) - assert len(routes.routes) == 1 - assert routes.routes["GET:/api/resource"]["hits"] == 0 + # Ensure it does not add the route again if it already exists + routes.ensure_route(method="GET", route="/api/resource") + assert len(routes.export()) == 1 + assert routes.get(method="GET", route="/api/resource")["hits"] == 0 def test_increment_route(): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(route="/api/resource")) - routes.increment_route(gen_route_metadata(route="/api/resource")) - assert len(routes.routes) == 1 - assert routes.routes["GET:/api/resource"]["hits"] == 1 + routes.ensure_route(method="GET", route="/api/resource") + routes.increment_route(method="GET", route="/api/resource") + assert len(routes.export()) == 1 + assert routes.get(method="GET", route="/api/resource")["hits"] == 1 def test_increment_route_twice(): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(route="/api/resource")) - routes.increment_route(gen_route_metadata(route="/api/resource")) - routes.increment_route(gen_route_metadata(route="/api/resource")) - assert len(routes.routes) == 1 - assert routes.routes["GET:/api/resource"]["hits"] == 2 + routes.ensure_route(method="GET", route="/api/resource") + routes.increment_route(method="GET", route="/api/resource") + routes.increment_route(method="GET", route="/api/resource") + + assert len(routes.export()) == 1 + assert routes.get(method="GET", route="/api/resource")["hits"] == 2 def test_clear_routes(): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(route="/api/resource")) + routes.ensure_route(method="GET", route="/api/resource") routes.clear() - assert len(routes.routes) == 0 + assert routes.empty() def test_manage_routes_size_eviction(): routes = Routes(max_size=2) - routes.initialize_route(gen_route_metadata(route="/api/resource1")) - routes.initialize_route(gen_route_metadata(route="/api/resource2")) - routes.initialize_route( - gen_route_metadata(route="/api/resource3") - ) # This should evict the least used route + routes.ensure_route(method="GET", route="/api/resource1") + routes.ensure_route(method="GET", route="/api/resource2") + routes.ensure_route(method="GET", route="/api/resource3") - assert len(routes.routes) == 2 - assert ( - "GET:/api/resource1" not in routes.routes - ) # Assuming resource1 is the least used + assert len(routes.export()) == 2 + assert routes.get(method="GET", route="/api/resource1") is None -def test_iterable(): +def test_api_discovery_for_new_routes(monkeypatch): routes = Routes(max_size=3) - routes.initialize_route(gen_route_metadata(method="GET", route="/api/resource1")) - routes.initialize_route(gen_route_metadata(method="POST", route="/api/resource2")) - routes.initialize_route(gen_route_metadata(method="PUT", route="/api/resource3")) + context1 = Context( + "GET", + "/api/resource1", + body={ + "user": { + "name": "John Doe", + "email": "john.doe@example.com", + "phone": 12345678, + }, + }, + ) + + # Ensure route is initialized + routes.ensure_route(method="GET", route="/api/resource1") - routes_list = list(routes) - assert len(routes_list) == 3 - assert { + # Check if the route is correctly added + routes_export = routes.export() + assert len(routes_export) == 1 + assert routes_export["GET:/api/resource1"] == { "method": "GET", "path": "/api/resource1", "hits": 0, - "hits_delta_since_sync": 0, - "apispec": {}, - } in routes_list - assert { - "method": "POST", - "path": "/api/resource2", - "hits": 0, - "hits_delta_since_sync": 0, "apispec": {}, - } in routes_list - assert { - "method": "PUT", - "path": "/api/resource3", - "hits": 0, - "hits_delta_since_sync": 0, - "apispec": {}, - } in routes_list - + } -def test_len(): - routes = Routes(max_size=3) - assert len(routes) == 0 + routes.update_route_with_apispec( + method="GET", route="/api/resource1", apispec=get_api_info(context1) + ) - routes.initialize_route(gen_route_metadata(route="/api/resource")) - assert len(routes) == 1 + # Verify the route is updated with apispec + routes_export = routes.export() + assert len(routes_export) == 1 + assert routes_export["GET:/api/resource1"]["apispec"] == { + "auth": None, + "body": { + "schema": { + "properties": { + "user": { + "properties": { + "email": {"type": "string"}, + "name": {"type": "string"}, + "phone": {"type": "number"}, + }, + "type": "object", + } + }, + "type": "object", + }, + "type": "form-urlencoded", + }, + "query": None, + } - routes.initialize_route(gen_route_metadata(method="POST", route="/api/resource2")) - assert len(routes) == 2 + routes_export = routes.export(include_apispecs=False) + assert len(routes_export) == 1 + assert routes_export["GET:/api/resource1"]["apispec"] == {} -def test_api_discovery_for_new_routes(monkeypatch): +def test_api_discovery_for_new_routes_without_ensure_route(monkeypatch): routes = Routes(max_size=3) - monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") + context1 = Context( "GET", "/api/resource1", @@ -181,51 +277,48 @@ def test_api_discovery_for_new_routes(monkeypatch): }, }, ) - routes.initialize_route(gen_route_metadata(route="/api/resource1")) - - routes_list = list(routes) - assert len(routes_list) == 1 - assert { - "method": "GET", - "path": "/api/resource1", - "hits": 0, - "hits_delta_since_sync": 0, - "apispec": {}, - } in routes_list - - apispec = get_api_info(context1) routes.update_route_with_apispec( - gen_route_metadata(route="/api/resource1"), apispec + method="GET", route="/api/resource1", apispec=get_api_info(context1) ) - routes_list = list(routes) - assert len(routes_list) == 1 - assert { - "method": "GET", - "path": "/api/resource1", - "hits": 0, - "hits_delta_since_sync": 0, - "apispec": apispec, - } in routes_list + # Verify the route is updated with apispec + routes_export = routes.export() + assert len(routes_export) == 1 + assert routes_export["GET:/api/resource1"]["apispec"] == { + "auth": None, + "body": { + "schema": { + "properties": { + "user": { + "properties": { + "email": {"type": "string"}, + "name": {"type": "string"}, + "phone": {"type": "number"}, + }, + "type": "object", + } + }, + "type": "object", + }, + "type": "form-urlencoded", + }, + "query": None, + } def test_api_discovery_existing_route_empty(monkeypatch): routes = Routes(max_size=3) - monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") - route_metadata = gen_route_metadata(route="/api/resource1") context1 = Context("GET", "/api/resource1") - routes.initialize_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context1)) - routes_list = list(routes) - assert { - "method": "GET", - "path": "/api/resource1", - "hits": 0, - "hits_delta_since_sync": 0, - "apispec": {}, - } in routes_list + # Update route with empty apispec + routes.update_route_with_apispec(method="GET", route="/api/resource1", apispec={}) + + # Verify the route is correctly added with empty apispec + routes_export = routes.export() + assert "GET:/api/resource1" in routes_export + assert routes_export["GET:/api/resource1"]["apispec"] == {} + # Mock context with body context2 = Context( "GET", "/api/resource1", @@ -237,47 +330,45 @@ def test_api_discovery_existing_route_empty(monkeypatch): }, }, ) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context2)) - routes_list = list(routes) - assert len(routes_list) == 1 - assert { - "method": "GET", - "path": "/api/resource1", - "hits": 1, - "hits_delta_since_sync": 1, + # Increment route hit count and update apispec + routes.increment_route(method="GET", route="/api/resource1") + routes.update_route_with_apispec( + method="GET", route="/api/resource1", apispec=get_api_info(context2) + ) + + # Verify the route is updated with new apispec and hit count + routes_export = routes.export() + assert len(routes_export) == 1 + assert routes_export["GET:/api/resource1"] == { "apispec": { + "auth": None, "body": { "schema": { "properties": { "user": { "properties": { - "email": { - "type": "string", - }, - "name": { - "type": "string", - }, + "email": {"type": "string"}, + "name": {"type": "string"}, "phone": {"type": "number"}, }, "type": "object", - }, + } }, "type": "object", }, "type": "form-urlencoded", }, - "auth": None, "query": None, }, - } in routes_list + "hits": 1, + "method": "GET", + "path": "/api/resource1", + } def test_api_discovery_merge_routes(monkeypatch): routes = Routes(max_size=3) - monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") - route_metadata = gen_route_metadata(route="/api/resource1") context1 = Context( "GET", @@ -296,19 +387,17 @@ def test_api_discovery_merge_routes(monkeypatch): }, content_type="application/json", ) - routes.initialize_route(route_metadata) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context1)) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context2)) + routes.increment_route("GET", "/api/resource1") + routes.update_route_with_apispec("GET", "/api/resource1", get_api_info(context1)) + routes.increment_route("GET", "/api/resource1") + routes.update_route_with_apispec("GET", "/api/resource1", get_api_info(context2)) - routes_list = list(routes) + routes_list = routes.export() assert len(routes_list) == 1 - assert { + assert routes_list["GET:/api/resource1"] == { "method": "GET", "path": "/api/resource1", "hits": 2, - "hits_delta_since_sync": 2, "apispec": { "body": { "schema": { @@ -334,41 +423,35 @@ def test_api_discovery_merge_routes(monkeypatch): "auth": None, "query": None, }, - } in routes_list + } def test_merge_body_schema(monkeypatch): - monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") - routes = Routes(200) - assert list(routes) == [] + assert len(routes.export()) == 0 - routes.initialize_route(gen_route_metadata(method="POST", route="/body")) - routes.increment_route(gen_route_metadata(method="POST", route="/body")) - assert list(routes) == [ - { + routes.increment_route("POST", "/body") + assert routes.export() == { + "POST:/body": { "method": "POST", "path": "/body", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, }, - ] + } context1 = Context( "POST", "/body", {"test": "abc", "arr": [1, 2, 3], "sub": {"y": 123}}, ) - route_metadata1 = gen_route_metadata(method="POST", route="/body") - routes.increment_route(route_metadata1) - routes.update_route_with_apispec(route_metadata1, get_api_info(context1)) - assert list(routes) == [ - { + + routes.update_route_with_apispec("POST", "/body", get_api_info(context1)) + assert routes.export() == { + "POST:/body": { "method": "POST", "path": "/body", - "hits": 2, - "hits_delta_since_sync": 2, + "hits": 1, "apispec": { "body": { "type": "form-urlencoded", @@ -393,22 +476,21 @@ def test_merge_body_schema(monkeypatch): "query": None, }, }, - ] + } context2 = Context( "POST", "/body", {"test": "abc", "arr": [1, 2, 3], "test2": 1, "sub": {"x": 123}}, ) - routes.increment_route(route_metadata1) - routes.update_route_with_apispec(route_metadata1, get_api_info(context2)) - routes.increment_route(route_metadata1) - assert list(routes) == [ - { + routes.increment_route("POST", "/body") + routes.update_route_with_apispec("POST", "/body", get_api_info(context2)) + routes.increment_route("POST", "/body") + assert routes.export() == { + "POST:/body": { "method": "POST", "path": "/body", - "hits": 4, - "hits_delta_since_sync": 4, + "hits": 3, "apispec": { "body": { "type": "form-urlencoded", @@ -435,7 +517,7 @@ def test_merge_body_schema(monkeypatch): "query": None, }, }, - ] + } def test_add_query_schema(monkeypatch): @@ -443,17 +525,14 @@ def test_add_query_schema(monkeypatch): routes = Routes(200) context = Context("GET", "/query", None, query={"test": "abc", "t": "123"}) - route_metadata = gen_route_metadata(route="/query") - routes.initialize_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context)) + routes.update_route_with_apispec("GET", "/query", get_api_info(context)) - assert list(routes) == [ - { + assert routes.export() == { + "GET:/query": { "method": "GET", "path": "/query", "hits": 0, - "hits_delta_since_sync": 0, "apispec": { "body": None, "auth": None, @@ -466,43 +545,39 @@ def test_add_query_schema(monkeypatch): }, }, }, - ] + } def test_merge_query_schema(monkeypatch): monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") routes = Routes(200) - assert list(routes) == [] - route_metadata = gen_route_metadata(route="/query") - routes.initialize_route(route_metadata) - routes.increment_route(route_metadata) - assert list(routes) == [ + assert len(routes.export()) == 0 + routes.increment_route("GET", "/query") + assert list(routes.export().values()) == [ { "method": "GET", "path": "/query", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, }, ] context1 = Context("GET", "/query", None, query={"test": "abc"}) context2 = Context("GET", "/query", None, query={"x": "123", "test": "abc"}) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context1)) + routes.increment_route("GET", "/query") + routes.update_route_with_apispec("GET", "/query", get_api_info(context1)) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context2)) + routes.increment_route("GET", "/query") + routes.update_route_with_apispec("GET", "/query", get_api_info(context2)) - routes.increment_route(route_metadata) + routes.increment_route("GET", "/query") - assert list(routes) == [ - { + assert routes.export() == { + "GET:/query": { "method": "GET", "path": "/query", "hits": 4, - "hits_delta_since_sync": 4, "apispec": { "query": { "type": "object", @@ -515,92 +590,80 @@ def test_merge_query_schema(monkeypatch): "body": None, }, }, - ] + } def test_add_auth_schema(monkeypatch): monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") routes = Routes(200) - route_metadata1 = gen_route_metadata(route="/auth") context1 = Context("GET", "/auth", headers={"AUTHORIZATION": "Bearer token"}) - route_metadata2 = gen_route_metadata(route="/auth2") context2 = Context("GET", "/auth2", cookies={"session": "test"}) - route_metadata3 = gen_route_metadata(route="/auth3") context3 = Context("GET", "/auth3", headers={"X_API_KEY": "token"}) - routes.initialize_route(route_metadata1) - routes.update_route_with_apispec(route_metadata1, get_api_info(context1)) - routes.initialize_route(route_metadata2) - routes.update_route_with_apispec(route_metadata2, get_api_info(context2)) - routes.initialize_route(route_metadata3) - routes.update_route_with_apispec(route_metadata3, get_api_info(context3)) + routes.update_route_with_apispec("GET", "/auth", get_api_info(context1)) + routes.update_route_with_apispec("GET", "/auth2", get_api_info(context2)) + routes.update_route_with_apispec("GET", "/auth3", get_api_info(context3)) - assert list(routes) == [ - { + assert routes.export() == { + "GET:/auth": { "method": "GET", "path": "/auth", "hits": 0, - "hits_delta_since_sync": 0, "apispec": { "body": None, "query": None, "auth": [{"type": "http", "scheme": "bearer"}], }, }, - { + "GET:/auth2": { "method": "GET", "path": "/auth2", "hits": 0, - "hits_delta_since_sync": 0, "apispec": { "body": None, "query": None, "auth": [{"type": "apiKey", "in": "cookie", "name": "session"}], }, }, - { + "GET:/auth3": { "method": "GET", "path": "/auth3", "hits": 0, - "hits_delta_since_sync": 0, "apispec": { "body": None, "query": None, "auth": [{"type": "apiKey", "in": "header", "name": "x-api-key"}], }, }, - ] + } def test_merge_auth_schema(monkeypatch): monkeypatch.setenv("AIKIDO_FEATURE_COLLECT_API_SCHEMA", "1") routes = Routes(200) - assert list(routes) == [] - route_metadata = gen_route_metadata(route="/auth") + assert routes.empty() - routes.initialize_route(route_metadata) - routes.increment_route(route_metadata) + routes.increment_route("GET", "/auth") context1 = Context("GET", "/auth", headers={"AUTHORIZATION": "Bearer token"}) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context1)) + routes.increment_route("GET", "/auth") + routes.update_route_with_apispec("GET", "/auth", get_api_info(context1)) context2 = Context("GET", "/auth", headers={"AUTHORIZATION": "Basic token"}) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context2)) + routes.increment_route("GET", "/auth") + routes.update_route_with_apispec("GET", "/auth", get_api_info(context2)) context3 = Context("GET", "/auth", headers={"X_API_KEY": "token"}) - routes.increment_route(route_metadata) - routes.update_route_with_apispec(route_metadata, get_api_info(context3)) - routes.increment_route(route_metadata) + routes.increment_route("GET", "/auth") + routes.update_route_with_apispec("GET", "/auth", get_api_info(context3)) + routes.increment_route("GET", "/auth") - assert list(routes)[0] == { + assert routes.export()["GET:/auth"] == { "method": "GET", "path": "/auth", "hits": 5, - "hits_delta_since_sync": 5, "apispec": { "auth": [ {"type": "http", "scheme": "bearer"}, @@ -611,4 +674,4 @@ def test_merge_auth_schema(monkeypatch): "query": None, }, } - assert len(list(routes)) == 1 + assert len(routes.export()) == 1 diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index d7fe8bb03..3b64b329a 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -2,7 +2,7 @@ import aikido_zen.background_process.comms as comms from aikido_zen.background_process.packages import PackagesStore -from aikido_zen.background_process.routes import Routes +from aikido_zen.storage.routes import Routes from aikido_zen.background_process.service_config import ServiceConfig from aikido_zen.storage.ai_statistics import AIStatistics from aikido_zen.storage.hostnames import Hostnames @@ -17,6 +17,10 @@ class ThreadCache: """ def __init__(self): + # We store both routes and routes_from_background, the latter contains all hits (useful for apispec) + self.routes = Routes(max_size=1000) + self.routes_from_background = Routes(max_size=1000) + self.hostnames = Hostnames(200) self.users = Users(1000) self.stats = Statistics() @@ -36,7 +40,6 @@ def get_endpoints(self): def reset(self): """Empties out all values of the cache""" - self.routes = Routes(max_size=1000) self.config = ServiceConfig( endpoints=[], blocked_uids=set(), @@ -45,6 +48,8 @@ def reset(self): received_any_stats=False, ) self.middleware_installed = False + self.routes.clear() + self.routes_from_background.clear() self.hostnames.clear() self.users.clear() self.stats.clear() @@ -59,7 +64,7 @@ def renew(self): res = comms.get_comms().send_data_to_bg_process( action="SYNC_DATA", obj={ - "current_routes": self.routes.get_routes_with_hits(), + "routes": self.routes.export(), "middleware_installed": self.middleware_installed, "hostnames": self.hostnames.as_array(), "users": self.users.as_array(), @@ -78,10 +83,7 @@ def renew(self): self.config = res["data"]["config"] # update routes - if isinstance(res["data"].get("routes"), dict): - self.routes.routes = res["data"]["routes"] - for route in self.routes.routes.values(): - route["hits_delta_since_sync"] = 0 + self.routes_from_background.import_from_record(res["data"].get("routes", {})) # For these 2 functions and the data they process, we rely on Python's GIL diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index c64e6fad2..8474eeea4 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -1,6 +1,6 @@ import pytest from unittest.mock import patch, MagicMock -from aikido_zen.background_process.routes import Routes +from aikido_zen.storage.routes import Routes from .thread_cache import ThreadCache, get_cache from .. import set_user from ..background_process.packages import PackagesStore @@ -181,14 +181,12 @@ def test_parses_routes_correctly( "method": "POST", "path": "/body", "hits": 20, - "hits_delta_since_sync": 25, "apispec": {}, }, "GET:/body": { "method": "GET", "path": "/body", "hits": 10, - "hits_delta_since_sync": 5, "apispec": {}, }, }, @@ -210,22 +208,21 @@ def test_parses_routes_correctly( } ] assert thread_cache.is_user_blocked("user123") - assert list(thread_cache.routes) == [ - { + assert thread_cache.routes.export() == {} + assert thread_cache.routes_from_background.export() == { + "POST:/body": { "method": "POST", "path": "/body", "hits": 20, - "hits_delta_since_sync": 0, "apispec": {}, }, - { + "GET:/body": { "method": "GET", "path": "/body", "hits": 10, - "hits_delta_since_sync": 0, "apispec": {}, }, - ] + } @patch("aikido_zen.background_process.comms.get_comms") @@ -242,8 +239,7 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach thread_cache.stats.on_detected_attack(blocked=True, operation="op1") thread_cache.stats.on_detected_attack(blocked=False, operation="op1") thread_cache.stats.on_detected_attack(blocked=False, operation="op2") - thread_cache.routes.initialize_route({"method": "GET", "route": "/test"}) - thread_cache.routes.increment_route({"method": "GET", "route": "/test"}) + thread_cache.routes.increment_route("GET", "/test") thread_cache.ai_stats.on_ai_call("openai", "gpt-4o", 6427, 200) thread_cache.ai_stats.on_ai_call("openai", "gpt-4o2", 424, 235) thread_cache.ai_stats.on_ai_call("openai", "gpt-4o2", 232, 932) @@ -271,12 +267,11 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach mock_comms.send_data_to_bg_process.assert_called_once_with( action="SYNC_DATA", obj={ - "current_routes": { + "routes": { "GET:/test": { "method": "GET", "path": "/test", "hits": 1, - "hits_delta_since_sync": 1, "apispec": {}, } }, @@ -363,7 +358,7 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): mock_comms.send_data_to_bg_process.assert_called_once_with( action="SYNC_DATA", obj={ - "current_routes": {}, + "routes": {}, "stats": { "startedAt": -1, "endedAt": -1, @@ -415,7 +410,7 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach mock_comms.send_data_to_bg_process.assert_called_once_with( action="SYNC_DATA", obj={ - "current_routes": {}, + "routes": {}, "stats": { "startedAt": -1, "endedAt": -1, @@ -443,7 +438,7 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache mock_get_comms.return_value = mock_comms # Setup initial state with a route but no requests - thread_cache.routes.initialize_route({"method": "GET", "route": "/test"}) + thread_cache.routes.increment_route("GET", "/test") # Call renew with patch( @@ -455,7 +450,14 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache mock_comms.send_data_to_bg_process.assert_called_once_with( action="SYNC_DATA", obj={ - "current_routes": {}, + "routes": { + "GET:/test": { + "method": "GET", + "path": "/test", + "hits": 1, + "apispec": {}, + } + }, "stats": { "startedAt": -1, "endedAt": -1, diff --git a/end2end/django_mysql_test.py b/end2end/django_mysql_test.py index db6060195..ae8d7ef3f 100644 --- a/end2end/django_mysql_test.py +++ b/end2end/django_mysql_test.py @@ -84,7 +84,6 @@ def test_initial_heartbeat(): [{ "apispec": {'body': {'type': 'form-urlencoded', 'schema': {'type': 'object', 'properties': {'dog_name': {'type': 'string'}}}}, 'query': None, 'auth': None}, "hits": 1, - "hits_delta_since_sync": 0, "method": "POST", "path": "/app/create" }],