jw-pkg/src/python/jw/pkg/lib/ExecContext.py

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

from __future__ import annotations

import abc, re, sys
from enum import Enum, auto
from typing import NamedTuple, TypeAlias, TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Self, Type
    from types import TracebackType

from .log import *

class InputMode(Enum):
    Interactive = auto()
    NonInteractive = auto()
    OptInteractive = auto()
    Auto = auto()

Input: TypeAlias = InputMode | bytes | str

class Result(NamedTuple):

    stdout: str|None
    stderr: str|None
    status: int|None

    def decode(self, encoding='UTF-8', errors='replace') -> Result:
        return Result(
            self.stdout.decode(encoding, errors=errors) if self.stdout is not None else None,
            self.stderr.decode(encoding, errors=errors) if self.stderr is not None else None,
            self.status
        )

class ExecContext(abc.ABC):

    class CallContext:

        def __init__(
                self,
                parent: ExecContext,
                title: str,
                cmd: list[str],
                cmd_input: Input,
                wd: str|None,
                log_prefix: str,
                throw: bool,
                verbose: bool
            ) -> None:
            self.__cmd = cmd
            self.__wd = wd
            self.__log_prefix = log_prefix
            self.__parent = parent
            self.__title = title
            self.__pretty_cmd: str|None = None
            self.__delim = title if title is not None else f'---- {parent.uri}: Running {self.pretty_cmd} -'
            delim_len = 120
            self.__delim += '-' * max(0, delim_len - len(self.__delim))

            # -- At the end of this dance, interactive needs to be either True
            # or False
            interactive: bool|None = None
            if not isinstance(cmd_input, InputMode):
                interactive = False
                self.__cmd_input = (
                    cmd_input if isinstance(cmd_input, bytes) else
                    cmd_input.encode(sys.stdout.encoding or "utf-8")
                )
            else:
                match cmd_input:
                    case InputMode.Interactive:
                        interactive = True
                    case InputMode.NonInteractive:
                        interactive = False
                    case InputMode.OptInteractive:
                        interactive = parent.interactive
                    case InputMode.Auto:
                        interactive = sys.stdin.isatty()
                if interactive is None:
                    interactive = parent.interactive
                if interactive is None:
                    interactive = sys.stdin.isatty()
                self.__cmd_input = None
            assert interactive in [ True, False ]
            self.__interactive = interactive

            self.__cmd_input = cmd_input if not isinstance(cmd_input, InputMode) else None
            self.__throw = throw
            self.__verbose = verbose if verbose is not None else parent.verbose_default

        def __enter__(self) -> CallContext:
            self.log_delim(start=True)
            return self

        def __exit__(
            self,
            exc_type: Type[BaseException]|None,
            exc_value: BaseException|None,
            traceback: TracebackType|None
        ) -> bool:
            self.log_delim(start=False)

        @property
        def log_prefix(self) -> str:
            return self.__log_prefix

        @property
        def interactive(self) -> bool:
            return self.__interactive

        @property
        def verbose(self) -> bool:
            return self.__verbose

        @property
        def cmd_input(self) -> bytes|None:
            return self.__cmd_input

        @property
        def throw(self) -> bool:
            return self.__throw

        @property
        def wd(self) -> str|None:
            return self.__wd

        @property
        def cmd(self) -> list[str]:
            return self.__cmd

        @property
        def pretty_cmd(self) -> str:
            if self.__pretty_cmd is None:
                from .util import pretty_cmd
                self.__pretty_cmd = pretty_cmd(self.__cmd, self.__wd)
            return self.__pretty_cmd

        def log(prio: int, *args, **kwargs) -> None:
            log(prio, self.__log_prefix, *args, **kwargs)

        def log_delim(self, start: bool) -> None:
            if not self.__verbose:
                return None
            if self.__interactive: # Don't log footer in interative mode
                if start:
                    log(NOTICE, self.__delim)
                return
            delim = ',' + self.__delim + ' >' if start else '`' + self.__delim + ' <'
            log(NOTICE, delim)

        def check_exit_code(self, result: Result) -> None:
            if result.status == 0:
                return
            if (self.__throw or self.__verbose):
                msg = f'Command exited with status {result.status}: {self.pretty_cmd}'
                if result.stderr:
                    msg += ': ' + result.decode().stderr.strip()
                if self.__throw:
                    raise RuntimeError(msg)

        def exception(self, result: Result, e: Exception) -> Result:
            log(ERR, self.__log_prefix, f'Failed to run {self.pretty_cmd}')
            if self.__throw:
                raise e
            return result

    def __init__(self, uri: str, interactive: bool|None=None, verbose_default=False):
        self.__uri = uri
        self.__interactive = interactive
        self.__verbose_default = verbose_default
        self.__log_name: str|None = None
        assert verbose_default is not None

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

    @property
    def uri(self) -> str:
        return self.__uri

    @property
    def log_name(self) -> str:
        if self.__log_name is None:
            from urllib.parse import urlparse
            parsed = urlparse(self.__uri)
            scheme = 'local' if parsed.scheme is None else parsed.scheme
            hostname = '' if parsed.hostname is None else ''
            self.__log_name = f'{scheme}://{hostname}'
        return self.__log_name

    @property
    def interactive(self) -> bool|None:
        return self.__interactive

    @property
    def verbose_default(self) -> bool:
        return self.__verbose_default

    @abc.abstractmethod
    async def _run(self, *args, **kwargs) -> Result:
        pass

    async def run(
        self,
        cmd: list[str],
        wd: str|None = None,
        throw: bool = True,
        verbose: bool|None = None,
        cmd_input: Input = InputMode.OptInteractive,
        env: dict[str, str]|None = None,
        title: str=None
    ) -> Result:
        """
        Run a command asynchronously and return its output

        Args:
            cmd:     Command and arguments
            wd:      Optional working directory
            throw:   Raise an exception on non-zero exit status if True
            verbose: Emit log output while the command runs
            cmd_input:
                - "InputMode.OptInteractive"  -> Let --interactive govern how to handle interactivity (default)
                - "InputMode.Interactive"     -> Inherit terminal stdin
                - "InputMode.Auto"            -> Inherit terminal stdin if it is a TTY
                - "InputMode.NonInteractive"  -> stdin from /dev/null
                - None                        -> Alias for InputMode.NonInteractive
                - otherwise                   -> Feed cmd_input to stdin
            env:     The environment the command should be run in

        Returns:
            A Result instance
              In PTY mode stderr is always None because PTY merges stdout/stderr.
        """

        # Note that in the calls to the wrapped method, cmd_input == None can
        # be returned by CallContext and is very much allowed
        assert cmd_input is not None

        ret = Result(None, None, 1)
        with self.CallContext(self, title=title, cmd=cmd, cmd_input=cmd_input, wd=wd,
                              log_prefix='|', throw=throw, verbose=verbose) as cc:
            try:
                ret = await self._run(
                    cmd=cc.cmd,
                    wd=wd,
                    verbose=cc.verbose,
                    cmd_input=cc.cmd_input,
                    env=env,
                    interactive=cc.interactive,
                    log_prefix=cc.log_prefix
                )
            except Exception as e:
                return cc.exception(ret, e)
            cc.check_exit_code(ret)
        return ret

    @abc.abstractmethod
    async def _sudo(self, *args, **kwargs) -> Result:
        pass

    async def sudo(
        self,
        cmd: list[str],
        mod_env: dict[str, str]|None=None,
        opts: list[str]|None=None,
        wd: str|None = None,
        throw: bool = True,
        verbose: bool|None = None,
        cmd_input: Input = InputMode.OptInteractive,
        env: dict[str, str]|None = None,
        title: str=None,
    ) -> Result:

        # Note that in the calls to the wrapped method, cmd_input == None can
        # be returned by CallContext and is very much allowed
        assert cmd_input is not None

        ret = Result(None, None, 1)
        if opts is None:
            opts = {}
        if mod_env is None:
            mod_env = {}
        with self.CallContext(self, title=title, cmd=cmd, cmd_input=cmd_input, wd=wd,
                              log_prefix='|', throw=throw, verbose=verbose) as cc:
            try:
                ret = await self._sudo(
                    cmd=cc.cmd,
                    mod_env=mod_env,
                    opts=opts,
                    wd=wd,
                    verbose=cc.verbose,
                    cmd_input=cc.cmd_input,
                    env=env,
                    interactive=cc.interactive,
                    log_prefix=cc.log_prefix,
                )
            except Exception as e:
                return cc.exception(ret, e)
            cc.check_exit_code(ret)
        return ret

    async def _get(
        self,
        path: str,
        wd: str|None,
        throw: bool,
        verbose: bool|None,
        title: str
    ) -> Result:
        ret = Result(None, None, 1)
        if wd is not None:
            path = wd + '/' + path
        with self.CallContext(self, title=title, cmd=['cat', path],
                              cmd_input=InputMode.NonInteractive, wd=None,
                              log_prefix='|', throw=throw, verbose=verbose) as cc:
            try:
                ret = await self._run(
                    cmd=cc.cmd,
                    wd=wd,
                    verbose=cc.verbose,
                    cmd_input=cc.cmd_input,
                    env=None,
                    interactive=cc.interactive,
                    log_prefix=cc.log_prefix
                )
            except Exception as e:
                return cc.exception(ret, e)
            cc.check_exit_code(ret)
        return ret

    async def get(
        self,
        path: str,
        wd: str|None = None,
        throw: bool = True,
        verbose: bool|None = None,
        title: str=None,
        owner: str|None=None,
        group: str|None=None,
        mode: str|None=None,
    ) -> Result:
        return await self._get(path, wd=wd, throw=throw, verbose=verbose, title=title)

    async def _put(
        self,
        content: bytes,
        path: str,
        wd: str|None,
        throw: bool,
        verbose: bool|None,
        title: str,
        owner: str|None,
        group: str|None,
        mode: str|None,
    ) -> Result:

        from .util import pretty_cmd

        async def __run(cmd: list[str], cmd_input: Input=InputMode.NonInteractive) -> Result:
            with self.CallContext(self, title=title, cmd=cmd, cmd_input=cmd_input, wd=None,
                                  log_prefix='|', throw=True, verbose=verbose) as cc:
                try:
                    ret = await self._run(
                        cmd=cc.cmd,
                        wd=cc.wd,
                        verbose=cc.verbose,
                        cmd_input=cc.cmd_input,
                        env=None,
                        interactive=cc.interactive,
                        log_prefix=cc.log_prefix
                    )
                except Exception as e:
                    return cc.exception(ret, e)
                cc.check_exit_code(ret)
            return ret

        ret = Result(None, None, 1)
        try:
            if wd is not None:
                path = wd + '/' + path
            tmp = (await __run(['mktemp', path + '.XXXXXX'])).stdout.decode().strip()
            cmds: list[dict[str, str|list[str]|bool]] = []
            cmds.append({'cmd': ['tee', tmp], 'cmd_input': content})
            if owner is not None and group is not None:
                cmds.append({'cmd': ['chown', f'{owner}:{group}', tmp]})
            elif owner is not None:
                cmds.append({'cmd': ['chown', owner, tmp]})
            elif group is not None:
                cmds.append({'cmd': ['chgrp', group, tmp]})
            if mode is not None:
                cmds.append({'cmd': ['chmod', mode, tmp]})
            cmds.append({'cmd': ['mv', tmp, path]})
            for cmd in cmds:
                log(DEBUG, f'{self.log_name}: Running {pretty_cmd(cmd['cmd'], wd)}')
                ret = await __run(**cmd)
            return ret
        except:
            if throw:
                raise
            return cc.exception(ret, e)
        return ret

    async def put(
        self,
        content: str,
        path: str,
        wd: str|None = None,
        throw: bool = True,
        verbose: bool|None = None,
        title: str=None,
        owner: str|None=None,
        group: str|None=None,
        mode: str|None=None,
    ) -> Result:
        return await self._put(content, path, wd=wd, throw=throw, verbose=verbose,
                               title=title, owner=owner, group=group, mode=mode)

    async def _close(self) -> None:
        pass

    async def close(self) -> None:
        await self._close()

    @classmethod
    def create(cls, uri: str, *args, **kwargs) -> Self:
        tokens = re.split(r'://', uri)
        schema = tokens[0] if tokens[0] != uri else 'file'
        match schema:
            case 'local' | 'file':
                from .ec.Local import Local
                return Local(uri, *args, **kwargs)
            case 'ssh':
                from .ec.SSHClient import ssh_client
                return ssh_client(uri, *args, **kwargs)
            case _:
                pass
        raise Exception(f'Can\'t create execution context for {uri} with unknown schema "{schema}"')