k8s_parser.py

# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import json
import logging
import pkgutil
from collections import defaultdict
from dataclasses import dataclass
from pathlib import PurePath
from typing import Any

from pants.backend.helm.utils.yaml import YamlPath
from pants.backend.python.subsystems.python_tool_base import PythonToolRequirementsBase
from pants.backend.python.target_types import EntryPoint
from pants.backend.python.util_rules import pex
from pants.backend.python.util_rules.pex import (
    VenvPex,
    VenvPexProcess,
    VenvPexRequest,
    create_venv_pex,
)
from pants.backend.python.util_rules.pex_environment import PexEnvironment
from pants.engine.engine_aware import EngineAwareParameter, EngineAwareReturnType
from pants.engine.fs import CreateDigest, FileContent, FileEntry
from pants.engine.intrinsics import create_digest, execute_process
from pants.engine.rules import collect_rules, implicitly, rule
from pants.option.option_types import DictOption
from pants.util.logging import LogLevel
from pants.util.strutil import pluralize, softwrap

logger = logging.getLogger(__name__)

_HELM_K8S_PARSER_SOURCE = "k8s_parser_main.py"
_HELM_K8S_PARSER_PACKAGE = "pants.backend.helm.subsystems"


class HelmKubeParserSubsystem(PythonToolRequirementsBase):
    options_scope = "helm-k8s-parser"
    help_short = "Analyses K8S manifests rendered by Helm."

    default_requirements = [
        "hikaru>=1.1.0",
        "hikaru-model-28",
        "hikaru-model-27",
        "hikaru-model-26",
        "hikaru-model-25",
        "hikaru-model-24",
        "hikaru-model-23",
    ]

    register_interpreter_constraints = True
    crd = DictOption[str](
        help=softwrap(
            """
            Additional custom resource definitions be made available to all Helm processes
            or during value interpolation.
            Example:
                [helm-k8s-parser.crd]
                "filename1"="classname1"
                "filename2"="classname2"
            """
        ),
        default={},
    )

    default_lockfile_resource = (_HELM_K8S_PARSER_PACKAGE, "k8s_parser.lock")


@dataclass(frozen=True)
class _HelmKubeParserTool:
    pex: VenvPex
    crd: dict[str] = defaultdict


@rule
async def build_k8s_parser_tool(
    k8s_parser: HelmKubeParserSubsystem,
    pex_environment: PexEnvironment,
) -> _HelmKubeParserTool:
    parser_sources = pkgutil.get_data(_HELM_K8S_PARSER_PACKAGE, _HELM_K8S_PARSER_SOURCE)

    if not parser_sources:
        raise ValueError(
            f"Unable to find source to {_HELM_K8S_PARSER_SOURCE!r} in {_HELM_K8S_PARSER_PACKAGE}"
        )

    parser_file_content = FileContent(
        path="__k8s_parser.py", content=parser_sources, is_executable=True
    )

    digest_sources = [parser_file_content]

    modulename_classname = []
    if k8s_parser.crd != {}:
        for file, classname in k8s_parser.crd.items():
            crd_sources = open(file, "rb").read()
            if not crd_sources:
                raise ValueError(
                    f"Unable to find source to customer resource definition in {_HELM_K8S_PARSER_PACKAGE}"
                )
            unique_name = f"_crd_source_{hash(file)}"
            parser_file_content_source = FileContent(
                path=f"{unique_name}.py", content=crd_sources, is_executable=False
            )
            digest_sources.append(parser_file_content_source)
            modulename_classname.append((unique_name, classname))
    parser_digest = await create_digest(CreateDigest(digest_sources))

    # We use copies of site packages because hikaru gets confused with symlinked packages
    # The core hikaru package tries to load the packages containing the kubernetes-versioned models
    # using the __path__ attribute of the core package,
    # which doesn't work when the packages are symlinked from inside the namespace-handling dirs in the PEX
    use_site_packages_copies = True

    parser_pex = await create_venv_pex(
        VenvPexRequest(
            k8s_parser.to_pex_request(
                main=EntryPoint(PurePath(parser_file_content.path).stem),
                sources=parser_digest,
            ),
            pex_environment.in_sandbox(working_directory=None),
            site_packages_copies=use_site_packages_copies,
        ),
        **implicitly(),
    )
    return _HelmKubeParserTool(parser_pex, json.dumps(modulename_classname))


@dataclass(frozen=True)
class ParseKubeManifestRequest(EngineAwareParameter):
    file: FileEntry

    def debug_hint(self) -> str | None:
        return self.file.path

    def metadata(self) -> dict[str, Any] | None:
        return {"file": self.file}


@dataclass(frozen=True)
class ParsedImageRefEntry:
    document_index: int
    path: YamlPath
    unparsed_image_ref: str


@dataclass(frozen=True)
class ParsedKubeManifest(EngineAwareReturnType):
    filename: str
    found_image_refs: tuple[ParsedImageRefEntry, ...]

    def level(self) -> LogLevel | None:
        return LogLevel.DEBUG

    def message(self) -> str | None:
        return f"Found {pluralize(len(self.found_image_refs), 'image reference')} in file {self.filename}"

    def metadata(self) -> dict[str, Any] | None:
        return {
            "filename": self.filename,
            "found_image_refs": self.found_image_refs,
        }


@rule(desc="Parse Kubernetes resource manifest")
async def parse_kube_manifest(
    request: ParseKubeManifestRequest, tool: _HelmKubeParserTool
) -> ParsedKubeManifest:
    file_digest = await create_digest(CreateDigest([request.file]))

    result = await execute_process(
        **implicitly(
            VenvPexProcess(
                tool.pex,
                argv=[request.file.path, tool.crd],
                input_digest=file_digest,
                description=f"Analyzing Kubernetes manifest {request.file.path}",
                level=LogLevel.DEBUG,
            )
        )
    )

    if result.exit_code == 0:
        output = result.stdout.decode("utf-8").splitlines()
        image_refs: list[ParsedImageRefEntry] = []
        for line in output:
            parts = line.split(",")
            if len(parts) != 3:
                raise Exception(
                    softwrap(
                        f"""Unexpected output from k8s parser when parsing file {request.file.path}:

                        {line}
                        """
                    )
                )

            image_refs.append(
                ParsedImageRefEntry(
                    document_index=int(parts[0]),
                    path=YamlPath.parse(parts[1]),
                    unparsed_image_ref=parts[2],
                )
            )

        return ParsedKubeManifest(filename=request.file.path, found_image_refs=tuple(image_refs))
    else:
        parser_error = result.stderr.decode("utf-8")
        raise Exception(
            softwrap(
                f"""
                Could not parse Kubernetes manifests in file: {request.file.path}.
                {parser_error}
                """
            )
        )


def rules():
    return [
        *collect_rules(),
        *pex.rules(),
    ]