sashatrask
450 lines
17 KiB
Python
450 lines
17 KiB
Python
import ctypes
|
|
import os
|
|
import select
|
|
import signal
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
from runtime_security import canonical_path
|
|
|
|
|
|
if os.name == 'nt':
|
|
from ctypes import wintypes
|
|
|
|
class _FILETIME(ctypes.Structure):
|
|
_fields_ = [('dwLowDateTime', wintypes.DWORD), ('dwHighDateTime', wintypes.DWORD)]
|
|
|
|
class _UNICODE_STRING(ctypes.Structure):
|
|
_fields_ = [
|
|
('Length', wintypes.USHORT),
|
|
('MaximumLength', wintypes.USHORT),
|
|
('Buffer', ctypes.c_void_p),
|
|
]
|
|
|
|
_P_DWORD = ctypes.POINTER(wintypes.DWORD)
|
|
_P_ULONG = ctypes.POINTER(wintypes.ULONG)
|
|
_P_BOOL = ctypes.POINTER(wintypes.BOOL)
|
|
_P_FILETIME = ctypes.POINTER(_FILETIME)
|
|
_P_UNICODE_STRING = ctypes.POINTER(_UNICODE_STRING)
|
|
_P_INT = ctypes.POINTER(ctypes.c_int)
|
|
_P_LPWSTR = ctypes.POINTER(wintypes.LPWSTR)
|
|
|
|
_KERNEL32 = ctypes.WinDLL('kernel32', use_last_error=True)
|
|
_NTDLL = ctypes.WinDLL('ntdll', use_last_error=True)
|
|
_SHELL32 = ctypes.WinDLL('shell32', use_last_error=True)
|
|
|
|
_GET_EXIT_CODE_PROCESS = _KERNEL32.GetExitCodeProcess
|
|
_GET_EXIT_CODE_PROCESS.argtypes = [wintypes.HANDLE, _P_DWORD]
|
|
_GET_EXIT_CODE_PROCESS.restype = wintypes.BOOL
|
|
_WAIT_FOR_SINGLE_OBJECT = _KERNEL32.WaitForSingleObject
|
|
_WAIT_FOR_SINGLE_OBJECT.argtypes = [wintypes.HANDLE, wintypes.DWORD]
|
|
_WAIT_FOR_SINGLE_OBJECT.restype = wintypes.DWORD
|
|
_CLOSE_HANDLE = _KERNEL32.CloseHandle
|
|
_CLOSE_HANDLE.argtypes = [wintypes.HANDLE]
|
|
_CLOSE_HANDLE.restype = wintypes.BOOL
|
|
_GET_PROCESS_TIMES = _KERNEL32.GetProcessTimes
|
|
_GET_PROCESS_TIMES.argtypes = [
|
|
wintypes.HANDLE, _P_FILETIME, _P_FILETIME, _P_FILETIME, _P_FILETIME,
|
|
]
|
|
_GET_PROCESS_TIMES.restype = wintypes.BOOL
|
|
_QUERY_FULL_PROCESS_IMAGE_NAME = _KERNEL32.QueryFullProcessImageNameW
|
|
_QUERY_FULL_PROCESS_IMAGE_NAME.argtypes = [
|
|
wintypes.HANDLE, wintypes.DWORD, wintypes.LPWSTR, _P_DWORD,
|
|
]
|
|
_QUERY_FULL_PROCESS_IMAGE_NAME.restype = wintypes.BOOL
|
|
_IS_PROCESS_IN_JOB = _KERNEL32.IsProcessInJob
|
|
_IS_PROCESS_IN_JOB.argtypes = [wintypes.HANDLE, wintypes.HANDLE, _P_BOOL]
|
|
_IS_PROCESS_IN_JOB.restype = wintypes.BOOL
|
|
_OPEN_PROCESS = _KERNEL32.OpenProcess
|
|
_OPEN_PROCESS.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
|
|
_OPEN_PROCESS.restype = wintypes.HANDLE
|
|
_GET_CURRENT_PROCESS = _KERNEL32.GetCurrentProcess
|
|
_GET_CURRENT_PROCESS.argtypes = []
|
|
_GET_CURRENT_PROCESS.restype = wintypes.HANDLE
|
|
_TERMINATE_PROCESS = _KERNEL32.TerminateProcess
|
|
_TERMINATE_PROCESS.argtypes = [wintypes.HANDLE, wintypes.UINT]
|
|
_TERMINATE_PROCESS.restype = wintypes.BOOL
|
|
_LOCAL_FREE = _KERNEL32.LocalFree
|
|
_LOCAL_FREE.argtypes = [wintypes.HLOCAL]
|
|
_LOCAL_FREE.restype = wintypes.HLOCAL
|
|
_NT_QUERY_INFORMATION_PROCESS = _NTDLL.NtQueryInformationProcess
|
|
_NT_QUERY_INFORMATION_PROCESS.argtypes = [
|
|
wintypes.HANDLE, wintypes.ULONG, ctypes.c_void_p, wintypes.ULONG, _P_ULONG,
|
|
]
|
|
_NT_QUERY_INFORMATION_PROCESS.restype = ctypes.c_long
|
|
_COMMAND_LINE_TO_ARGV = _SHELL32.CommandLineToArgvW
|
|
_COMMAND_LINE_TO_ARGV.argtypes = [wintypes.LPCWSTR, _P_INT]
|
|
_COMMAND_LINE_TO_ARGV.restype = _P_LPWSTR
|
|
else:
|
|
_FILETIME = _UNICODE_STRING = None
|
|
_KERNEL32 = _NTDLL = _SHELL32 = None
|
|
|
|
|
|
class ProcessIdentityError(OSError):
|
|
pass
|
|
|
|
|
|
class ProcessExitedError(ProcessIdentityError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProcessIdentity:
|
|
pid: int
|
|
creation_time: str
|
|
creation_time_unix: float
|
|
executable: str
|
|
in_job: object
|
|
|
|
def as_dict(self):
|
|
return {
|
|
'pid': int(self.pid),
|
|
'creation_time': str(self.creation_time),
|
|
'creation_time_unix': float(self.creation_time_unix),
|
|
'executable': str(self.executable),
|
|
'in_job': self.in_job,
|
|
}
|
|
|
|
|
|
class RetainedProcess:
|
|
def __init__(self, identity, handle=None, pidfd=None):
|
|
self.identity = identity
|
|
self._handle = handle
|
|
self._pidfd = pidfd
|
|
self._closed = False
|
|
|
|
@property
|
|
def pid(self):
|
|
return self.identity.pid
|
|
|
|
def is_running(self):
|
|
if self._closed:
|
|
return False
|
|
if os.name == 'nt':
|
|
result = _WAIT_FOR_SINGLE_OBJECT(self._handle, 0)
|
|
if result == 258:
|
|
return True
|
|
if result == 0:
|
|
return False
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
try:
|
|
current = _posix_identity(self.pid)
|
|
return current.creation_time == self.identity.creation_time
|
|
except ProcessIdentityError:
|
|
return False
|
|
|
|
def wait(self, timeout):
|
|
timeout = max(0.0, float(timeout))
|
|
if os.name == 'nt':
|
|
milliseconds = min(int(timeout * 1000), 0xFFFFFFFE)
|
|
result = _WAIT_FOR_SINGLE_OBJECT(self._handle, milliseconds)
|
|
if result == 0:
|
|
return True
|
|
if result == 258:
|
|
return False
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if not self.is_running():
|
|
return True
|
|
time.sleep(min(0.05, max(0.0, deadline - time.monotonic())))
|
|
return not self.is_running()
|
|
|
|
def exit_code(self):
|
|
if self._closed:
|
|
raise ProcessIdentityError('retained process handle is closed')
|
|
if os.name != 'nt':
|
|
return None
|
|
code = wintypes.DWORD()
|
|
if not _GET_EXIT_CODE_PROCESS(self._handle, ctypes.byref(code)):
|
|
raise ProcessIdentityError(f'unable to read process exit status: {ctypes.WinError(ctypes.get_last_error())}')
|
|
if code.value == 259:
|
|
return None
|
|
return int(code.value)
|
|
|
|
def terminate(self):
|
|
if self._closed:
|
|
raise ProcessIdentityError('retained process handle is closed')
|
|
if os.name == 'nt':
|
|
if not _TERMINATE_PROCESS(self._handle, 1):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
return
|
|
sender = getattr(signal, 'pidfd_send_signal', None)
|
|
if self._pidfd is not None and sender is not None:
|
|
sender(self._pidfd, signal.SIGTERM, None, 0)
|
|
return
|
|
current = _posix_identity(self.pid)
|
|
if (
|
|
current.creation_time != self.identity.creation_time
|
|
or current.executable != self.identity.executable
|
|
):
|
|
raise ProcessIdentityError(f'process identity changed before signaling PID {self.pid}')
|
|
os.kill(self.pid, signal.SIGTERM)
|
|
|
|
def command_line(self):
|
|
if self._closed:
|
|
raise ProcessIdentityError('retained process handle is closed')
|
|
if os.name != 'nt':
|
|
try:
|
|
with open(f'/proc/{self.pid}/cmdline', 'rb') as handle:
|
|
return [item.decode(errors='surrogateescape') for item in handle.read().split(b'\0') if item]
|
|
except OSError as exc:
|
|
raise ProcessIdentityError(f'unable to read process {self.pid} command line') from exc
|
|
|
|
needed = wintypes.ULONG()
|
|
_NT_QUERY_INFORMATION_PROCESS(self._handle, 60, None, 0, ctypes.byref(needed))
|
|
if not needed.value:
|
|
raise ProcessIdentityError(f'unable to size process {self.pid} command line')
|
|
buffer = ctypes.create_string_buffer(needed.value)
|
|
status = _NT_QUERY_INFORMATION_PROCESS(
|
|
self._handle, 60, buffer, needed.value, ctypes.byref(needed),
|
|
)
|
|
if status < 0:
|
|
raise ProcessIdentityError(f'unable to read process {self.pid} command line (NTSTATUS 0x{status & 0xFFFFFFFF:08X})')
|
|
value = ctypes.cast(buffer, _P_UNICODE_STRING).contents
|
|
command = ctypes.wstring_at(value.Buffer, value.Length // ctypes.sizeof(ctypes.c_wchar))
|
|
argc = ctypes.c_int()
|
|
argv = _COMMAND_LINE_TO_ARGV(command, ctypes.byref(argc))
|
|
if not argv:
|
|
raise ProcessIdentityError(f'unable to parse process {self.pid} command line')
|
|
try:
|
|
return [argv[index] for index in range(argc.value)]
|
|
finally:
|
|
_LOCAL_FREE(argv)
|
|
|
|
def close(self):
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
if os.name == 'nt' and self._handle:
|
|
_CLOSE_HANDLE(self._handle)
|
|
elif self._pidfd is not None:
|
|
try:
|
|
os.close(self._pidfd)
|
|
except OSError:
|
|
pass
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, value, traceback):
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.close()
|
|
except BaseException:
|
|
pass
|
|
|
|
|
|
def _windows_identity(handle, pid):
|
|
creation = _FILETIME()
|
|
ignored_exit = _FILETIME()
|
|
ignored_kernel = _FILETIME()
|
|
ignored_user = _FILETIME()
|
|
if not _GET_PROCESS_TIMES(
|
|
handle, ctypes.byref(creation), ctypes.byref(ignored_exit),
|
|
ctypes.byref(ignored_kernel), ctypes.byref(ignored_user),
|
|
):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
filetime = (int(creation.dwHighDateTime) << 32) | int(creation.dwLowDateTime)
|
|
path_buffer = ctypes.create_unicode_buffer(32768)
|
|
path_size = wintypes.DWORD(len(path_buffer))
|
|
if not _QUERY_FULL_PROCESS_IMAGE_NAME(handle, 0, path_buffer, ctypes.byref(path_size)):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
in_job = wintypes.BOOL()
|
|
if not _IS_PROCESS_IN_JOB(handle, None, ctypes.byref(in_job)):
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
unix_time = (filetime - 116444736000000000) / 10000000.0
|
|
return ProcessIdentity(
|
|
pid=int(pid),
|
|
creation_time=f'windows-filetime:{filetime}',
|
|
creation_time_unix=unix_time,
|
|
executable=canonical_path(path_buffer.value),
|
|
in_job=bool(in_job.value),
|
|
)
|
|
|
|
|
|
def _posix_identity(pid):
|
|
stat_path = f'/proc/{int(pid)}/stat'
|
|
try:
|
|
with open(stat_path, 'r', encoding='ascii') as handle:
|
|
value = handle.read()
|
|
close_paren = value.rfind(')')
|
|
fields = value[close_paren + 2:].split()
|
|
start_ticks = int(fields[19])
|
|
executable = canonical_path(os.readlink(f'/proc/{int(pid)}/exe'))
|
|
clock_ticks = int(os.sysconf('SC_CLK_TCK'))
|
|
boot_time = None
|
|
with open('/proc/stat', 'r', encoding='ascii') as handle:
|
|
for line in handle:
|
|
if line.startswith('btime '):
|
|
boot_time = float(line.split()[1])
|
|
break
|
|
if boot_time is None:
|
|
raise ValueError('boot time unavailable')
|
|
except (OSError, ValueError, IndexError) as exc:
|
|
raise ProcessIdentityError(f'unable to inspect process {pid}') from exc
|
|
return ProcessIdentity(
|
|
pid=int(pid),
|
|
creation_time=f'proc-start-ticks:{start_ticks}',
|
|
creation_time_unix=boot_time + (start_ticks / float(clock_ticks)),
|
|
executable=executable,
|
|
in_job=False,
|
|
)
|
|
|
|
|
|
def _pidfd_live(pidfd):
|
|
poller = select.poll()
|
|
poller.register(pidfd, select.POLLIN)
|
|
return not bool(poller.poll(0))
|
|
|
|
|
|
def open_process(pid, *, terminate=False):
|
|
pid = int(pid)
|
|
if pid <= 0:
|
|
raise ProcessIdentityError(f'invalid process ID: {pid}')
|
|
if os.name == 'nt':
|
|
rights = 0x00100000 | 0x00001000
|
|
if terminate:
|
|
rights |= 0x00000001
|
|
handle = _OPEN_PROCESS(rights, False, pid)
|
|
if not handle:
|
|
native_error = ctypes.WinError(ctypes.get_last_error())
|
|
raise ProcessIdentityError(f'unable to open process {pid}: {native_error}') from native_error
|
|
try:
|
|
wait_result = _WAIT_FOR_SINGLE_OBJECT(handle, 0)
|
|
if wait_result == 0:
|
|
exit_code = wintypes.DWORD()
|
|
code = int(exit_code.value) if _GET_EXIT_CODE_PROCESS(
|
|
handle, ctypes.byref(exit_code),
|
|
) else -1
|
|
raise ProcessExitedError(
|
|
f'process {pid} has already exited with code {code}'
|
|
)
|
|
if wait_result != 258:
|
|
raise ProcessIdentityError(
|
|
f'unable to wait on process {pid}: {ctypes.WinError(ctypes.get_last_error())}'
|
|
)
|
|
exit_code = wintypes.DWORD()
|
|
if not _GET_EXIT_CODE_PROCESS(handle, ctypes.byref(exit_code)):
|
|
native_error = ctypes.WinError(ctypes.get_last_error())
|
|
raise ProcessIdentityError(
|
|
f'unable to read process {pid} exit status: {native_error}'
|
|
) from native_error
|
|
try:
|
|
identity = _windows_identity(handle, pid)
|
|
except OSError as exc:
|
|
retry_exit_code = wintypes.DWORD()
|
|
if (
|
|
_GET_EXIT_CODE_PROCESS(handle, ctypes.byref(retry_exit_code))
|
|
and retry_exit_code.value != 259
|
|
):
|
|
raise ProcessExitedError(
|
|
f'process {pid} exited during identity inspection '
|
|
f'with code {int(retry_exit_code.value)}'
|
|
) from exc
|
|
raise ProcessIdentityError(f'unable to inspect process {pid}') from exc
|
|
final_wait = _WAIT_FOR_SINGLE_OBJECT(handle, 0)
|
|
if final_wait == 0:
|
|
raise ProcessExitedError(
|
|
f'process {pid} exited during identity inspection'
|
|
)
|
|
if final_wait != 258:
|
|
raise ProcessIdentityError(
|
|
f'unable to confirm process {pid} liveness: '
|
|
f'{ctypes.WinError(ctypes.get_last_error())}'
|
|
)
|
|
return RetainedProcess(identity, handle=handle)
|
|
except BaseException:
|
|
_CLOSE_HANDLE(handle)
|
|
raise
|
|
pidfd = None
|
|
if hasattr(os, 'pidfd_open'):
|
|
try:
|
|
pidfd = os.pidfd_open(pid, 0)
|
|
except ProcessLookupError as exc:
|
|
raise ProcessExitedError(f'process {pid} has already exited') from exc
|
|
except OSError as exc:
|
|
raise ProcessIdentityError(
|
|
f'unable to pin process {pid} with pidfd',
|
|
) from exc
|
|
try:
|
|
if pidfd is not None and not _pidfd_live(pidfd):
|
|
raise ProcessExitedError(f'process {pid} exited before identity binding')
|
|
identity = _posix_identity(pid)
|
|
if pidfd is not None and not _pidfd_live(pidfd):
|
|
raise ProcessExitedError(f'process {pid} exited during identity binding')
|
|
verified = _posix_identity(pid)
|
|
if (
|
|
verified.creation_time != identity.creation_time
|
|
or verified.executable != identity.executable
|
|
):
|
|
raise ProcessIdentityError(
|
|
f'process {pid} identity changed during pidfd binding',
|
|
)
|
|
if pidfd is not None and not _pidfd_live(pidfd):
|
|
raise ProcessExitedError(f'process {pid} exited after identity binding')
|
|
return RetainedProcess(identity, pidfd=pidfd)
|
|
except BaseException:
|
|
if pidfd is not None:
|
|
os.close(pidfd)
|
|
raise
|
|
|
|
|
|
def current_process_identity():
|
|
if os.name == 'nt':
|
|
return _windows_identity(_GET_CURRENT_PROCESS(), os.getpid())
|
|
return _posix_identity(os.getpid())
|
|
|
|
|
|
def verify_retained_process(pid, creation_time, executable, *, terminate=False):
|
|
process = open_process(pid, terminate=True) if terminate else open_process(pid)
|
|
expected_executable = canonical_path(executable)
|
|
if process.identity.creation_time != str(creation_time) or process.identity.executable != expected_executable:
|
|
process.close()
|
|
raise ProcessIdentityError(f'process identity mismatch for PID {pid}')
|
|
return process
|
|
|
|
|
|
def serialize_process_identity(identity):
|
|
if isinstance(identity, ProcessIdentity):
|
|
return identity.as_dict()
|
|
raise TypeError('expected ProcessIdentity')
|
|
|
|
|
|
def exact_process_identity_state(pid, creation_time, executable):
|
|
"""Return alive, dead, reused, or unknown without PID-only inference."""
|
|
try:
|
|
pid = int(pid)
|
|
except (TypeError, ValueError):
|
|
return 'unknown'
|
|
if pid <= 0 or not creation_time or not executable:
|
|
return 'unknown'
|
|
try:
|
|
process = open_process(pid)
|
|
except ProcessExitedError:
|
|
return 'dead'
|
|
except ProcessIdentityError as exc:
|
|
cause = exc.__cause__
|
|
winerror = getattr(cause, 'winerror', None) or getattr(exc, 'winerror', None)
|
|
errno_value = getattr(cause, 'errno', None) or getattr(exc, 'errno', None)
|
|
if os.name == 'nt' and winerror in (87, 1168):
|
|
return 'dead'
|
|
if os.name != 'nt' and errno_value in (2, 3):
|
|
return 'dead'
|
|
return 'unknown'
|
|
try:
|
|
if not process.is_running():
|
|
return 'dead'
|
|
if (
|
|
str(process.identity.creation_time) != str(creation_time)
|
|
or canonical_path(process.identity.executable) != canonical_path(executable)
|
|
):
|
|
return 'reused'
|
|
return 'alive'
|
|
except (OSError, ValueError):
|
|
return 'unknown'
|
|
finally:
|
|
process.close()
|