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 @@
     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 @@
             "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/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):

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 @@
         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 @@
 
             # 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 @@
 
 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 @@
         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

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