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

from __future__ import annotations

import re
import sys

from contextlib import suppress
from pathlib import Path
from typing import TYPE_CHECKING

from ....lib.ec.secret_store.Base import Base as SecretStore
from ....lib.ec.secret_store.KvFile import KvFile
from ....lib.ec.secret_store.layout import (
    DEFAULT_PREFIX,
    FT_SECRET,
    FT_SECRET_KEY,
    SECRET_SUFFIX,
    host_prefix,
)
from ....lib.FileContext import FileContext
from ....lib.log import DEBUG, NOTICE, WARNING, log
from .FilesContext import FilesContext

if TYPE_CHECKING:
    from collections.abc import Collection

    from ....lib.Distro import Distro
    from .base import Attrs

class DistroContext(FilesContext):

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

    async def match_files(self, pkg_names: Collection[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: Collection[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: Collection[str], ignore_missing: bool = False
    ) -> list[str]:
        ret = []
        for tmpl in await self.list_template_files(pkg_names):
            path = str(Path(tmpl).with_suffix(SECRET_SUFFIX))
            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: Collection[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: Collection[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: Collection[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: Collection[str],
        only_missing: bool = False,
        verbose: bool = False,
    ) -> None:

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

        fc = FileContext.create(src_uri, chroot = True)
        await fc.open()
        try:
            hostname = self.ctx.uri.hostname
            if not hostname:
                raise Exception('Have no hostname to find secrets in the store')
            if isinstance(fc, SecretStore):
                if not fc.whole:
                    raise Exception(
                        f'src {src_uri} must name the whole store, not a '
                        f'rooted subtree'
                    )
                src: FileContext = fc
                # -- Store candidates: the host tree, then the default
                #    tree
                prefixes: tuple[str, ...] = (host_prefix(hostname), DEFAULT_PREFIX)
            else:
                # -- A file tree in the on-host layout: wrap it in the kvfile store
                src = KvFile(fc, distro = self.__distro)
                prefixes = ('', )

            # -- The store and the kvfile source are alike addressed
            #    by the on-host file name
            suffix = SECRET_SUFFIX

            for entry in await self.list_compilation_targets(pkg_names):
                # -- A present entry, even an empty one, wins over the
                #    fallback
                entry_rel = f'{entry.lstrip("/")}{suffix}'
                addr = entry_rel
                keys: list[str] | None = None
                for pfx in prefixes:
                    addr = f'{pfx}/{entry_rel}' if pfx else entry_rel
                    try:
                        keys = list(await src.readdir(addr))
                        break
                    except FileNotFoundError:
                        continue
                if keys is None:
                    log(NOTICE, f'not extracted: {entry}')
                    continue
                log(NOTICE, f'extracted:     {entry}')
                pairs: list[tuple[str, bytes]] = []
                for key in keys:
                    val = await src.get(f'{addr}/{key}')
                    pairs.append((key, val.stdout))
                dst = self.store
                with suppress(FileNotFoundError):
                    await dst.unlink(entry_rel)
                await dst.mkdir(entry_rel, ftype = FT_SECRET)
                for key, content in pairs:
                    await dst.put(f'{entry_rel}/{key}', content, ftype = FT_SECRET_KEY)
        finally:
            await fc.close()