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

# -*- coding: utf-8 -*-

from __future__ import annotations

import re

from pathlib import Path

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Iterable
    from ....lib.FileContext import FileContext

from ....lib.log import *
from ....lib.util import run_cmd
from ....lib.TarIo import TarIo
from ....lib.ProcFilterGpg import ProcFilterGpg

from .base import Attrs
from .FilesContext import FilesContext

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]:
        for path in await self.list_compilation_targets(pkg_names):
            try:
                self.ctx.stat(path)
                log(NOTICE, f'Removing {path}')
                await self.ctx.unlink(path)
            except FileNotFoundError as e:
                log(DEBUG, f'Compilation target {path} doesn\'t exist (ignored)')
                continue

    async def compile_template_files(self, pkg_names: Iterable[str], default_attrs: Attrs) -> list[str]:
        missing = 0
        for target in await self.list_compilation_targets(pkg_names):
            if not await self.compile_template_file(target, default_attrs):
                missing += 1
        if missing > 0:
            log(WARNING, f'{missing} missing secrets found. You might want to add them and run sudo {app.cmdline} again')

    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):
            ec = self.ctx
            from ....lib.ec.Local import Local
            from ....lib.util import get
            if not isinstance(ec, Local):
                ec = Local() # Security: Use a local exec context for decrypting and filtering secrets
            return (await get(src_uri, content_filter=ProcFilterGpg(ec=ec))).stdout

        def _matches_host_prefix(path: str) -> bool:
            return re.match(r'^' + root_in_tar, path)

        def _crop_host_prefix(path: str) -> bool:
            return re.sub(r'^' + root_in_tar, '', path)

        def _crop_default_prefix(path: str) -> bool:
            return re.sub(r'^' + default, '', path)

        def _matches_default_prefix(path: str) -> bool:
            return re.match(r'^default', path)

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

        from .tar import filter as tar_filter, rewrite as tar_rewrite, extract as tar_extract, merge as tar_merge
        if only_missing:
            raise NotImplementedError('--only-missing is not yet implemented')

        secret_paths = await self.list_secret_paths(pkg_names)

        host_root_in_tar = '/'.join(reversed(self.ctx.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)
        blob_host_filtered       = tar_filter(blob_all, lambda p: re.match(host_rx, p), 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), 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 not path in extracted_paths:
               log(NOTICE, f'not extracted: {path}')
           else:
               log(NOTICE, f'extracted:     {path}')