Add container images from Kustomziations to the cluster manifest

flux_local/git_repo.py

@@ -218,6 +218,26 @@ class ResourceVisitor:
     """
 
 
+@dataclass
+class DocumentVisitor:
+    """Invoked when a document is visited so the caller can intercept.
+
+    This is similar to a resource visitor, but it visits the unparsed documents
+    since they may not have explicit schemas in this project.
+    """
+
+    kinds: list[str]
+    """The resource kinds of documents to visit."""
+
+    func: Callable[[str, dict[str, Any]], None]
+    """Function called with the resource and optional content.
+
+    The function arguments are:
+      - parent: The namespaced name of the Fluxtomization or HelmRelease
+      - doc: The resource object (e.g. Pod, ConfigMap, HelmRelease, etc)
+    """
+
+
 @dataclass
 class MetadataSelector:
     """A filter for objects to select from the cluster."""
@@ -304,6 +324,9 @@ class ResourceSelector:
     cluster_policy: MetadataSelector = field(default_factory=MetadataSelector)
     """ClusterPolicy objects to return."""
 
+    doc_visitor: DocumentVisitor | None = None
+    """Raw objects to visit."""
+
 
 def is_allowed_source(sources: list[Source]) -> Callable[[Kustomization], bool]:
     """Return true if this Kustomization is from an allowed source."""
@@ -484,20 +507,23 @@ def node_name(ks: Kustomization) -> str:
 async def build_kustomization(
     kustomization: Kustomization,
     cluster_path: Path,
-    root: Path,
-    kustomization_selector: MetadataSelector,
-    helm_release_selector: MetadataSelector,
-    helm_repo_selector: MetadataSelector,
-    cluster_policy_selector: MetadataSelector,
+    selector: ResourceSelector,
     kustomize_flags: list[str],
     builder: CachableBuilder,
 ) -> tuple[Iterable[HelmRepository], Iterable[HelmRelease], Iterable[ClusterPolicy]]:
     """Build helm objects for the Kustomization."""
+
+    root: Path = selector.path.root
+    kustomization_selector: MetadataSelector = selector.kustomization
+    helm_repo_selector: MetadataSelector = selector.helm_repo
+    helm_release_selector: MetadataSelector = selector.helm_release
+    cluster_policy_selector: MetadataSelector = selector.cluster_policy
     if (
         not kustomization_selector.enabled
-        and not helm_release_selector.enabled
         and not helm_repo_selector.enabled
+        and not helm_release_selector.enabled
         and not cluster_policy_selector.enabled
+        and not selector.doc_visitor
     ):
         return ([], [], [])
 
@@ -519,23 +545,37 @@ async def build_kustomization(
             ) from err
 
         if kustomization_selector.visitor:
-            if kustomization_selector.visitor:
-                await kustomization_selector.visitor.func(
-                    cluster_path,
-                    Path(kustomization.path),
-                    kustomization,
-                    cmd,
-                )
+            await kustomization_selector.visitor.func(
+                cluster_path,
+                Path(kustomization.path),
+                kustomization,
+                cmd,
+            )
 
-        if (
-            not helm_release_selector.enabled
-            and not helm_repo_selector.enabled
-            and not cluster_policy_selector.enabled
-        ):
+        kinds = []
+        if helm_repo_selector.enabled:
+            kinds.append(HELM_REPO_KIND)
+        if helm_release_selector.enabled:
+            kinds.append(HELM_RELEASE_KIND)
+        if cluster_policy_selector.enabled:
+            kinds.append(CLUSTER_POLICY_KIND)
+        if selector.doc_visitor:
+            kinds.extend(selector.doc_visitor.kinds)
+        if not kinds:
             return ([], [], [])
-        docs = await cmd.grep(
-            f"kind=^({HELM_REPO_KIND}|{HELM_RELEASE_KIND}|{CLUSTER_POLICY_KIND})$"
-        ).objects(target_namespace=kustomization.target_namespace)
+
+        regexp = f"kind=^({'|'.join(kinds)})$"
+        docs = await cmd.grep(regexp).objects(
+            target_namespace=kustomization.target_namespace
+        )
+
+        if selector.doc_visitor:
+            doc_kinds = set(selector.doc_visitor.kinds)
+            for doc in docs:
+                if doc.get("kind") not in doc_kinds:
+                    continue
+                selector.doc_visitor.func(kustomization.namespaced_name, doc)
+
         return (
             filter(
                 helm_repo_selector.predicate,
@@ -609,11 +649,7 @@ async def update_kustomization(cluster: Cluster) -> None:
                     build_kustomization(
                         kustomization,
                         Path(cluster.path),
-                        selector.path.root,
-                        selector.kustomization,
-                        selector.helm_release,
-                        selector.helm_repo,
-                        selector.cluster_policy,
+                        selector,
                         options.kustomize_flags,
                         builder,
                     )

flux_local/image.py

@@ -0,0 +1,74 @@
+"""Helper functions for working with container images."""
+
+import logging
+from typing import Any
+
+from . import git_repo, manifest
+
+_LOGGER = logging.getLogger(__name__)
+
+
+# Object types that may have container images.
+KINDS = [
+    "Pod",
+    "Deployment",
+    "StatefulSet",
+    "ReplicaSet",
+    "DaemonSet",
+    "CronJob",
+    "Job",
+    "ReplicationController",
+]
+IMAGE_KEY = "image"
+
+
+def _extract_images(doc: dict[str, Any]) -> set[str]:
+    """Extract the image from a Kubernetes object."""
+    images: set[str] = set({})
+    for key, value in doc.items():
+        if key == IMAGE_KEY:
+            images.add(value)
+        elif isinstance(value, dict):
+            images.update(_extract_images(value))
+        elif isinstance(value, list):
+            for item in value:
+                if isinstance(item, dict):
+                    images.update(_extract_images(item))
+    return images
+
+
+class ImageVisitor:
+    """Helper that visits container image related objects.
+
+    This tracks the container images used by the kustomizations and HelmReleases
+    so they can be dumped for further verification.
+    """
+
+    def __init__(self) -> None:
+        """Initialize ImageVisitor."""
+        self.images: dict[str, set[str]] = {}
+
+    def repo_visitor(self) -> git_repo.DocumentVisitor:
+        """Return a git_repo.DocumentVisitor that points to this object."""
+
+        def add_image(name: str, doc: dict[str, Any]) -> None:
+            """Visitor function to find relevant images and record them for later inspection.
+
+            Updates the image set with the images found in the document.
+            """
+            images = _extract_images(doc)
+            if not images:
+                return
+            if name in self.images:
+                self.images[name].update(images)
+            else:
+                self.images[name] = set(images)
+
+        return git_repo.DocumentVisitor(kinds=KINDS, func=add_image)
+
+    def update_manifest(self, manifest: manifest.Manifest) -> None:
+        """Update the manifest with the images found in the repo."""
+        for cluster in manifest.clusters:
+            for kustomization in cluster.kustomizations:
+                if images := self.images.get(kustomization.namespaced_name):
+                    kustomization.images = list(images)

flux_local/manifest.py

@@ -57,15 +57,15 @@ class BaseManifest(BaseModel):
 
     _COMPACT_EXCLUDE_FIELDS: dict[str, Any] = {}
 
-    def compact_dict(self, exclude: dict[str, Any] | None = None) -> dict[str, Any]:
+    def compact_dict(self, exclude: dict[str, Any] | None = None, include: dict[str, Any] | None = None) -> dict[str, Any]:
         """Return a compact dictionary representation of the object.
 
         This is similar to `dict()` but with a specific implementation for serializing
         with variable fields removed.
         """
         if exclude is None:
             exclude = self._COMPACT_EXCLUDE_FIELDS
-        return self.dict(exclude=exclude)  # type: ignore[arg-type]
+        return self.dict(exclude=exclude, exclude_unset=True, exclude_none=True, exclude_defaults=True)  # type: ignore[arg-type]
 
     @classmethod
     def parse_yaml(cls, content: str) -> "BaseManifest":
@@ -293,6 +293,9 @@ class Kustomization(BaseManifest):
     contents: dict[str, Any] | None = None
     """Contents of the raw Kustomization document."""
 
+    images: list[str] = Field(default_factory=list)
+    """The list of images referenced in the kustomization."""
+
     @classmethod
     def parse_doc(cls, doc: dict[str, Any]) -> "Kustomization":
         """Parse a partial Kustomization from a kubernetes resource."""

flux_local/tool/get.py

@@ -1,10 +1,11 @@
 """Flux-local get action."""
 
 import logging
-from argparse import ArgumentParser, _SubParsersAction as SubParsersAction
+from argparse import ArgumentParser, BooleanOptionalAction, _SubParsersAction as SubParsersAction
 from typing import cast, Any
+import sys
 
-from flux_local import git_repo
+from flux_local import git_repo, image
 
 from .format import PrintFormatter, YamlFormatter
 from . import selector
@@ -149,6 +150,13 @@ def register(
             ),
         )
         selector.add_cluster_selector_flags(args)
+        args.add_argument(
+            "--enable-images",
+            type=str,
+            default=False,
+            action=BooleanOptionalAction,
+            help="Output container images when traversing the cluster",
+        )
         args.add_argument(
             "--output",
             "-o",
@@ -162,16 +170,42 @@ def register(
     async def run(  # type: ignore[no-untyped-def]
         self,
         output: str,
+        enable_images: bool,
         **kwargs,  # pylint: disable=unused-argument
     ) -> None:
         """Async Action implementation."""
         query = selector.build_cluster_selector(**kwargs)
         query.helm_release.enabled = output == "yaml"
+
+        image_visitor: image.ImageVisitor | None = None
+        if enable_images:
+            if output != "yaml":
+                print(
+                    "Flag --enable-images only works with --output yaml",
+                    file=sys.stderr,
+                )
+                return
+            image_visitor = image.ImageVisitor()
+            query.doc_visitor = image_visitor.repo_visitor()
+
         manifest = await git_repo.build_manifest(
             selector=query, options=selector.options(**kwargs)
         )
         if output == "yaml":
-            YamlFormatter().print([manifest.compact_dict()])
+            include: dict[str, Any] | None = None
+            if image_visitor:
+                image_visitor.update_manifest(manifest)
+                include = {
+                    "clusters": {
+                        "__all__": {
+                            "kustomizations": {
+                                "__all__": True,
+                                #"images": True,
+                            }
+                        }
+                    }
+                }
+            YamlFormatter().print([manifest.compact_dict(include=include)])
             return
 
         cols = ["path", "kustomizations"]

tests/__snapshots__/test_git_repo.ambr

@@ -39,7 +39,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([
@@ -140,7 +139,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([
@@ -226,7 +224,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([
@@ -352,7 +349,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([
@@ -438,7 +434,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([
@@ -919,7 +914,6 @@
             'cluster_policies': list([
               dict({
                 'name': 'test-allow-policy',
-                'namespace': None,
               }),
             ]),
             'helm_releases': list([

tests/test_image.py

@@ -0,0 +1,41 @@
+"""Tests for image."""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+from syrupy.assertion import SnapshotAssertion
+
+from flux_local.git_repo import build_manifest, ResourceSelector, PathSelector
+from flux_local.image import ImageVisitor
+
+TESTDATA = Path("tests/testdata/cluster8")
+CWD = Path.cwd()
+
+
+@pytest.mark.parametrize(
+    ("test_path", "expected"),
+    [
+        (
+            CWD / "tests/testdata/cluster8",
+            {"flux-system/apps": {"alpine", "busybox"}},
+        ),
+        (
+            CWD / "tests/testdata/cluster7",
+            {},
+        ),
+    ],
+)
+async def test_image_visitor(
+    snapshot: SnapshotAssertion, test_path: str, expected: dict[str, Any]
+) -> None:
+    """Tests for building the manifest."""
+
+    image_visitor = ImageVisitor()
+    query = ResourceSelector(
+        path=PathSelector(Path(test_path)),
+        doc_visitor=image_visitor.repo_visitor(),
+    )
+    await build_manifest(selector=query)
+
+    assert image_visitor.images == expected