jw-pkg/src/python/jw/pkg/cmds/secrets/lib/DistroContext.py

from __future__ import annotations

import re
import sys

from pathlib import Path
from typing import TYPE_CHECKING

from ....lib.log import DEBUG, NOTICE, WARNING, log
from ....lib.ProcFilterGpg import ProcFilterGpg
from .FilesContext import FilesContext

if TYPE_CHECKING:
    from .base import Attrs
    from typing import Iterable

    from ....lib.Distro import Distro

class DistroContext(FilesContext):

    def __init__(self, distro: Distro) -> None:
        super().__init__(distro.ctx)
        self.__distro = distro

    async def match_files(self, pkg_names: Iterable[str], pattern: str) -> list[str]:
        ret: list[str] = []
        for pkg_name in pkg_names:
            for path in await self.__distro.pkg_files(pkg_name):
                if re.match(pattern, path):
                    ret.append(path)
        return ret

    async def list_template_files(self, pkg_names: Iterable[str]) -> list[str]:
        if not pkg_names:
            pkg_names = [p.name for p in await self.__distro.select()]
        return await self.match_files(pkg_names, pattern = r'.*\.jw-tmpl$')

    async def list_secret_paths(
        self, pkg_names: Iterable[str], ignore_missing: bool = False
    ) -> list[str]:
        ret = []
        for tmpl in await self.list_template_files(pkg_names):
            path = str(Path(tmpl).with_suffix('.jw-secret'))
            if ignore_missing and not await self.ctx.file_exists(path):
                continue
            ret.append(path)
        return ret

    async def list_compilation_targets(
        self, pkg_names: Iterable[str], ignore_missing: bool = False
    ) -> list[str]:
        ret = []
        for tmpl in await self.list_template_files(pkg_names):
            path = tmpl.removesuffix('.jw-tmpl')
            if ignore_missing and not await self.ctx.file_exists(path):
                continue
            ret.append(path)
        return ret

    async def remove_compilation_targets(self, pkg_names: Iterable[str]) -> list[str]:
        ret: list[str] = []
        for path in await self.list_compilation_targets(pkg_names):
            try:
                await self.ctx.stat(path)
                log(NOTICE, f'Removing {path}')
                await self.ctx.unlink(path)
                ret.append(path)
            except FileNotFoundError:
                log(DEBUG, f"Compilation target {path} doesn't exist (ignored)")
                continue
        return ret

    async def compile_template_files(
        self, pkg_names: Iterable[str], default_attrs: Attrs
    ) -> list[str]:
        ret: list[str] = []
        missing = 0
        for target in await self.list_compilation_targets(pkg_names):
            if await self.compile_template_file(target, default_attrs):
                ret.append(target)
            else:
                missing += 1
        if missing > 0:
            from ....lib.util import pretty_cmd

            cmdline = pretty_cmd(sys.argv)
            log(
                WARNING,
                (
                    f'{missing} missing secrets found. You might want to add them and '
                    f'run sudo {cmdline} again'
                ),
            )
        return ret

    async def install(
        self,
        src_uri: str,
        pkg_names: Iterable[str],
        only_missing: bool = False,
        verbose: bool = False,
    ) -> None:

        async def _read_secret_tar_blob(src_uri: str) -> bytes:
            ec = self.ctx
            from ....lib.ec.Local import Local
            from ....lib.util import get

            if not isinstance(ec, Local):
                # Security: Use a local exec context for decrypting and
                # filtering secrets
                ec = Local()
            return (await get(src_uri, content_filter = ProcFilterGpg(ec = ec))).stdout

        def _is_needed_secret(path: str) -> bool:
            return path in secret_paths

        from .tar import extract as tar_extract
        from .tar import filter as tar_filter
        from .tar import merge as tar_merge
        from .tar import rewrite as tar_rewrite

        if only_missing:
            raise NotImplementedError('--only-missing is not yet implemented')

        secret_paths = await self.list_secret_paths(pkg_names)
        hostname = self.ctx.uri.hostname
        if not hostname:
            raise Exception('Have no hostname to find secrets in tar file')
        host_root_in_tar = '/'.join(reversed(hostname.split('.')))
        host_rx = re.compile(r'^' + host_root_in_tar)
        default_rx = re.compile(r'^default')

        filtered_paths: list[str] = []
        extracted_paths: list[str] = []

        blob_all = await _read_secret_tar_blob(src_uri)
        if blob_all is None:
            raise Exception(f'Tar blob {src_uri} is empty')
        blob_host_filtered = tar_filter(
            blob_all, lambda p: re.match(host_rx, p) is not None, filtered_paths
        )
        blob_host_transformed = tar_rewrite(
            blob_host_filtered, lambda p: re.sub(host_rx, '', p)
        )
        blob_default_filtered = tar_filter(
            blob_all, lambda p: re.match(default_rx, p) is not None, filtered_paths
        )
        blob_default_transformed = tar_rewrite(
            blob_default_filtered, lambda p: re.sub(default_rx, '', p)
        )
        blob_secret_material = tar_merge(
            [blob_host_transformed, blob_default_transformed], overwrite = False
        )
        blob_needed = tar_filter(
            blob_secret_material, _is_needed_secret, extracted_paths
        )

        await tar_extract(self.ctx, blob_needed, root = '/', verbose = verbose)
        for path in secret_paths:
            if path not in extracted_paths:
                log(NOTICE, f'not extracted: {path}')
            else:
                log(NOTICE, f'extracted:     {path}')