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

from __future__ import annotations

import json
import os
import sys

from argparse import Namespace
from enum import Enum, auto
from typing import TYPE_CHECKING, Iterable, TypeVar, cast

from .base import Input, InputMode, Result
from .log import DEBUG, ERR, log
from .Uri import Uri

if TYPE_CHECKING:
    from .ExecContext import ExecContext
    from .FileContext import FileContext
    from .ProcFilter import ProcFilter, ProcPipeline

T = TypeVar('T')

class AskpassKey(Enum):
    Username = auto()
    Password = auto()

def pretty_cmd(cmd: list[str] | None = None, wd = None):
    if cmd is None:
        cmd = sys.argv
    tokens = [cmd[0]]
    for token in cmd[1:]:
        if token.find(' ') != -1:
            token = '"' + token + '"'
        tokens.append(token)
    ret = ' '.join(tokens)
    if wd is not None:
        ret += f' in {wd}'
    return ret

# See ExecContext.run() for what this function does
async def run_cmd(
    *args,
    ec: ExecContext | None = None,
    verbose: bool | None = None,
    cmd_input: Input = InputMode.NonInteractive,
    **kwargs,
) -> Result:
    if verbose is None:
        verbose = False if ec is None else ec.verbose_default
    if ec is None:
        from .ec.Local import Local

        interactive = cmd_input == InputMode.Interactive
        ec = Local(verbose_default = verbose, interactive = interactive)
    kwargs['verbose'] = verbose
    return await ec.run(*args, **kwargs)

async def run_curl(
    args: list[str],
    wd = None,
    throw = None,
    verbose = None,
    cmd_input = InputMode.NonInteractive,
    ec: ExecContext | None = None,
    decode = False,
) -> Result:
    if verbose is None:
        verbose = False if ec is None else ec.verbose_default
    cmd = ['curl']
    if not verbose:
        cmd.append('-s')
    cmd.extend(args)
    return await run_cmd(
        cmd, wd = wd, throw = throw, verbose = verbose, cmd_input = cmd_input, ec = ec
    )

async def run_curl_into(
    expected_type: type[T],
    args: list[str],
    **kwargs,
) -> T:
    result = await run_curl(args, **kwargs)
    stdout = result.stdout_str
    try:
        ret = json.loads(stdout)
    except Exception as e:
        log(
            ERR,
            f'Failed to parse {len(stdout)} bytes of Curl output ({str(e)})',
            file = sys.stderr,
        )
    raise
    if not isinstance(ret, expected_type):
        raise TypeError(
            f'Expected {expected_type.__name__}, got {type(ret).__name__} from Curl'
        )
    return cast(T, ret)

async def run_askpass(
    askpass_env: list[str],
    key: AskpassKey,
    host: str | None = None,
    ec: ExecContext | None = None,
    throw: bool = False,
) -> str | None:
    if host is not None:  # Currently unsupported
        raise NotImplementedError(f'Tried to run askpass with host "{host}"')
    for var in askpass_env:
        exe = os.getenv(var)
        if exe is None:
            continue
        exe_arg = ''
        match var:
            case 'GIT_ASKPASS':
                match key:
                    case AskpassKey.Username:
                        exe_arg += 'Username'
                    case AskpassKey.Password:
                        exe_arg += 'Password'
            case 'SSH_ASKPASS':
                match key:
                    case AskpassKey.Username:
                        continue  # Can't get user name from SSH_ASKPASS
                    case AskpassKey.Password:
                        exe_arg += 'Password'
        result = await run_cmd([exe, exe_arg], throw = throw, ec = ec)
        if result.status == 0 and result.stdout_or_none is not None:
            ret = result.stdout_str_or_none
            if ret:
                return ret
    msg = (
        f"Trying to get user data from {', '.join(askpass_env)} didn't produce anything"
    )
    if throw:
        raise Exception(msg)
    log(DEBUG, msg)
    return None

async def run_sudo(
    cmd: list[str],
    *args,
    interactive: bool = True,
    ec: ExecContext | None = None,
    **kwargs,
):
    if ec is None:
        from .ec.Local import Local

        ec = Local(interactive = interactive)
    return await ec.sudo(cmd, *args, **kwargs)

async def get(
    uri: str | Uri,
    *args,
    ctx: FileContext | None = None,
    content_filter: ProcFilter | list[ProcFilter] | ProcPipeline | None = None,
    **kwargs,
) -> Result:
    uri = Uri.pimp(uri)
    if ctx is None or uri.id != ctx.uri.id:
        from .FileContext import FileContext

        ctx = FileContext.create(uri)
    from .ProcFilter import run as run_pipeline

    return await run_pipeline(await ctx.get(uri.path, *args, **kwargs), content_filter)

async def copy(
    src_uri: str | Iterable[str],
    dst: str | FileContext,
    owner: str | None = None,
    group: str | None = None,
    mode: int | None = None,
    throw = True,
) -> Exception | str | list[str]:
    if not isinstance(src_uri, str):
        ret: list[str] = []
        for uri in src_uri:  # TODO: Group identical netlocs into one CopyContext
            rr = await copy(uri, dst, owner, group, mode, throw)
            if isinstance(rr, Exception):
                return rr
            if isinstance(rr, list):
                ret.extend(rr)
            if isinstance(rr, str):
                ret.append(rr)
            else:
                raise Exception(f'copy() returned unexpected type {type(rr)}')
        return ret
    from .CopyContext import CopyContext

    async with CopyContext(src_uri, dst) as ctx:
        try:
            result = await ctx.src.get(ctx.src.root, throw = True)
            dst_path = ctx.dst.root
            if await ctx.dst.is_dir(ctx.dst.root):
                dst_path += '/' + os.path.basename(src_uri)
            await ctx.dst.put(
                path = dst_path,
                content = result.stdout,
                owner = owner,
                group = group,
                mode = mode,
                throw = True,
            )
            return dst_path
        except Exception as e:
            if throw:
                raise
            log(ERR, f'Failed to copy {src_uri} -> {dst} ({str(e)})')
            return e
    assert False, 'Unreachable code'

async def get_username(
    args: Namespace | None = None,
    url: str | None = None,
    askpass_env: list[str] = [],
    ec: ExecContext | None = None,
) -> str | None:  # export
    url_user = None if url is None else Uri(url).username
    if args is not None:
        if args.username is not None:
            if url_user is not None and url_user != args.username:
                raise Exception(
                    f'Username mismatch: called with --username="{args.username}", '
                    f'URL has user name "{url_user}"'
                )
            return args.username
    if url_user is not None:
        return url_user
    return await run_askpass(askpass_env, AskpassKey.Username, ec = ec)

async def get_password(
    args: Namespace | None = None,
    url: str | None = None,
    askpass_env: list[str] = [],
    ec: ExecContext | None = None,
) -> str | None:  # export
    if args is None and url is None and not askpass_env:
        raise Exception(
            'Neither URL nor command-line arguments nor askpass environment variable '
            "available, can't get password"
        )
    if args is not None and hasattr(args, 'password'):
        # use getattr(), because we don't necessarily want to have insecure
        # --password among options
        ret = getattr(args, 'password')
        if ret is not None:
            return ret
    if url is not None:
        ret = Uri(url).password
        if ret is not None:
            return ret
    return await run_askpass(askpass_env, AskpassKey.Password, ec = ec)

async def get_profile_env(
    throw: bool = True,
    keep: Iterable[str] | bool = False,
    ec: ExecContext | None = None,
) -> dict[str, str]:  # export
    """
    Get a fresh environment from /etc/profile

    Args:
        keep:
          - False            -> Don't keep anything
          - True             -> Keep what's in the current environment
          - List of strings  -> Keep those variables

    Returns:
        Dictionary with fresh environment
    """
    mod_env: dict[str, str] | None = None
    if (not keep) or isinstance(keep, Iterable):
        mod_env = {
            'HOME': os.environ.get('HOME', '/'),
            'USER': os.environ.get('USER', ''),
            'PATH': '/usr/bin:/bin',
        }

    # Run bash as a login shell, which sources /etc/profile, then print
    # environment as NUL-separated key=value pairs
    cmd = ['/usr/bin/env', '-i', '/bin/bash', '-lc', 'env -0']
    result = await run_cmd(
        cmd, throw = throw, verbose = True, mod_env = mod_env, ec = ec
    )
    ret: dict[str, str] = {}
    stdout = result.stdout_or_none
    if stdout is not None:
        for entry in stdout.rstrip(b'\0').split(b'\0'):
            if not entry:
                continue
            bkey, bval = entry.split(b'=', 1)
            ret[bkey.decode()] = bval.decode()
    if isinstance(keep, Iterable):
        for key in keep:
            val = os.getenv(key)
            if val is not None:
                ret[key] = val
    return ret