luckiest-co 1.0.47 → 1.0.49

package/skills/trends/scripts/lib/windows_job.py

@@ -0,0 +1,138 @@
+"""Windows process jobs used to retain descendants after launcher processes exit."""
+
+from __future__ import annotations
+
+import ctypes
+from ctypes import wintypes
+
+CREATE_SUSPENDED = 0x4
+TH32CS_SNAPTHREAD = 0x4
+THREAD_SUSPEND_RESUME = 0x2
+JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
+JOB_OBJECT_EXTENDED_LIMIT_INFORMATION = 9
+
+
+class _BasicLimit(ctypes.Structure):
+    _fields_ = [
+        ("per_process_time", ctypes.c_int64),
+        ("per_job_time", ctypes.c_int64),
+        ("flags", wintypes.DWORD),
+        ("min_working_set", ctypes.c_size_t),
+        ("max_working_set", ctypes.c_size_t),
+        ("active_processes", 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_operations", "write_operations", "other_operations",
+        "read_bytes", "write_bytes", "other_bytes",
+    )]
+
+
+class _ExtendedLimit(ctypes.Structure):
+    _fields_ = [
+        ("basic", _BasicLimit),
+        ("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),
+    ]
+
+
+class _ThreadEntry(ctypes.Structure):
+    _fields_ = [
+        ("size", wintypes.DWORD),
+        ("usage", wintypes.DWORD),
+        ("thread_id", wintypes.DWORD),
+        ("owner_pid", wintypes.DWORD),
+        ("base_priority", wintypes.LONG),
+        ("delta_priority", wintypes.LONG),
+        ("flags", wintypes.DWORD),
+    ]
+
+
+def _kernel32():
+    kernel = ctypes.WinDLL("kernel32", use_last_error=True)
+    kernel.CreateJobObjectW.argtypes = (ctypes.c_void_p, wintypes.LPCWSTR)
+    kernel.CreateJobObjectW.restype = wintypes.HANDLE
+    kernel.SetInformationJobObject.argtypes = (wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD)
+    kernel.SetInformationJobObject.restype = wintypes.BOOL
+    kernel.AssignProcessToJobObject.argtypes = (wintypes.HANDLE, wintypes.HANDLE)
+    kernel.AssignProcessToJobObject.restype = wintypes.BOOL
+    kernel.TerminateJobObject.argtypes = (wintypes.HANDLE, wintypes.UINT)
+    kernel.TerminateJobObject.restype = wintypes.BOOL
+    kernel.CreateToolhelp32Snapshot.argtypes = (wintypes.DWORD, wintypes.DWORD)
+    kernel.CreateToolhelp32Snapshot.restype = wintypes.HANDLE
+    kernel.Thread32First.argtypes = (wintypes.HANDLE, ctypes.POINTER(_ThreadEntry))
+    kernel.Thread32First.restype = wintypes.BOOL
+    kernel.Thread32Next.argtypes = (wintypes.HANDLE, ctypes.POINTER(_ThreadEntry))
+    kernel.Thread32Next.restype = wintypes.BOOL
+    kernel.OpenThread.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD)
+    kernel.OpenThread.restype = wintypes.HANDLE
+    kernel.ResumeThread.argtypes = (wintypes.HANDLE,)
+    kernel.ResumeThread.restype = wintypes.DWORD
+    kernel.CloseHandle.argtypes = (wintypes.HANDLE,)
+    kernel.CloseHandle.restype = wintypes.BOOL
+    return kernel
+
+
+class WindowsJob:
+    def __init__(self):
+        self.kernel = _kernel32()
+        self.handle = self.kernel.CreateJobObjectW(None, None)
+        if not self.handle:
+            raise ctypes.WinError(ctypes.get_last_error())
+        limits = _ExtendedLimit()
+        limits.basic.flags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
+        if not self.kernel.SetInformationJobObject(
+            self.handle, JOB_OBJECT_EXTENDED_LIMIT_INFORMATION, ctypes.byref(limits), ctypes.sizeof(limits)
+        ):
+            error = ctypes.get_last_error()
+            self.close()
+            raise ctypes.WinError(error)
+
+    def assign(self, process_handle: int) -> None:
+        if not self.kernel.AssignProcessToJobObject(self.handle, process_handle):
+            raise ctypes.WinError(ctypes.get_last_error())
+
+    def terminate(self) -> bool:
+        return bool(self.handle and self.kernel.TerminateJobObject(self.handle, 1))
+
+    def close(self) -> None:
+        if self.handle:
+            self.kernel.CloseHandle(self.handle)
+            self.handle = None
+
+
+def resume_suspended_process(pid: int) -> None:
+    kernel = _kernel32()
+    snapshot = kernel.CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
+    if not snapshot or snapshot == ctypes.c_void_p(-1).value:
+        raise ctypes.WinError(ctypes.get_last_error())
+    try:
+        entry = _ThreadEntry()
+        entry.size = ctypes.sizeof(entry)
+        found = kernel.Thread32First(snapshot, ctypes.byref(entry))
+        while found:
+            if entry.owner_pid == pid:
+                thread = kernel.OpenThread(THREAD_SUSPEND_RESUME, False, entry.thread_id)
+                if not thread:
+                    raise ctypes.WinError(ctypes.get_last_error())
+                try:
+                    previous = kernel.ResumeThread(thread)
+                    if previous == 0xFFFFFFFF:
+                        raise ctypes.WinError(ctypes.get_last_error())
+                    if previous != 1:
+                        raise OSError(f"Unexpected suspend count {previous} for process {pid}")
+                    return
+                finally:
+                    kernel.CloseHandle(thread)
+            found = kernel.Thread32Next(snapshot, ctypes.byref(entry))
+        raise OSError(f"No primary thread found for suspended process {pid}")
+    finally:
+        kernel.CloseHandle(snapshot)