From 9938e721f5410a55bc8111b4567bf882502ae0f7 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:40:00 +0200 Subject: [PATCH 01/14] thread_cache use PackagesStore --- aikido_zen/thread/thread_cache.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 1270ee03..e4ba91bf 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -1,6 +1,7 @@ """Exports class ThreadConfig""" 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.background_process.service_config import ServiceConfig from aikido_zen.storage.ai_statistics import AIStatistics @@ -48,6 +49,7 @@ def reset(self): self.users.clear() self.stats.clear() self.ai_stats.clear() + PackagesStore.clear() def renew(self): if not comms.get_comms(): @@ -63,6 +65,7 @@ def renew(self): "users": self.users.as_array(), "stats": self.stats.get_record(), "ai_stats": self.ai_stats.get_stats(), + "packages": PackagesStore.get_packages() }, receive=True, ) From 7e322d86a6d117a3302682c564f8ea15739260eb Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:40:35 +0200 Subject: [PATCH 02/14] Create new builtins_import wrapper --- aikido_zen/sinks/builtins_import.py | 42 +++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 aikido_zen/sinks/builtins_import.py diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py new file mode 100644 index 00000000..a39486bc --- /dev/null +++ b/aikido_zen/sinks/builtins_import.py @@ -0,0 +1,42 @@ +import importlib.metadata +from importlib.metadata import PackageNotFoundError + +from aikido_zen.background_process.packages import PackagesStore +from aikido_zen.sinks import on_import, patch_function, after + + +@after +def _import(func, instance, args, kwargs, return_value): + if not hasattr(return_value, "__file__"): + return # Would be built-in into the interpreter (system package) + + if not hasattr(return_value, "__package__"): + return + name = getattr(return_value, "__package__") + + if not name or "." in name: + # Make sure the name exists and that it's not a submodule + return + if name == "importlib_metadata": + # Avoid circular dependencies + return + + if PackagesStore.get_package(name): + return + + version = None + try: + version = importlib.metadata.version(name) + except PackageNotFoundError: + pass + if version: + PackagesStore.add_package(name, version) + + +@on_import("builtins") +def patch(m): + """ + patching module builtins + - patches builtins.__import__ + """ + patch_function(m, "__import__", _import) From 5983b95fd56366cef7b8632b4ad212d121684afb Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:40:52 +0200 Subject: [PATCH 03/14] Add import for `builtins_import` --- aikido_zen/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py index 47388d61..168b4489 100644 --- a/aikido_zen/__init__.py +++ b/aikido_zen/__init__.py @@ -36,6 +36,8 @@ def protect(mode="daemon"): if mode == "daemon_disabled": logger.debug("Not starting the background process, daemon disabled.") + import aikido_zen.sinks.builtins_import + # Import sources import aikido_zen.sources.django import aikido_zen.sources.flask From 0ca9862164b75dd8a6a5af414bdfa2b1abf96a28 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:41:03 +0200 Subject: [PATCH 04/14] Send packages through heartbeat --- .../cloud_connection_manager/send_heartbeat.py | 5 +++++ 1 file changed, 5 insertions(+) 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 56a4cb9f..d96a4f41 100644 --- a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py +++ b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py @@ -1,5 +1,6 @@ """Exports the send_heartbeat function""" +from aikido_zen.background_process.packages import PackagesStore from aikido_zen.helpers.logging import logger from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms @@ -16,12 +17,15 @@ def send_heartbeat(connection_manager): routes = list(connection_manager.routes) outgoing_domains = connection_manager.hostnames.as_array() ai_stats = connection_manager.ai_stats.get_stats() + packages = PackagesStore.get_packages() connection_manager.statistics.clear() connection_manager.users.clear() connection_manager.routes.clear() connection_manager.hostnames.clear() connection_manager.ai_stats.clear() + PackagesStore.clear() + res = connection_manager.api.report( connection_manager.token, { @@ -31,6 +35,7 @@ def send_heartbeat(connection_manager): "stats": stats, "ai": ai_stats, "hostnames": outgoing_domains, + "packages": packages, "routes": routes, "users": users, "middlewareInstalled": connection_manager.middleware_installed, From 1f7117d63a63069a912bcebc3117761399948850 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:41:22 +0200 Subject: [PATCH 05/14] sync_data merge packages --- aikido_zen/background_process/commands/sync_data.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aikido_zen/background_process/commands/sync_data.py b/aikido_zen/background_process/commands/sync_data.py index 9137a4f8..bdaaada7 100644 --- a/aikido_zen/background_process/commands/sync_data.py +++ b/aikido_zen/background_process/commands/sync_data.py @@ -1,6 +1,7 @@ """Exports process_renew_config""" from aikido_zen.api_discovery.update_route_info import update_route_info +from aikido_zen.background_process.packages import PackagesStore from aikido_zen.helpers.logging import logger @@ -49,6 +50,9 @@ def process_sync_data(connection_manager, data, conn, queue=None): # Sync ai stats connection_manager.ai_stats.import_list(data.get("ai_stats", [])) + # Sync packages + PackagesStore.import_list(data.get("packages", [])) + if connection_manager.conf.last_updated_at > 0: # Only report data if the config has been fetched. return { From 7050b09271d0ccdca1f8aa9b5030a2ab7200c170 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:41:42 +0200 Subject: [PATCH 06/14] Update packages.py to allow more versatile PackagesStore --- aikido_zen/background_process/packages.py | 30 ++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index e8712e89..ec4cd93b 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -3,6 +3,8 @@ import importlib.metadata as importlib_metadata from packaging.version import Version + +from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms from aikido_zen.helpers.logging import logger # If any version is supported, this constant can be used @@ -59,14 +61,22 @@ class PackagesStore: @staticmethod def get_packages(): global packages - return packages + result = [] + for package in packages.values(): + if package.get("cleared", False): + continue + result.append(dict(package)) + return result @staticmethod - def add_package(package, version, supported): + def add_package(package, version, supported=None): global packages packages[package] = { + "name": package, "version": version, - "supported": bool(supported), + "requiredAt": get_unixtime_ms(), + "supported": supported, + "cleared": False, } @staticmethod @@ -75,3 +85,17 @@ def get_package(package_name): if package_name in packages: return packages[package_name] return None + + @staticmethod + def clear(): + global packages + for package in packages: + packages[package]["cleared"] = True + + @staticmethod + def import_list(imported_packages): + for package in imported_packages: + if PackagesStore.get_package(package["name"]): + continue + global packages + packages[package["name"]] = package From e39e5ad76008f5ea003a6e9fdcdd09b7333abb96 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:41:51 +0200 Subject: [PATCH 07/14] linting on packages --- aikido_zen/thread/thread_cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index e4ba91bf..f6a4d77f 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -65,7 +65,7 @@ def renew(self): "users": self.users.as_array(), "stats": self.stats.get_record(), "ai_stats": self.ai_stats.get_stats(), - "packages": PackagesStore.get_packages() + "packages": PackagesStore.get_packages(), }, receive=True, ) From 4854b77f7b9df744e0836d90312a373fd13d1bd1 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:48:49 +0200 Subject: [PATCH 08/14] Update thread_cache test cases for new "packages" object --- aikido_zen/thread/thread_cache_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index 17f5074f..b970df0b 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -313,6 +313,7 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "middleware_installed": False, "hostnames": [], "users": [], + "packages": [], }, receive=True, ) @@ -358,6 +359,7 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): "middleware_installed": False, "hostnames": [], "ai_stats": [], + "packages": [], "users": [ { "id": "123", @@ -410,6 +412,7 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach "hostnames": [], "users": [], "ai_stats": [], + "packages": [], }, receive=True, ) @@ -449,6 +452,7 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache "hostnames": [], "users": [], "ai_stats": [], + "packages": [], }, receive=True, ) From ae3e54649affa4c0038caf2f96d8d9d649a7b509 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:50:33 +0200 Subject: [PATCH 09/14] Clear up the cleared flag on packages --- .../cloud_connection_manager/send_heartbeat.py | 2 +- aikido_zen/background_process/packages.py | 5 ++++- aikido_zen/thread/thread_cache.py | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) 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 d96a4f41..9a1ff988 100644 --- a/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py +++ b/aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py @@ -17,7 +17,7 @@ def send_heartbeat(connection_manager): routes = list(connection_manager.routes) outgoing_domains = connection_manager.hostnames.as_array() ai_stats = connection_manager.ai_stats.get_stats() - packages = PackagesStore.get_packages() + packages = PackagesStore.export() connection_manager.statistics.clear() connection_manager.users.clear() diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index ec4cd93b..c333a351 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -59,7 +59,7 @@ def is_version_supported(version, required_version): class PackagesStore: @staticmethod - def get_packages(): + def export(): global packages result = [] for package in packages.values(): @@ -88,6 +88,9 @@ def get_package(package_name): @staticmethod def clear(): + # To clear we set the `cleared` attribute to True + # This is to ensure you can still get the packages + # But that they will not show up during an export global packages for package in packages: packages[package]["cleared"] = True diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index f6a4d77f..d7fe8bb0 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -65,7 +65,7 @@ def renew(self): "users": self.users.as_array(), "stats": self.stats.get_record(), "ai_stats": self.ai_stats.get_stats(), - "packages": PackagesStore.get_packages(), + "packages": PackagesStore.export(), }, receive=True, ) From 12ec19b614685833786b0e0cdbd90cfe13d80740 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:52:41 +0200 Subject: [PATCH 10/14] Update sync_data test cases --- .../background_process/commands/sync_data_test.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/aikido_zen/background_process/commands/sync_data_test.py b/aikido_zen/background_process/commands/sync_data_test.py index de5958c9..b12451a0 100644 --- a/aikido_zen/background_process/commands/sync_data_test.py +++ b/aikido_zen/background_process/commands/sync_data_test.py @@ -5,6 +5,7 @@ from .sync_data import process_sync_data from aikido_zen.background_process.routes import Routes from aikido_zen.helpers.iplist import IPList +from ..packages import PackagesStore from ...storage.hostnames import Hostnames from ...storage.statistics import Statistics @@ -62,6 +63,13 @@ def test_process_sync_data_initialization(setup_connection_manager): }, "middleware_installed": False, "hostnames": test_hostnames.as_array(), + "packages": [ + { + "name": "test-package", + "version": "2.2.0", + "cleared": False, + } + ], } result = process_sync_data(connection_manager, data, None) @@ -96,6 +104,11 @@ def test_process_sync_data_initialization(setup_connection_manager): {"hits": 15, "hostname": "example2.com", "port": 443}, {"hits": 1, "hostname": "bumblebee.com", "port": 8080}, ] + assert PackagesStore.get_package("test-package") == { + "name": "test-package", + "version": "2.2.0", + "cleared": False, + } def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_manager): From d32c9b8f9a905e024551e02d795b6b6b178ad994 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 17:58:52 +0200 Subject: [PATCH 11/14] Update thread_cache_test test cases --- aikido_zen/background_process/packages.py | 4 ++-- aikido_zen/thread/thread_cache_test.py | 22 +++++++++++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index c333a351..0cbdb5e0 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -4,7 +4,7 @@ from packaging.version import Version -from aikido_zen.helpers.get_current_unixtime_ms import get_unixtime_ms +import aikido_zen.helpers.get_current_unixtime_ms as t from aikido_zen.helpers.logging import logger # If any version is supported, this constant can be used @@ -74,7 +74,7 @@ def add_package(package, version, supported=None): packages[package] = { "name": package, "version": version, - "requiredAt": get_unixtime_ms(), + "requiredAt": t.get_unixtime_ms(), "supported": supported, "cleared": False, } diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index b970df0b..00e404b3 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -3,6 +3,7 @@ from aikido_zen.background_process.routes import Routes from .thread_cache import ThreadCache, get_cache from .. import set_user +from ..background_process.packages import PackagesStore from ..background_process.service_config import ServiceConfig from ..context import current_context, Context from aikido_zen.helpers.iplist import IPList @@ -252,9 +253,20 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach with patch( "aikido_zen.helpers.get_current_unixtime_ms.get_unixtime_ms", return_value=-1 ): + PackagesStore.add_package("test-package-4", "4.3.0") + PackagesStore.clear() + PackagesStore.add_package("test-package-1", "4.3.0") thread_cache.renew() assert thread_cache.ai_stats.empty() + assert PackagesStore.get_package("test-package-1") == { + "cleared": True, + "name": "test-package-1", + "requiredAt": -1, + "supported": None, + "version": "4.3.0", + } + assert PackagesStore.export() == [] # Assert that send_data_to_bg_process was called with the correct arguments mock_comms.send_data_to_bg_process.assert_called_once_with( @@ -313,7 +325,15 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "middleware_installed": False, "hostnames": [], "users": [], - "packages": [], + "packages": [ + { + "name": "test-package-1", + "version": "4.3.0", + "requiredAt": -1, + "supported": None, + "cleared": False, + } + ], }, receive=True, ) From 8a7a760137459cb6e12826ab37fc9b3e7bd70edb Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 18:21:09 +0200 Subject: [PATCH 12/14] Default value of supported should be true --- aikido_zen/background_process/packages.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index 0cbdb5e0..22f448f7 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -69,7 +69,7 @@ def export(): return result @staticmethod - def add_package(package, version, supported=None): + def add_package(package, version, supported=True): global packages packages[package] = { "name": package, From 15e480366a8ee64faafb9df1170b0a45273ef18d Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 18:22:28 +0200 Subject: [PATCH 13/14] remove supported flag from packages.py --- aikido_zen/background_process/packages.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index 22f448f7..bc11dadf 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -22,7 +22,8 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=N for package in packages: # Checks if we already looked up the package : if PackagesStore.get_package(package) is not None: - return PackagesStore.get_package(package)["supported"] + package_version = PackagesStore.get_package(package)["version"] + return is_version_supported(package_version, required_version) # Safely get the package version, with an exception for when the package was not found try: @@ -32,7 +33,7 @@ def is_package_compatible(package=None, required_version=ANY_VERSION, packages=N # Check support and store package for later supported = is_version_supported(package_version, required_version) - PackagesStore.add_package(package, package_version, supported) + PackagesStore.add_package(package, package_version) if supported: logger.debug( @@ -69,13 +70,12 @@ def export(): return result @staticmethod - def add_package(package, version, supported=True): + def add_package(package, version): global packages packages[package] = { "name": package, "version": version, "requiredAt": t.get_unixtime_ms(), - "supported": supported, "cleared": False, } From 8ae6266430890407564f6dad9737df9047190db0 Mon Sep 17 00:00:00 2001 From: BitterPanda Date: Wed, 11 Jun 2025 18:28:03 +0200 Subject: [PATCH 14/14] Fix test cases still using "supported --- aikido_zen/thread/thread_cache_test.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index 00e404b3..c64e6fad 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -263,7 +263,6 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "cleared": True, "name": "test-package-1", "requiredAt": -1, - "supported": None, "version": "4.3.0", } assert PackagesStore.export() == [] @@ -330,7 +329,6 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "name": "test-package-1", "version": "4.3.0", "requiredAt": -1, - "supported": None, "cleared": False, } ],