Skip to content

Commit

Permalink
Clean up Windows implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
WardBrian committed May 14, 2024
1 parent fc200c5 commit 7a598cc
Showing 1 changed file with 57 additions and 90 deletions.
147 changes: 57 additions & 90 deletions dllist/windows.py
Original file line number Diff line number Diff line change
@@ -1,99 +1,66 @@
import ctypes
import warnings
from ctypes.wintypes import BOOL, DWORD, HANDLE, HMODULE, LPDWORD, LPWSTR
from typing import List, Optional, Sequence, Tuple

# https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-enumprocessmodulesex

LIST_MODULES_ALL = 3 # 0x03


def get_current_process() -> HANDLE:
_get_current_process = ctypes.windll.kernel32.GetCurrentProcess
_get_current_process.restype = HANDLE
return _get_current_process()


def get_module_filename(hModule: HMODULE) -> Optional[str]:
_get_module_filename = ctypes.windll.kernel32.GetModuleFileNameW
_get_module_filename.argtypes = [HMODULE, LPWSTR, DWORD]
_get_module_filename.restype = DWORD

nSize = 32768 # MAX_PATH
try:
lpFilename = ctypes.create_unicode_buffer(nSize)
if _get_module_filename(hModule, lpFilename, nSize) != 0:
return lpFilename.value
else:
warnings.warn(
f"Failed to get module file name for module {hModule}", stacklevel=2
)
except:
warnings.warn(
f"Failed to get module file name for module {hModule}", stacklevel=2
)
from ctypes import wintypes
from typing import List, Optional

# https://learn.microsoft.com/windows/win32/api/psapi/nf-psapi-enumprocessmodules
# https://learn.microsoft.com/windows/win32/api/libloaderapi/nf-libloaderapi-getmodulefilenamew

_kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
_psapi = ctypes.WinDLL('psapi', use_last_error=True)

_kernel32.GetCurrentProcess.restype = wintypes.HANDLE

_kernel32.GetModuleFileNameW.restype = wintypes.DWORD
_kernel32.GetModuleFileNameW.argtypes = (
wintypes.HMODULE,
wintypes.LPWSTR,
wintypes.DWORD,
)

_psapi.EnumProcessModules.restype = wintypes.BOOL
_psapi.EnumProcessModules.argtypes = (
wintypes.HANDLE,
ctypes.POINTER(wintypes.HMODULE),
wintypes.DWORD,
wintypes.LPDWORD,
)

def get_module_filename(hModule: wintypes.HMODULE) -> Optional[str]:
name = (wintypes.WCHAR * 32767)() # UNICODE_STRING_MAX_CHARS
if _kernel32.GetModuleFileNameW(hModule, name, len(name)):
return name.value
error = ctypes.get_last_error()
warnings.warn(f"Failed to get module file name for module {hModule}: "
f"GetModuleFileNameW failed with error code {error}",
stacklevel=2)
return None


def get_process_module_handles_partial(
hProcess: HANDLE, maxbuffsize: int
) -> Tuple[Sequence[HMODULE], int]:
_enumerate_loaded_modules = ctypes.windll.psapi.EnumProcessModulesEx
_enumerate_loaded_modules.argtypes = [
HANDLE,
ctypes.POINTER(HMODULE),
DWORD,
LPDWORD,
DWORD,
]
_enumerate_loaded_modules.restype = BOOL
cb = DWORD(maxbuffsize * ctypes.sizeof(HMODULE))
cbNeeded = DWORD(0)

hModules = (HMODULE * maxbuffsize)()
success = _enumerate_loaded_modules(
hProcess,
hModules,
cb,
ctypes.byref(cbNeeded),
LIST_MODULES_ALL,
)

if not success:
warnings.warn(
"Unable to list loaded libraries: "
f"EnumProcessModulesEx failed with error code {ctypes.GetLastError()}",
stacklevel=3,
)
return [], 0

bufsize_needed = cbNeeded.value // ctypes.sizeof(HMODULE)

return hModules, bufsize_needed


def get_process_module_handles() -> Sequence[HMODULE]:
hProcess = get_current_process()
first_attempt = 1024
hModules, buffer_needed = get_process_module_handles_partial(
hProcess, maxbuffsize=first_attempt
)
if buffer_needed > first_attempt:
# We need a bigger buffer, but luckily we know how big it needs to be
hModules, buffer_needed = get_process_module_handles_partial(
hProcess, maxbuffsize=buffer_needed
)

# skip first entry, which is the executable itself,
# and trim the list to the number of modules actually loaded
return hModules[1:buffer_needed]
def get_module_handles() -> List[int]:
hProcess = _kernel32.GetCurrentProcess()
cbNeeded = wintypes.DWORD()
n = 1024
while True:
modules = (wintypes.HMODULE * n)()
if not _psapi.EnumProcessModules(hProcess,
modules,
ctypes.sizeof(modules),
ctypes.byref(cbNeeded)):
break
n = cbNeeded.value // ctypes.sizeof(wintypes.HMODULE)
if n <= len(modules):
return modules[:n]
error = ctypes.get_last_error()
warnings.warn("Unable to list loaded libraries: EnumProcessModules "
f"failed with error code {error}",
stacklevel=2)
return []


def _platform_specific_dllist() -> List[str]:
hModules = get_process_module_handles()

libraries = [
name for hMod in hModules if (name := get_module_filename(hMod)) is not None
]

# skip first entry, which is the executable itself
modules = get_module_handles()[1:]
libraries = [name for h in modules
if (name := get_module_filename(h)) is not None]
return libraries

0 comments on commit 7a598cc

Please sign in to comment.