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 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..9a1ff988 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.export() 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, 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 { 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): diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index e8712e89..bc11dadf 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 + +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 @@ -20,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: @@ -30,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( @@ -57,16 +60,23 @@ def is_version_supported(version, required_version): class PackagesStore: @staticmethod - def get_packages(): + def export(): 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): global packages packages[package] = { + "name": package, "version": version, - "supported": bool(supported), + "requiredAt": t.get_unixtime_ms(), + "cleared": False, } @staticmethod @@ -75,3 +85,20 @@ def get_package(package_name): if package_name in packages: return packages[package_name] return None + + @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 + + @staticmethod + def import_list(imported_packages): + for package in imported_packages: + if PackagesStore.get_package(package["name"]): + continue + global packages + packages[package["name"]] = package 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) diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 1270ee03..d7fe8bb0 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.export(), }, receive=True, ) diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index 17f5074f..c64e6fad 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,19 @@ 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, + "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,6 +324,14 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach "middleware_installed": False, "hostnames": [], "users": [], + "packages": [ + { + "name": "test-package-1", + "version": "4.3.0", + "requiredAt": -1, + "cleared": False, + } + ], }, receive=True, ) @@ -358,6 +377,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 +430,7 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach "hostnames": [], "users": [], "ai_stats": [], + "packages": [], }, receive=True, ) @@ -449,6 +470,7 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache "hostnames": [], "users": [], "ai_stats": [], + "packages": [], }, receive=True, )