Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,449 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user