distinct_agent/isolation/windows.py

"""Windows process restriction via a kernel Job Object.

This is a real implementation, exercised by ``tests/test_agent_isolation.py``
against real child processes on a real Windows kernel.  The tests do not check
that the API calls returned success; they check that a child which tries to
exceed a limit actually fails.

**The ordering matters more than the limits.** The child is created suspended,
assigned to the job, and only then resumed.  Assigning after the child is
already running leaves a window in which it can fork or allocate freely, and a
containment boundary with a race in it is not a boundary.

**What this does not do.** The child still runs under the operator's access
token, so it can read whatever the operator can read.  Narrowing that needs a
restricted token or an AppContainer profile, neither of which is implemented.
Do not describe this layer as a security sandbox; it is containment and
resource control.
"""

from __future__ import annotations

import ctypes
import subprocess
import sys
from collections.abc import Sequence
from ctypes import wintypes
from typing import Any

from .base import IsolationUnavailable, ProcessRestriction, RestrictionReport

MECHANISM = "windows-job-object"

# --- Win32 constants -------------------------------------------------------

JobObjectBasicUIRestrictions = 4
JobObjectExtendedLimitInformation = 9

JOB_OBJECT_LIMIT_PROCESS_TIME = 0x0000_0002
JOB_OBJECT_LIMIT_JOB_TIME = 0x0000_0004
JOB_OBJECT_LIMIT_ACTIVE_PROCESS = 0x0000_0008
JOB_OBJECT_LIMIT_PROCESS_MEMORY = 0x0000_0100
JOB_OBJECT_LIMIT_JOB_MEMORY = 0x0000_0200
JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0000_0400
JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x0000_2000

# Deliberately never set: either flag would let a child leave the job.
JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x0000_0800
JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK = 0x0000_1000

JOB_OBJECT_UILIMIT_HANDLES = 0x0000_0001
JOB_OBJECT_UILIMIT_READCLIPBOARD = 0x0000_0002
JOB_OBJECT_UILIMIT_WRITECLIPBOARD = 0x0000_0004
JOB_OBJECT_UILIMIT_SYSTEMPARAMETERS = 0x0000_0008
JOB_OBJECT_UILIMIT_DISPLAYSETTINGS = 0x0000_0010
JOB_OBJECT_UILIMIT_GLOBALATOMS = 0x0000_0020
JOB_OBJECT_UILIMIT_DESKTOP = 0x0000_0040
JOB_OBJECT_UILIMIT_EXITWINDOWS = 0x0000_0080

ALL_UI_RESTRICTIONS = (
    JOB_OBJECT_UILIMIT_HANDLES
    | JOB_OBJECT_UILIMIT_READCLIPBOARD
    | JOB_OBJECT_UILIMIT_WRITECLIPBOARD
    | JOB_OBJECT_UILIMIT_SYSTEMPARAMETERS
    | JOB_OBJECT_UILIMIT_DISPLAYSETTINGS
    | JOB_OBJECT_UILIMIT_GLOBALATOMS
    | JOB_OBJECT_UILIMIT_DESKTOP
    | JOB_OBJECT_UILIMIT_EXITWINDOWS
)

CREATE_SUSPENDED = 0x0000_0004
CREATE_NO_WINDOW = 0x0800_0000
CREATE_BREAKAWAY_FROM_JOB = 0x0100_0000

TH32CS_SNAPTHREAD = 0x0000_0004
THREAD_SUSPEND_RESUME = 0x0002
INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value

ULONG_PTR = ctypes.c_size_t
HUNDRED_NANOSECONDS_PER_SECOND = 10_000_000


# --- Structures ------------------------------------------------------------


class IO_COUNTERS(ctypes.Structure):
    _fields_ = [
        ("ReadOperationCount", ctypes.c_ulonglong),
        ("WriteOperationCount", ctypes.c_ulonglong),
        ("OtherOperationCount", ctypes.c_ulonglong),
        ("ReadTransferCount", ctypes.c_ulonglong),
        ("WriteTransferCount", ctypes.c_ulonglong),
        ("OtherTransferCount", ctypes.c_ulonglong),
    ]


class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("PerProcessUserTimeLimit", ctypes.c_int64),
        ("PerJobUserTimeLimit", ctypes.c_int64),
        ("LimitFlags", wintypes.DWORD),
        ("MinimumWorkingSetSize", ctypes.c_size_t),
        ("MaximumWorkingSetSize", ctypes.c_size_t),
        ("ActiveProcessLimit", wintypes.DWORD),
        ("Affinity", ULONG_PTR),
        ("PriorityClass", wintypes.DWORD),
        ("SchedulingClass", wintypes.DWORD),
    ]


class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION),
        ("IoInfo", IO_COUNTERS),
        ("ProcessMemoryLimit", ctypes.c_size_t),
        ("JobMemoryLimit", ctypes.c_size_t),
        ("PeakProcessMemoryUsed", ctypes.c_size_t),
        ("PeakJobMemoryUsed", ctypes.c_size_t),
    ]


class JOBOBJECT_BASIC_UI_RESTRICTIONS(ctypes.Structure):
    _fields_ = [("UIRestrictionsClass", wintypes.DWORD)]


class THREADENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wintypes.DWORD),
        ("cntUsage", wintypes.DWORD),
        ("th32ThreadID", wintypes.DWORD),
        ("th32OwnerProcessID", wintypes.DWORD),
        ("tpBasePri", wintypes.LONG),
        ("tpDeltaPri", wintypes.LONG),
        ("dwFlags", wintypes.DWORD),
    ]


def _kernel32() -> Any:
    if sys.platform != "win32":  # pragma: no cover - guarded by callers
        raise IsolationUnavailable("Job Objects exist only on Windows")
    library = ctypes.WinDLL("kernel32", use_last_error=True)
    library.CreateJobObjectW.restype = wintypes.HANDLE
    library.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
    library.SetInformationJobObject.restype = wintypes.BOOL
    library.SetInformationJobObject.argtypes = [
        wintypes.HANDLE,
        ctypes.c_int,
        wintypes.LPVOID,
        wintypes.DWORD,
    ]
    library.AssignProcessToJobObject.restype = wintypes.BOOL
    library.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
    library.IsProcessInJob.restype = wintypes.BOOL
    library.IsProcessInJob.argtypes = [
        wintypes.HANDLE,
        wintypes.HANDLE,
        ctypes.POINTER(wintypes.BOOL),
    ]
    library.CloseHandle.restype = wintypes.BOOL
    library.CloseHandle.argtypes = [wintypes.HANDLE]
    library.CreateToolhelp32Snapshot.restype = wintypes.HANDLE
    library.CreateToolhelp32Snapshot.argtypes = [wintypes.DWORD, wintypes.DWORD]
    library.Thread32First.restype = wintypes.BOOL
    library.Thread32First.argtypes = [wintypes.HANDLE, ctypes.POINTER(THREADENTRY32)]
    library.Thread32Next.restype = wintypes.BOOL
    library.Thread32Next.argtypes = [wintypes.HANDLE, ctypes.POINTER(THREADENTRY32)]
    library.OpenThread.restype = wintypes.HANDLE
    library.OpenThread.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
    library.ResumeThread.restype = wintypes.DWORD
    library.ResumeThread.argtypes = [wintypes.HANDLE]
    return library


class JobObject:
    """An anonymous Job Object with every limit already applied."""

    def __init__(self, restriction: ProcessRestriction) -> None:
        self.restriction = restriction
        self.kernel32 = _kernel32()
        handle = self.kernel32.CreateJobObjectW(None, None)
        if not handle:
            raise IsolationUnavailable(
                f"CreateJobObject failed (error {ctypes.get_last_error()})"
            )
        self.handle = handle
        self.limits: dict[str, Any] = {}
        try:
            self._apply_limits()
            if restriction.restrict_ui:
                self._apply_ui_restrictions()
        except BaseException:
            self.close()
            raise

    def _apply_limits(self) -> None:
        restriction = self.restriction
        info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
        flags = (
            JOB_OBJECT_LIMIT_ACTIVE_PROCESS
            | JOB_OBJECT_LIMIT_PROCESS_MEMORY
            | JOB_OBJECT_LIMIT_JOB_MEMORY
            | JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION
            | JOB_OBJECT_LIMIT_PROCESS_TIME
            | JOB_OBJECT_LIMIT_JOB_TIME
        )
        if restriction.kill_on_close:
            flags |= JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
        # Assert rather than assume: a breakaway flag would silently undo the
        # whole boundary, so its absence is checked, not trusted.
        assert not flags & JOB_OBJECT_LIMIT_BREAKAWAY_OK
        assert not flags & JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK

        cpu_ticks = int(restriction.cpu_seconds * HUNDRED_NANOSECONDS_PER_SECOND)
        info.BasicLimitInformation.LimitFlags = flags
        info.BasicLimitInformation.ActiveProcessLimit = restriction.max_processes
        info.BasicLimitInformation.PerProcessUserTimeLimit = cpu_ticks
        info.BasicLimitInformation.PerJobUserTimeLimit = cpu_ticks
        info.ProcessMemoryLimit = restriction.memory_bytes
        info.JobMemoryLimit = restriction.memory_bytes

        ok = self.kernel32.SetInformationJobObject(
            self.handle,
            JobObjectExtendedLimitInformation,
            ctypes.byref(info),
            ctypes.sizeof(info),
        )
        if not ok:
            raise IsolationUnavailable(
                f"SetInformationJobObject(limits) failed (error {ctypes.get_last_error()})"
            )
        self.limits = {
            "max_processes": restriction.max_processes,
            "memory_bytes": restriction.memory_bytes,
            "cpu_seconds": restriction.cpu_seconds,
            "kill_on_close": restriction.kill_on_close,
            "breakaway_allowed": False,
        }

    def _apply_ui_restrictions(self) -> None:
        info = JOBOBJECT_BASIC_UI_RESTRICTIONS()
        info.UIRestrictionsClass = ALL_UI_RESTRICTIONS
        ok = self.kernel32.SetInformationJobObject(
            self.handle,
            JobObjectBasicUIRestrictions,
            ctypes.byref(info),
            ctypes.sizeof(info),
        )
        if not ok:
            raise IsolationUnavailable(
                f"SetInformationJobObject(UI) failed (error {ctypes.get_last_error()})"
            )
        self.limits["ui_restricted"] = True

    def assign(self, process_handle: int) -> None:
        ok = self.kernel32.AssignProcessToJobObject(
            self.handle, wintypes.HANDLE(process_handle)
        )
        if not ok:
            raise IsolationUnavailable(
                f"AssignProcessToJobObject failed (error {ctypes.get_last_error()})"
            )

    def contains(self, process_handle: int) -> bool:
        """Ask the kernel whether the process really is in this job."""

        result = wintypes.BOOL(0)
        ok = self.kernel32.IsProcessInJob(
            wintypes.HANDLE(process_handle), self.handle, ctypes.byref(result)
        )
        if not ok:
            raise IsolationUnavailable(
                f"IsProcessInJob failed (error {ctypes.get_last_error()})"
            )
        return bool(result.value)

    def close(self) -> None:
        handle, self.handle = getattr(self, "handle", None), None
        if handle:
            self.kernel32.CloseHandle(handle)

    def __enter__(self) -> JobObject:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()


def _resume_process(kernel32: Any, pid: int) -> int:
    """Resume every thread of a suspended process; return how many."""

    snapshot = kernel32.CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
    if snapshot == INVALID_HANDLE_VALUE or not snapshot:
        raise IsolationUnavailable(
            f"CreateToolhelp32Snapshot failed (error {ctypes.get_last_error()})"
        )
    resumed = 0
    try:
        entry = THREADENTRY32()
        entry.dwSize = ctypes.sizeof(THREADENTRY32)
        found = kernel32.Thread32First(snapshot, ctypes.byref(entry))
        while found:
            if entry.th32OwnerProcessID == pid:
                thread = kernel32.OpenThread(THREAD_SUSPEND_RESUME, False, entry.th32ThreadID)
                if thread:
                    try:
                        if kernel32.ResumeThread(thread) != 0xFFFF_FFFF:
                            resumed += 1
                    finally:
                        kernel32.CloseHandle(thread)
            found = kernel32.Thread32Next(snapshot, ctypes.byref(entry))
    finally:
        kernel32.CloseHandle(snapshot)
    if resumed == 0:
        raise IsolationUnavailable("no suspended thread could be resumed")
    return resumed


def spawn_in_job_object(
    command: Sequence[str],
    restriction: ProcessRestriction,
    **popen_kwargs: Any,
) -> tuple[subprocess.Popen, RestrictionReport, Any]:
    """Create the child suspended, assign it, verify it, then let it run."""

    try:
        job = JobObject(restriction)
    except IsolationUnavailable as exc:
        if restriction.mode == "required":
            raise
        from .base import _NullGuard

        return (
            subprocess.Popen(command, **popen_kwargs),
            RestrictionReport(
                platform="win32", applied=False, mechanism=MECHANISM, reason=str(exc)
            ),
            _NullGuard(),
        )

    flags = int(popen_kwargs.pop("creationflags", 0))
    # A child that breaks away is a child outside the boundary. Strip the flag
    # rather than trusting a caller not to have passed it.
    flags &= ~CREATE_BREAKAWAY_FROM_JOB
    flags |= CREATE_SUSPENDED | CREATE_NO_WINDOW

    process: subprocess.Popen | None = None
    try:
        process = subprocess.Popen(command, creationflags=flags, **popen_kwargs)
        handle = getattr(process, "_handle", None)
        if not isinstance(handle, int) or handle == 0:
            # No OS handle, no job assignment, no claim of containment.
            raise IsolationUnavailable("the spawned process exposed no OS handle")
        job.assign(handle)
        if not job.contains(handle):
            raise IsolationUnavailable("the kernel does not report the child inside the job")
        _resume_process(job.kernel32, process.pid)
    except BaseException as exc:
        if process is not None and process.poll() is None:
            process.kill()
        job.close()
        if restriction.mode == "required" or not isinstance(exc, IsolationUnavailable):
            raise
        from .base import _NullGuard

        return (
            subprocess.Popen(command, **popen_kwargs),
            RestrictionReport(
                platform="win32", applied=False, mechanism=MECHANISM, reason=str(exc)
            ),
            _NullGuard(),
        )

    return (
        process,
        RestrictionReport(
            platform="win32", applied=True, mechanism=MECHANISM, limits=dict(job.limits)
        ),
        job,
    )


def describe_windows_support() -> RestrictionReport:
    """Probe support by building and discarding a job, touching no child."""

    try:
        with JobObject(ProcessRestriction(mode="auto")) as job:
            return RestrictionReport(
                platform="win32", applied=True, mechanism=MECHANISM, limits=dict(job.limits)
            )
    except Exception as exc:
        return RestrictionReport(
            platform="win32", applied=False, mechanism=MECHANISM, reason=str(exc)
        )