Python SCA: Report packages

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

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,

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 {

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
@@ -57,16 +59,24 @@ 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, 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,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

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)

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,
         )

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,
     )