73 lines
3.4 KiB
Python
73 lines
3.4 KiB
Python
"""Small Win32 helpers, imported lazily on Windows."""
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
import os
|
|
|
|
class Blob(ctypes.Structure):
|
|
_fields_ = [('size', wintypes.DWORD), ('ptr', ctypes.POINTER(ctypes.c_byte))]
|
|
|
|
def protect(raw, decrypt=False, machine=False):
|
|
if os.name != 'nt':
|
|
raise RuntimeError('Windows DPAPI is required')
|
|
buffer = ctypes.create_string_buffer(raw)
|
|
src = Blob(len(raw), ctypes.cast(buffer, ctypes.POINTER(ctypes.c_byte)))
|
|
dst = Blob()
|
|
crypt = ctypes.WinDLL('crypt32', use_last_error=True)
|
|
if decrypt:
|
|
ok = crypt.CryptUnprotectData(ctypes.byref(src), None, None, None, None, 1, ctypes.byref(dst))
|
|
else:
|
|
ok = crypt.CryptProtectData(ctypes.byref(src), 'Monitor', None, None, None,
|
|
1 | (4 if machine else 0), ctypes.byref(dst))
|
|
if not ok:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
try:
|
|
return ctypes.string_at(dst.ptr, dst.size)
|
|
finally:
|
|
kernel = ctypes.WinDLL('kernel32', use_last_error=True)
|
|
kernel.LocalFree.argtypes = [ctypes.c_void_p]
|
|
kernel.LocalFree(dst.ptr)
|
|
|
|
def message(text, error=False):
|
|
ctypes.windll.user32.MessageBoxW(None, str(text), 'Monitor', 0x10 if error else 0x40)
|
|
|
|
class Job:
|
|
"""Kill all owned descendants when the supervisor exits."""
|
|
def __init__(self):
|
|
class Basic(ctypes.Structure):
|
|
_fields_ = [('process_time', ctypes.c_int64), ('job_time', ctypes.c_int64),
|
|
('flags', wintypes.DWORD), ('min_ws', ctypes.c_size_t),
|
|
('max_ws', ctypes.c_size_t), ('active', wintypes.DWORD),
|
|
('affinity', ctypes.c_size_t), ('priority', wintypes.DWORD),
|
|
('scheduling', wintypes.DWORD)]
|
|
class IO(ctypes.Structure):
|
|
_fields_ = [(n, ctypes.c_uint64) for n in ('ro','wo','oo','rb','wb','ob')]
|
|
class Extended(ctypes.Structure):
|
|
_fields_ = [('basic', Basic), ('io', IO), ('process_memory', ctypes.c_size_t),
|
|
('job_memory', ctypes.c_size_t), ('peak_process', ctypes.c_size_t),
|
|
('peak_job', ctypes.c_size_t)]
|
|
self.kernel = ctypes.WinDLL('kernel32', use_last_error=True)
|
|
self.kernel.CreateJobObjectW.restype = wintypes.HANDLE
|
|
self.kernel.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]
|
|
self.kernel.SetInformationJobObject.argtypes = [wintypes.HANDLE, ctypes.c_int,
|
|
ctypes.c_void_p, wintypes.DWORD]
|
|
self.kernel.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
|
self.kernel.CloseHandle.argtypes = [wintypes.HANDLE]
|
|
self.handle = self.kernel.CreateJobObjectW(None, None)
|
|
if not self.handle:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
info = Extended()
|
|
info.basic.flags = 0x2000 # JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
|
if not self.kernel.SetInformationJobObject(self.handle, 9, ctypes.byref(info), ctypes.sizeof(info)):
|
|
self.close()
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
|
|
def assign(self, process):
|
|
if not self.kernel.AssignProcessToJobObject(self.handle, wintypes.HANDLE(int(process._handle))):
|
|
process.kill()
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
|
|
def close(self):
|
|
if self.handle:
|
|
self.kernel.CloseHandle(self.handle)
|
|
self.handle = None
|