"""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 base64
import errno
import fcntl
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 termios
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
MAX_TERMINAL_BYTES = 8192
MAX_TERMINAL_FRAME = 16384
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_terminal_size(rows: object, cols: object) -> None:
    if type(rows) is not int or type(cols) is not int or not (2 <= rows <= 200 and 2 <= cols <= 400):
        raise ProtocolError("terminal dimensions must be integers: 2–200 rows and 2–400 columns")


def terminal_data(value: object) -> bytes:
    if not isinstance(value, str) or len(value) > 4 * ((MAX_TERMINAL_BYTES + 2) // 3):
        raise ProtocolError("terminal data exceeds its size limit")
    try:
        data = base64.b64decode(value, validate=True)
    except (ValueError, UnicodeError) as error:
        raise ProtocolError("invalid terminal data") from error
    if len(data) > MAX_TERMINAL_BYTES:
        raise ProtocolError("terminal data exceeds its size limit")
    return data


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"):
                stream = getattr(process, name)
                if stream is not None:
                    stream.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) -> None:
    """Enter the job cgroup and account in a fresh interpreter, then replace it."""
    (scope / "cgroup.procs").write_text(str(os.getpid()), encoding="ascii")
    if command is None:
        # Popen's start_new_session already performed setsid in this fresh child.
        fcntl.ioctl(0, termios.TIOCSCTTY, 0)
    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",
    }
    if command is None:
        shell = account.pw_shell or "/bin/sh"
        environment.update(SHELL=shell, TERM="xterm-256color", COLORTERM="truecolor")
        os.execve(shell, [shell, "-i"], environment)
    else:
        os.execve("/bin/sh", ["sh", "-c", command], environment)


def run_terminal(connection: socket.socket, username: str, rows: int, cols: int, *, scope_factory=JobScope) -> None:
    """Bridge one real guest PTY until shell exit, explicit close, or host loss."""
    validate_terminal_size(rows, cols)
    scope = scope_factory()
    process = None
    master = slave = None
    try:
        master, slave = os.openpty()
        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        process = subprocess.Popen(
            [sys.executable, "-I", str(Path(__file__).resolve()), "--terminal-child", str(scope.path), username],
            stdin=slave, stdout=slave, stderr=slave, start_new_session=True,
            env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8"},
        )
        os.close(slave)
        slave = None
        os.set_blocking(master, False)
        send_frame(connection, {"type": "ready"}, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
        pending = bytearray()
        with selectors.DefaultSelector() as selector:
            selector.register(connection, selectors.EVENT_READ, "host")
            selector.register(master, selectors.EVENT_READ, "pty")
            finished = False
            while not finished:
                for key, events in selector.select(0.1):
                    if key.data == "host":
                        message = receive_frame(connection, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
                        if message == {"type": "close"}:
                            finished = True
                            break
                        if set(message) == {"type", "data"} and message["type"] == "input":
                            pending.extend(terminal_data(message["data"]))
                            if len(pending) > 65536:
                                raise ProtocolError("terminal input backlog exceeded")
                        elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize":
                            validate_terminal_size(message["rows"], message["cols"])
                            fcntl.ioctl(master, termios.TIOCSWINSZ, struct.pack("HHHH", message["rows"], message["cols"], 0, 0))
                        else:
                            raise ProtocolError("invalid terminal control message")
                        selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty")
                    else:
                        if events & selectors.EVENT_READ:
                            try:
                                data = os.read(master, MAX_TERMINAL_BYTES)
                            except OSError as error:
                                if error.errno != errno.EIO:
                                    raise
                                data = b""
                            if not data:
                                finished = True
                                break
                            send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")},
                                       MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
                        if events & selectors.EVENT_WRITE and pending:
                            try:
                                del pending[:os.write(master, pending)]
                            except BlockingIOError:
                                pass
                            selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty")
                if _exited(process):
                    # Drain the leader's final output before terminating any background jobs.
                    for _ in range(16):
                        try:
                            data = os.read(master, MAX_TERMINAL_BYTES)
                        except (BlockingIOError, OSError):
                            break
                        if not data:
                            break
                        send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")},
                                   MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
                    break
    finally:
        try:
            _cleanup_job(process, scope)
        finally:
            for descriptor in (master, slave):
                if descriptor is not None:
                    os.close(descriptor)
    send_frame(connection, {"type": "exit", "exit_code": process.returncode}, MAX_TERMINAL_FRAME,
               time.monotonic() + CONNECT_TIMEOUT_SECONDS)


def serve_connection(connection: socket.socket, config: dict, *, runner=run_command, terminal_runner=run_terminal) -> 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", "terminal"}:
                if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION:
                    raise ProtocolError("unsupported guest protocol version")
                size = request["terminal"]
                if not isinstance(size, dict) or set(size) != {"rows", "cols"}:
                    raise ProtocolError("invalid terminal request")
                validate_terminal_size(size["rows"], size["cols"])
                terminal_runner(connection, config["username"], size["rows"], size["cols"])
                return
            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) == 4 and sys.argv[1] == "--terminal-child":
        try:
            child_main(Path(sys.argv[2]), sys.argv[3], None)
        except (OSError, KeyError) as error:
            print(f"uvm-agent: could not launch guest terminal: {error}", file=sys.stderr)
            sys.exit(126)
        return
    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()