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

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

from __future__ import annotations

from typing import TYPE_CHECKING, Iterable

if TYPE_CHECKING:
    from typing import Sequence
    from .ExecContext import ExecContext
    from .ProcFilter import ProcFilter, ProcPipeline

import os, sys, json

from argparse import Namespace
from enum import Enum, auto

from .log import *
from .base import InputMode
from .Uri import Uri

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

def pretty_cmd(cmd: list[str], wd=None):
    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)
    return await ec.run(verbose=verbose, *args, **kwargs)

async def run_curl(args: list[str], parse_json: bool=False, wd=None, throw=None, verbose=None, cmd_input=InputMode.NonInteractive, ec: ExecContext|None=None, decode=False) -> dict|str: # export
    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)
    if parse_json:
        decode = True
    output = await run_cmd(cmd, wd=wd, throw=throw, verbose=verbose, cmd_input=cmd_input, ec=ec)
    stdout, stderr, status = output.decode() if decode else output
    if not parse_json:
        ret = stdout
    else:
        try:
            ret = json.loads(stdout)
        except Exception as e:
            size = 'unknown number of'
            try:
                size = len(stdout)
            except:
                pass
            log(ERR, f'Failed to parse {size} bytes output of command '
                  + f'>{pretty_cmd(cmd, wd)}< ({str(e)}): "{stdout}"', file=sys.stderr)
            raise
    return ret, stderr, status

async def run_askpass(askpass_env: list[str], key: AskpassKey, host: str|None=None, ec: ExecContext|None=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'
        ret, stderr, status  = await run_cmd([exe, exe_arg], throw=False, ec=ec).decode()
        if ret is not None:
            return ret
    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 = ret.append(await copy(uri, dst, owner, group, mode, throw))
            if isinstance(rr, Exception):
                return rr
        return ret
    from .CopyContext import CopyContext
    async with CopyContext(src_uri, dst) as ctx:
        try:
            content = (await ctx.src.get(ctx.src.root, throw=True)).stdout
            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=content, 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: # 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}", 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: # export
    if args is None and url is None and not askpass_env:
        raise Exception(f'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 keep == False 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] = {}
    for entry in result.stdout.rstrip(b"\0").split(b"\0"):
        if not entry:
            continue
        key, val = entry.split(b"=", 1)
        ret[key.decode()] = val.decode()
    if isinstance(keep, Iterable):
        for key in keep:
            val = os.getenv(key)
            if val is not None:
                ret[key] = val
    return ret