jw-pkg/src/python/jw/pkg/lib/ec/Local.py

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

import os, sys, subprocess, asyncio

from ..ExecContext import ExecContext as Base
from ..ExecContext import Result

from ..log import *
from ..util import pretty_cmd

class Local(Base):

    def __init__(self, uri='local', *args, **kwargs) -> None:
        super().__init__(uri, *args, **kwargs)

    async def _run(
        self,
        cmd: list[str],
        wd: str|None,
        verbose: bool,
        cmd_input: bytes|None,
        env: dict[str, str]|None,
        interactive: bool,
        log_prefix: str
    ) -> Result:

        def __log(prio, *args, verbose=verbose):
            if verbose:
                log(prio, log_prefix, *args)

        def __make_pty_reader(collector: list[bytes], enc_for_verbose: str):
            def _read(fd):
                ret = os.read(fd, 1024)
                if not ret:
                    return ret
                collector.append(ret)
                return ret
            return _read

        cwd: str|None = None
        if wd is not None:
            cwd = os.getcwd()
            os.chdir(wd)

        try:

            # -- interactive mode

            if interactive:

                import pty

                def _spawn():
                    # Apply env in PTY mode by temporarily updating os.environ around spawn.
                    if env:
                        old_env = os.environ.copy()
                        try:
                            os.environ.update(env)
                            return pty.spawn(cmd, master_read=reader)
                        finally:
                            os.environ.clear()
                            os.environ.update(old_env)
                    return pty.spawn(cmd, master_read=reader)

                stdout_chunks: list[bytes] = []
                enc_for_verbose = sys.stdout.encoding or "utf-8"
                reader = __make_pty_reader(stdout_chunks, enc_for_verbose)

                exit_code = await asyncio.to_thread(_spawn)

                # PTY merges stdout/stderr
                stdout = b"".join(stdout_chunks) if stdout_chunks else None
                return Result(stdout, None, exit_code)

            # -- non-interactive mode
            stdin = asyncio.subprocess.DEVNULL if cmd_input is None else asyncio.subprocess.PIPE

            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=stdin,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
            )

            stdout_parts: list[bytes] = []
            stderr_parts: list[bytes] = []

            # -- decoding for verbose output in pipe mode
            stdout_log_enc = sys.stdout.encoding or "utf-8"
            stderr_log_enc = sys.stderr.encoding or "utf-8"

            async def read_stream(stream, prio, collector: list[bytes], log_enc: str):
                buf = b""
                while True:
                    chunk = await stream.read(4096)
                    if not chunk:
                        break
                    collector.append(chunk)
                    if verbose:
                        buf += chunk
                        while b"\n" in buf:
                            line, buf = buf.split(b"\n", 1)
                            __log(prio, line.decode(log_enc, errors="replace"))
                if verbose and buf:
                    # flush trailing partial line (no newline)
                    __log(prio, buf.decode(log_enc, errors="replace"))

            tasks = [
                asyncio.create_task(
                    read_stream(proc.stdout, NOTICE, stdout_parts, stdout_log_enc)
                ),
                asyncio.create_task(
                    read_stream(proc.stderr, ERR, stderr_parts, stderr_log_enc)
                ),
            ]

            if stdin is asyncio.subprocess.PIPE:
                proc.stdin.write(cmd_input)
                await proc.stdin.drain()
                proc.stdin.close()

            exit_code = await proc.wait()
            await asyncio.gather(*tasks)

            stdout = b"".join(stdout_parts) if stdout_parts else None
            stderr = b"".join(stderr_parts) if stderr_parts else None

            return Result(stdout, stderr, exit_code)

        finally:
            if cwd is not None:
                os.chdir(cwd)

    async def _sudo(self, cmd: list[str], mod_env: dict[str, str], opts: list[str], *args, **kwargs) -> Result:
        env: dict[str, str]|None = None
        cmd_input: bytes|None = None
        if mod_env:
            env = os.environ.copy()
            env.update(mod_env)
        cmdline = []
        if os.getuid() != 0:
            cmdline.append('/usr/bin/sudo')
            if env is not None:
                cmdline.append('--preserve-env=' + ','.join(mod_env.keys()))
            cmdline.extend(opts)
        cmdline.extend(cmd)
        return await self._run(cmdline, *args, **kwargs)