#!/usr/bin/env python3
"""Secure Asyncio Subprocess Runner

A practical Python 3 template for launching subprocesses with asyncio while
keeping security and reliability controls explicit.

Features:
- Argument-list execution by default (no shell)
- Optional shell mode for rare cases where shell syntax is required
- Input validation for timeout and output limits
- Exit-code checking
- Controlled stdout/stderr capture
- Cleanup on timeout or cancellation
- argparse-based CLI for safe, repeatable use

Usage examples:
  python secure_asyncio_subprocess_runner.py --cmd /usr/bin/uptime
  python secure_asyncio_subprocess_runner.py --cmd /usr/bin/ls --arg -l --arg /tmp
  python secure_asyncio_subprocess_runner.py --cmd /usr/bin/rsync --arg -a --arg /src --arg /dst --timeout 60
  python secure_asyncio_subprocess_runner.py --shell --cmd "printf 'hello\n'"

Notes:
- Avoid --shell unless you need shell features such as pipes, redirects, or globbing.
- Do not pass secrets in command-line arguments unless there is no safer alternative.
- Output is capped by default to reduce memory and log exposure risk.
"""

from __future__ import annotations

import argparse
import asyncio
import dataclasses
import os
import shlex
import sys
from typing import List, Optional, Sequence, Tuple


DEFAULT_TIMEOUT = 30.0
DEFAULT_MAX_OUTPUT_BYTES = 64 * 1024


@dataclasses.dataclass
class RunResult:
    returncode: int
    stdout: str
    stderr: str
    timed_out: bool = False


class CommandValidationError(ValueError):
    pass


class OutputLimitExceeded(RuntimeError):
    pass


async def _read_limited(stream: asyncio.StreamReader, limit: int) -> bytes:
    """Read a stream with a hard byte limit."""
    chunks: List[bytes] = []
    total = 0

    while True:
        chunk = await stream.read(4096)
        if not chunk:
            break
        total += len(chunk)
        if total > limit:
            raise OutputLimitExceeded(f"output exceeded limit of {limit} bytes")
        chunks.append(chunk)

    return b"".join(chunks)


def _validate_timeout(value: float) -> float:
    if value <= 0:
        raise CommandValidationError("timeout must be greater than zero")
    return value


def _validate_max_output_bytes(value: int) -> int:
    if value <= 0:
        raise CommandValidationError("max output bytes must be greater than zero")
    return value


def _validate_executable(path: str) -> str:
    if not path or not path.strip():
        raise CommandValidationError("command path must not be empty")
    if any(c in path for c in ["\x00", "\n", "\r"]):
        raise CommandValidationError("command path contains invalid characters")
    return path


def _validate_args(args: Sequence[str]) -> List[str]:
    validated: List[str] = []
    for arg in args:
        if arg is None:
            raise CommandValidationError("argument must not be None")
        if any(c in arg for c in ["\x00", "\r", "\n"]):
            raise CommandValidationError(f"invalid characters in argument: {arg!r}")
        validated.append(arg)
    return validated


def build_command(cmd: str, args: Sequence[str], shell: bool) -> Sequence[str] | str:
    """Build a command safely.

    - In non-shell mode, returns a list suitable for create_subprocess_exec.
    - In shell mode, returns a string suitable for create_subprocess_shell.
    """
    _validate_executable(cmd)
    validated_args = _validate_args(args)

    if shell:
        # Shell mode is intentionally explicit and should be used sparingly.
        return " ".join([shlex.quote(cmd), *[shlex.quote(a) for a in validated_args]])

    return [cmd, *validated_args]


async def run_command(
    cmd: str,
    args: Sequence[str],
    timeout: float = DEFAULT_TIMEOUT,
    max_output_bytes: int = DEFAULT_MAX_OUTPUT_BYTES,
    shell: bool = False,
) -> RunResult:
    """Run a subprocess with timeout, output limits, and cleanup."""
    timeout = _validate_timeout(timeout)
    max_output_bytes = _validate_max_output_bytes(max_output_bytes)
    command = build_command(cmd, args, shell=shell)

    if shell:
        proc = await asyncio.create_subprocess_shell(
            command,  # type: ignore[arg-type]
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=_safe_env(),
        )
    else:
        proc = await asyncio.create_subprocess_exec(
            *command,  # type: ignore[arg-type]
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=_safe_env(),
        )

    try:
        stdout_task = asyncio.create_task(_read_limited(proc.stdout, max_output_bytes))
        stderr_task = asyncio.create_task(_read_limited(proc.stderr, max_output_bytes))

        try:
            await asyncio.wait_for(proc.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            return RunResult(
                returncode=proc.returncode if proc.returncode is not None else -1,
                stdout=(await _safe_cancel(stdout_task)).decode(errors="replace"),
                stderr=(await _safe_cancel(stderr_task)).decode(errors="replace"),
                timed_out=True,
            )

        stdout = await stdout_task
        stderr = await stderr_task
        return RunResult(
            returncode=proc.returncode if proc.returncode is not None else -1,
            stdout=stdout.decode(errors="replace"),
            stderr=stderr.decode(errors="replace"),
            timed_out=False,
        )

    except asyncio.CancelledError:
        if proc.returncode is None:
            proc.kill()
            with contextlib.suppress(Exception):
                await proc.wait()
        raise
    except OutputLimitExceeded:
        if proc.returncode is None:
            proc.kill()
            with contextlib.suppress(Exception):
                await proc.wait()
        raise


async def _safe_cancel(task: asyncio.Task[bytes]) -> bytes:
    """Cancel a task and return any result if already available."""
    if task.done():
        return task.result()
    task.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        return await task
    return b""


def _safe_env() -> dict[str, str]:
    """Return a minimal environment.

    Customize this for your environment if a child process needs specific
    variables. Keeping the environment small reduces surprise dependencies.
    """
    env = {}
    for key in ("PATH", "LANG", "LC_ALL"):
        value = os.environ.get(key)
        if value:
            env[key] = value
    return env


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run a subprocess safely with asyncio and explicit controls."
    )
    parser.add_argument(
        "--cmd",
        required=True,
        help="Executable path or command string when --shell is used.",
    )
    parser.add_argument(
        "--arg",
        action="append",
        default=[],
        help="Append a single argument. Repeat for multiple arguments.",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=DEFAULT_TIMEOUT,
        help=f"Timeout in seconds (default: {DEFAULT_TIMEOUT}).",
    )
    parser.add_argument(
        "--max-output-bytes",
        type=int,
        default=DEFAULT_MAX_OUTPUT_BYTES,
        help=f"Maximum bytes to capture from each stream (default: {DEFAULT_MAX_OUTPUT_BYTES}).",
    )
    parser.add_argument(
        "--shell",
        action="store_true",
        help="Use shell execution. Avoid unless shell syntax is required.",
    )
    return parser.parse_args(argv)


async def main_async(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)

    try:
        result = await run_command(
            cmd=args.cmd,
            args=args.arg,
            timeout=args.timeout,
            max_output_bytes=args.max_output_bytes,
            shell=args.shell,
        )
    except CommandValidationError as exc:
        print(f"validation error: {exc}", file=sys.stderr)
        return 2
    except OutputLimitExceeded as exc:
        print(f"output limit exceeded: {exc}", file=sys.stderr)
        return 3
    except FileNotFoundError:
        print("command not found", file=sys.stderr)
        return 127
    except PermissionError:
        print("permission denied when starting command", file=sys.stderr)
        return 126

    if result.timed_out:
        print("command timed out", file=sys.stderr)
        if result.stderr:
            print(result.stderr, file=sys.stderr, end="" if result.stderr.endswith("\n") else "\n")
        return 124

    if result.stdout:
        print(result.stdout, end="" if result.stdout.endswith("\n") else "\n")

    if result.stderr:
        print(result.stderr, file=sys.stderr, end="" if result.stderr.endswith("\n") else "\n")

    if result.returncode != 0:
        print(f"command failed with exit code {result.returncode}", file=sys.stderr)

    return result.returncode


if __name__ == "__main__":
    try:
        import contextlib
        raise SystemExit(asyncio.run(main_async()))
    except KeyboardInterrupt:
        print("interrupted", file=sys.stderr)
        raise SystemExit(130)