124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
"""Own Windows capture processes even when the task worker is forcibly terminated."""
|
|
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
|
|
class ProcessJob:
|
|
"""A non-inheritable job handle that owns a Windows process tree."""
|
|
|
|
def __init__(self):
|
|
class BasicLimits(ctypes.Structure):
|
|
_fields_ = [
|
|
("process_time", ctypes.c_int64),
|
|
("job_time", ctypes.c_int64),
|
|
("flags", wintypes.DWORD),
|
|
("minimum_working_set", ctypes.c_size_t),
|
|
("maximum_working_set", ctypes.c_size_t),
|
|
("active_process_limit", wintypes.DWORD),
|
|
("affinity", ctypes.c_size_t),
|
|
("priority", wintypes.DWORD),
|
|
("scheduling", wintypes.DWORD),
|
|
]
|
|
|
|
class IoCounters(ctypes.Structure):
|
|
_fields_ = [
|
|
(name, ctypes.c_uint64)
|
|
for name in (
|
|
"read_ops",
|
|
"write_ops",
|
|
"other_ops",
|
|
"read_bytes",
|
|
"write_bytes",
|
|
"other_bytes",
|
|
)
|
|
]
|
|
|
|
class ExtendedLimits(ctypes.Structure):
|
|
_fields_ = [
|
|
("basic", BasicLimits),
|
|
("io", IoCounters),
|
|
("process_memory", ctypes.c_size_t),
|
|
("job_memory", ctypes.c_size_t),
|
|
("peak_process_memory", ctypes.c_size_t),
|
|
("peak_job_memory", ctypes.c_size_t),
|
|
]
|
|
|
|
self._api = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
for name, arguments, result in (
|
|
("CreateJobObjectW", [wintypes.LPVOID, wintypes.LPCWSTR], wintypes.HANDLE),
|
|
(
|
|
"SetInformationJobObject",
|
|
[wintypes.HANDLE, ctypes.c_int, wintypes.LPVOID, wintypes.DWORD],
|
|
wintypes.BOOL,
|
|
),
|
|
("OpenProcess", [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD], wintypes.HANDLE),
|
|
("AssignProcessToJobObject", [wintypes.HANDLE, wintypes.HANDLE], wintypes.BOOL),
|
|
("TerminateJobObject", [wintypes.HANDLE, wintypes.UINT], wintypes.BOOL),
|
|
(
|
|
"QueryInformationJobObject",
|
|
[wintypes.HANDLE, ctypes.c_int, wintypes.LPVOID, wintypes.DWORD, wintypes.LPVOID],
|
|
wintypes.BOOL,
|
|
),
|
|
("CloseHandle", [wintypes.HANDLE], wintypes.BOOL),
|
|
):
|
|
function = getattr(self._api, name)
|
|
function.argtypes, function.restype = arguments, result
|
|
self._handle = self._api.CreateJobObjectW(None, None)
|
|
if not self._handle:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
limits = ExtendedLimits()
|
|
limits.basic.flags = 0x00002000 # JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
|
if not self._api.SetInformationJobObject(
|
|
self._handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)
|
|
):
|
|
error = ctypes.get_last_error()
|
|
self.close()
|
|
raise ctypes.WinError(error)
|
|
|
|
def assign(self, pid: int):
|
|
process = self._api.OpenProcess(0x0100 | 0x0001, False, pid)
|
|
if not process:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
try:
|
|
if not self._api.AssignProcessToJobObject(self._handle, process):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
finally:
|
|
self._api.CloseHandle(process)
|
|
|
|
def close(self):
|
|
if self._handle:
|
|
self._api.CloseHandle(self._handle)
|
|
self._handle = None
|
|
|
|
def terminate(self):
|
|
if self._handle and not self._api.TerminateJobObject(self._handle, 1):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
|
|
def active_processes(self) -> int:
|
|
if not self._handle:
|
|
return 0
|
|
|
|
class Accounting(ctypes.Structure):
|
|
_fields_ = [
|
|
("user_time", ctypes.c_int64),
|
|
("kernel_time", ctypes.c_int64),
|
|
("period_user_time", ctypes.c_int64),
|
|
("period_kernel_time", ctypes.c_int64),
|
|
("page_faults", wintypes.DWORD),
|
|
("total", wintypes.DWORD),
|
|
("active", wintypes.DWORD),
|
|
("terminated", wintypes.DWORD),
|
|
]
|
|
|
|
accounting = Accounting()
|
|
if not self._api.QueryInformationJobObject(
|
|
self._handle, 1, ctypes.byref(accounting), ctypes.sizeof(accounting), None
|
|
):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
return accounting.active
|
|
|
|
|
|
# Keep the recording implementation and its callers on the same ownership primitive.
|
|
CaptureProcessJob = ProcessJob
|