164 lines
4.4 KiB
Python
164 lines
4.4 KiB
Python
import contextlib
|
|
import ctypes
|
|
import os
|
|
|
|
from ctypes.wintypes import (
|
|
BOOL,
|
|
CHAR,
|
|
DWORD,
|
|
HANDLE,
|
|
LONG,
|
|
LPWSTR,
|
|
MAX_PATH,
|
|
PDWORD,
|
|
ULONG,
|
|
)
|
|
|
|
from shellingham._core import SHELL_NAMES
|
|
|
|
|
|
INVALID_HANDLE_VALUE = HANDLE(-1).value
|
|
ERROR_NO_MORE_FILES = 18
|
|
ERROR_INSUFFICIENT_BUFFER = 122
|
|
TH32CS_SNAPPROCESS = 2
|
|
PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
|
|
|
|
|
|
kernel32 = ctypes.windll.kernel32
|
|
|
|
|
|
def _check_handle(error_val=0):
|
|
def check(ret, func, args):
|
|
if ret == error_val:
|
|
raise ctypes.WinError()
|
|
return ret
|
|
|
|
return check
|
|
|
|
|
|
def _check_expected(expected):
|
|
def check(ret, func, args):
|
|
if ret:
|
|
return True
|
|
code = ctypes.GetLastError()
|
|
if code == expected:
|
|
return False
|
|
raise ctypes.WinError(code)
|
|
|
|
return check
|
|
|
|
|
|
class ProcessEntry32(ctypes.Structure):
|
|
_fields_ = (
|
|
("dwSize", DWORD),
|
|
("cntUsage", DWORD),
|
|
("th32ProcessID", DWORD),
|
|
("th32DefaultHeapID", ctypes.POINTER(ULONG)),
|
|
("th32ModuleID", DWORD),
|
|
("cntThreads", DWORD),
|
|
("th32ParentProcessID", DWORD),
|
|
("pcPriClassBase", LONG),
|
|
("dwFlags", DWORD),
|
|
("szExeFile", CHAR * MAX_PATH),
|
|
)
|
|
|
|
|
|
kernel32.CloseHandle.argtypes = [HANDLE]
|
|
kernel32.CloseHandle.restype = BOOL
|
|
|
|
kernel32.CreateToolhelp32Snapshot.argtypes = [DWORD, DWORD]
|
|
kernel32.CreateToolhelp32Snapshot.restype = HANDLE
|
|
kernel32.CreateToolhelp32Snapshot.errcheck = _check_handle( # type: ignore
|
|
INVALID_HANDLE_VALUE,
|
|
)
|
|
|
|
kernel32.Process32First.argtypes = [HANDLE, ctypes.POINTER(ProcessEntry32)]
|
|
kernel32.Process32First.restype = BOOL
|
|
kernel32.Process32First.errcheck = _check_expected( # type: ignore
|
|
ERROR_NO_MORE_FILES,
|
|
)
|
|
|
|
kernel32.Process32Next.argtypes = [HANDLE, ctypes.POINTER(ProcessEntry32)]
|
|
kernel32.Process32Next.restype = BOOL
|
|
kernel32.Process32Next.errcheck = _check_expected( # type: ignore
|
|
ERROR_NO_MORE_FILES,
|
|
)
|
|
|
|
kernel32.GetCurrentProcessId.argtypes = []
|
|
kernel32.GetCurrentProcessId.restype = DWORD
|
|
|
|
kernel32.OpenProcess.argtypes = [DWORD, BOOL, DWORD]
|
|
kernel32.OpenProcess.restype = HANDLE
|
|
kernel32.OpenProcess.errcheck = _check_handle( # type: ignore
|
|
INVALID_HANDLE_VALUE,
|
|
)
|
|
|
|
kernel32.QueryFullProcessImageNameW.argtypes = [HANDLE, DWORD, LPWSTR, PDWORD]
|
|
kernel32.QueryFullProcessImageNameW.restype = BOOL
|
|
kernel32.QueryFullProcessImageNameW.errcheck = _check_expected( # type: ignore
|
|
ERROR_INSUFFICIENT_BUFFER,
|
|
)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _handle(f, *args, **kwargs):
|
|
handle = f(*args, **kwargs)
|
|
try:
|
|
yield handle
|
|
finally:
|
|
kernel32.CloseHandle(handle)
|
|
|
|
|
|
def _iter_processes():
|
|
f = kernel32.CreateToolhelp32Snapshot
|
|
with _handle(f, TH32CS_SNAPPROCESS, 0) as snap:
|
|
entry = ProcessEntry32()
|
|
entry.dwSize = ctypes.sizeof(entry)
|
|
ret = kernel32.Process32First(snap, entry)
|
|
while ret:
|
|
yield entry
|
|
ret = kernel32.Process32Next(snap, entry)
|
|
|
|
|
|
def _get_full_path(proch):
|
|
size = DWORD(MAX_PATH)
|
|
while True:
|
|
path_buff = ctypes.create_unicode_buffer("", size.value)
|
|
if kernel32.QueryFullProcessImageNameW(proch, 0, path_buff, size):
|
|
return path_buff.value
|
|
size.value *= 2
|
|
|
|
|
|
def get_shell(pid=None, max_depth=10):
|
|
proc_map = {
|
|
proc.th32ProcessID: (proc.th32ParentProcessID, proc.szExeFile)
|
|
for proc in _iter_processes()
|
|
}
|
|
pid = pid or os.getpid()
|
|
|
|
for _ in range(0, max_depth + 1):
|
|
try:
|
|
ppid, executable = proc_map[pid]
|
|
except KeyError: # No such process? Give up.
|
|
break
|
|
|
|
# The executable name would be encoded with the current code page if
|
|
# we're in ANSI mode (usually). Try to decode it into str/unicode,
|
|
# replacing invalid characters to be safe (not thoeratically necessary,
|
|
# I think). Note that we need to use 'mbcs' instead of encoding
|
|
# settings from sys because this is from the Windows API, not Python
|
|
# internals (which those settings reflect). (pypa/pipenv#3382)
|
|
if isinstance(executable, bytes):
|
|
executable = executable.decode("mbcs", "replace")
|
|
|
|
name = executable.rpartition(".")[0].lower()
|
|
if name not in SHELL_NAMES:
|
|
pid = ppid
|
|
continue
|
|
|
|
key = PROCESS_QUERY_LIMITED_INFORMATION
|
|
with _handle(kernel32.OpenProcess, key, 0, pid) as proch:
|
|
return (name, _get_full_path(proch))
|
|
|
|
return None
|