"""Standalone, stdlib-only guest execution agent (Python 3.10+).

This file is copied into guest images and must not import the host application.
The host also imports its wire-format helpers; importing it never starts a server.
"""

from __future__ import annotations

import json
import math
import os
import pwd
import re
import selectors
import signal
import socket
import stat
import struct
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path


PROTOCOL_VERSION = 1
AGENT_PORT = 1070
DEFAULT_TIMEOUT_SECONDS = 60.0
MAX_TIMEOUT_SECONDS = 3600.0
MAX_COMMAND_BYTES = 65536
MAX_OUTPUT_BYTES = 1024 * 1024  # Combined stdout and stderr, before UTF-8 decoding.
MAX_REQUEST_BYTES = 6 * MAX_COMMAND_BYTES + 4096
MAX_RESPONSE_BYTES = 6 * MAX_OUTPUT_BYTES + 4096
CONNECT_TIMEOUT_SECONDS = 5.0
CLEANUP_SECONDS = 5.0
MAX_CONCURRENT_COMMANDS = 8
CONFIG_PATH = Path("/etc/uvm-agent.json")


class ProtocolError(ValueError):
    """An invalid or oversized wire message; never include command text."""


class ExecutionFailure(Exception):
    def __init__(self, message: str, *, outcome_unknown: bool = False) -> None:
        super().__init__(message)
        self.outcome_unknown = outcome_unknown


def validate_command(command: object, sudo: object, timeout_seconds: object) -> None:
    if not isinstance(command, str) or not command.strip() or "\0" in command:
        raise ProtocolError("command must be non-empty text without NUL characters")
    try:
        size = len(command.encode("utf-8"))
    except UnicodeError as error:
        raise ProtocolError("command must be valid UTF-8 text") from error
    if size > MAX_COMMAND_BYTES:
        raise ProtocolError(f"command exceeds {MAX_COMMAND_BYTES} UTF-8 bytes")
    if type(sudo) is not bool:
        raise ProtocolError("sudo must be a boolean")
    try:
        valid_timeout = (
            type(timeout_seconds) in (int, float)
            and math.isfinite(timeout_seconds)
            and 0 < timeout_seconds <= MAX_TIMEOUT_SECONDS
        )
    except OverflowError:
        valid_timeout = False
    if not valid_timeout:
        raise ProtocolError(f"timeout_seconds must be finite and greater than 0, up to {MAX_TIMEOUT_SECONDS:g}")


def _remaining(deadline: float) -> float:
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        raise TimeoutError("guest protocol deadline exceeded")
    return remaining


def receive_exact(connection: socket.socket, size: int, deadline: float) -> bytes:
    parts = bytearray()
    while len(parts) < size:
        connection.settimeout(_remaining(deadline))
        chunk = connection.recv(size - len(parts))
        if not chunk:
            raise EOFError("guest connection closed")
        parts.extend(chunk)
    return bytes(parts)


def receive_frame(connection: socket.socket, limit: int, deadline: float) -> dict:
    size, = struct.unpack("!I", receive_exact(connection, 4, deadline))
    if not 0 < size <= limit:
        raise ProtocolError("invalid guest message length")
    try:
        message = json.loads(receive_exact(connection, size, deadline))
    except (ValueError, UnicodeError, RecursionError) as error:
        raise ProtocolError("invalid guest JSON message") from error
    if not isinstance(message, dict):
        raise ProtocolError("guest message must be an object")
    return message


def send_frame(connection: socket.socket, message: dict, limit: int, deadline: float) -> None:
    payload = json.dumps(message, ensure_ascii=True, allow_nan=False, separators=(",", ":")).encode("ascii")
    if len(payload) > limit:
        raise ProtocolError("guest message exceeds its size limit")
    connection.settimeout(_remaining(deadline))
    connection.sendall(struct.pack("!I", len(payload)) + payload)


class JobScope:
    """A delegated guest cgroup, keeping detached descendants in the same job."""

    def __init__(self) -> None:
        entries = Path("/proc/self/cgroup").read_text(encoding="ascii").splitlines()
        unified = next((entry[3:] for entry in entries if entry.startswith("0::")), None)
        if unified is None or ".." in Path(unified).parts:
            raise ExecutionFailure("guest execution requires a delegated cgroup v2 hierarchy")
        parent = Path("/sys/fs/cgroup") / unified.lstrip("/")
        self.path = parent / f"uvm-exec-{uuid.uuid4().hex}"
        self.path.mkdir(mode=0o700)
        if not (self.path / "cgroup.kill").is_file():
            self.path.rmdir()
            raise ExecutionFailure("guest execution requires cgroup.kill (Linux 5.14 or newer)")

    def kill(self) -> None:
        (self.path / "cgroup.kill").write_text("1", encoding="ascii")

    def close(self) -> None:
        # cgroup.events reports descendants as well as processes in this cgroup.
        deadline = time.monotonic() + 1.0
        while "populated 1" in (self.path / "cgroup.events").read_text(encoding="ascii"):
            if time.monotonic() >= deadline:
                raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=True)
            time.sleep(0.01)
        for child in sorted(self.path.rglob("*"), key=lambda path: len(path.parts), reverse=True):
            if child.is_dir():
                child.rmdir()
        self.path.rmdir()


def _signal_group(process: subprocess.Popen, number: int) -> None:
    # Do not reap the group leader until after signaling: its PID pins the PGID.
    try:
        os.killpg(process.pid, number)
    except ProcessLookupError:
        pass


def _exited(process: subprocess.Popen) -> bool:
    return os.waitid(os.P_PID, process.pid, os.WEXITED | os.WNOHANG | os.WNOWAIT) is not None


def _cleanup_job(process: subprocess.Popen | None, scope: JobScope) -> None:
    failed = False
    try:
        scope.kill()
    except (OSError, ExecutionFailure):
        failed = True
    if process is not None:
        try:
            _signal_group(process, signal.SIGKILL)
            process.wait(timeout=1.0)
        except (OSError, subprocess.TimeoutExpired):
            failed = True
        finally:
            for name in ("stdout", "stderr"):
                getattr(process, name).close()
    try:
        scope.close()
    except (OSError, ExecutionFailure):
        failed = True
    if failed:
        raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=process is not None)


def run_command(command: str, username: str, timeout_seconds: float, *, scope_factory=JobScope) -> dict:
    """Run one non-interactive guest job, drain both pipes, and bound its lifetime."""
    scope = scope_factory()
    process = None
    buffers = {"stdout": bytearray(), "stderr": bytearray()}
    captured = 0
    truncated = False
    timed_out = False
    try:
        # The fresh child joins its scope BEFORE dropping credentials or executing
        # user text. Avoid preexec_fn in the threaded agent.
        process = subprocess.Popen(
            [sys.executable, "-I", str(Path(__file__).resolve()), "--child", str(scope.path), username, command],
            stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            start_new_session=True, env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8"},
        )
        deadline = time.monotonic() + timeout_seconds
        terminating = None
        killed = False
        with selectors.DefaultSelector() as selector:
            for name in buffers:
                stream = getattr(process, name)
                os.set_blocking(stream.fileno(), False)
                selector.register(stream, selectors.EVENT_READ, name)
            while True:
                now = time.monotonic()
                exited = _exited(process)
                if terminating is None and (exited or now >= deadline):
                    timed_out = not exited
                    terminating = now
                    _signal_group(process, signal.SIGTERM if timed_out else signal.SIGKILL)
                if terminating is not None:
                    if not killed and (exited or now >= terminating + 0.5):
                        scope.kill()
                        _signal_group(process, signal.SIGKILL)
                        killed = True
                    if killed and exited and not selector.get_map():
                        break
                    if now >= terminating + 3.0:
                        raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=True)
                for key, _events in selector.select(0.05):
                    chunk = os.read(key.fileobj.fileno(), 65536)
                    if not chunk:
                        selector.unregister(key.fileobj)
                        continue
                    keep = min(len(chunk), MAX_OUTPUT_BYTES - captured)
                    buffers[key.data].extend(chunk[:keep])
                    captured += keep
                    truncated |= keep < len(chunk)
        # Finally below kills the scope and reaps the leader before returning.
    except OSError as error:
        raise ExecutionFailure("guest command execution failed", outcome_unknown=process is not None) from error
    finally:
        _cleanup_job(process, scope)
    return {
        "exit_code": None if timed_out else process.returncode,
        "stdout": buffers["stdout"].decode("utf-8", errors="replace"),
        "stderr": buffers["stderr"].decode("utf-8", errors="replace"),
        "timed_out": timed_out,
        "truncated": truncated,
    }


def child_main(scope: Path, username: str, command: str) -> None:
    """Enter the job cgroup and account in a fresh interpreter, then replace it."""
    (scope / "cgroup.procs").write_text(str(os.getpid()), encoding="ascii")
    account = pwd.getpwnam(username)
    if os.geteuid() == 0:
        os.initgroups(account.pw_name, account.pw_gid)
        os.setgid(account.pw_gid)
        os.setuid(account.pw_uid)
    elif (account.pw_uid, account.pw_gid) != (os.geteuid(), os.getegid()):
        raise PermissionError("guest agent cannot assume the requested account")
    os.chdir(account.pw_dir)
    os.umask(0o022)
    environment = {
        "HOME": account.pw_dir, "USER": account.pw_name, "LOGNAME": account.pw_name,
        "SHELL": "/bin/sh", "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
        "LANG": "C.UTF-8",
    }
    os.execve("/bin/sh", ["sh", "-c", command], environment)


def serve_connection(connection: socket.socket, config: dict, *, runner=run_command) -> None:
    with connection:
        try:
            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
            send_frame(connection, {"version": PROTOCOL_VERSION, **config}, 4096, deadline)
            request = receive_frame(connection, MAX_REQUEST_BYTES, deadline)
            if set(request) != {"version", "command", "sudo", "timeout_seconds"}:
                raise ProtocolError("invalid execution request fields")
            if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION:
                raise ProtocolError("unsupported guest protocol version")
            validate_command(request["command"], request["sudo"], request["timeout_seconds"])
            username = "root" if request["sudo"] else config["username"]
            result = runner(request["command"], username, request["timeout_seconds"])
            send_frame(connection, result, MAX_RESPONSE_BYTES, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
        except (ProtocolError, ExecutionFailure, OSError, KeyError) as error:
            # Never log request text, guest output, or credentials to the console.
            try:
                send_frame(connection, {
                    "error": "guest execution failed: " + str(error),
                    "outcome_unknown": isinstance(error, ExecutionFailure) and error.outcome_unknown,
                }, 4096, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
            except (OSError, ProtocolError):
                pass
        except EOFError:
            pass


def load_config(path: Path = CONFIG_PATH) -> dict:
    descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
    with os.fdopen(descriptor, encoding="utf-8") as source:
        metadata = os.fstat(source.fileno())
        if (not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != 0
                or metadata.st_mode & 0o077 or metadata.st_size > 4096):
            raise ValueError("guest agent configuration must be a root-owned private regular file")
        config = json.load(source)
    if (not isinstance(config, dict) or set(config) != {"vm_id", "username"}
            or not isinstance(config["vm_id"], str)
            or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,80}", config["vm_id"]) is None
            or not isinstance(config["username"], str)
            or re.fullmatch(r"[a-z_][a-z0-9_-]{0,31}\$?", config["username"]) is None):
        raise ValueError("invalid guest agent configuration")
    pwd.getpwnam(config["username"])
    return config


def main() -> None:
    if len(sys.argv) == 5 and sys.argv[1] == "--child":
        try:
            child_main(Path(sys.argv[2]), sys.argv[3], sys.argv[4])
        except (OSError, KeyError) as error:
            print(f"uvm-agent: could not launch guest command: {error}", file=sys.stderr)
            sys.exit(126)
        return
    if len(sys.argv) != 1 or os.geteuid() != 0:
        raise SystemExit("uvm-agent must run as a root systemd service")
    config = load_config()
    slots = threading.BoundedSemaphore(MAX_CONCURRENT_COMMANDS)

    def handle(connection: socket.socket) -> None:
        try:
            serve_connection(connection, config)
        finally:
            slots.release()

    with socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) as listener:
        listener.bind((socket.VMADDR_CID_ANY, AGENT_PORT))
        listener.listen(MAX_CONCURRENT_COMMANDS)
        while True:
            connection, peer = listener.accept()
            # Guest-local callers must not use the root agent to elevate themselves.
            if peer[0] != socket.VMADDR_CID_HOST:
                connection.close()
                continue
            if not slots.acquire(blocking=False):
                connection.close()
                continue
            threading.Thread(target=handle, args=(connection,), daemon=True).start()


if __name__ == "__main__":
    main()