diff options
author | Giampaolo Rodola <g.rodola@gmail.com> | 2019-02-19 12:52:48 -0800 |
---|---|---|
committer | GitHub <noreply@github.com> | 2019-02-19 12:52:48 -0800 |
commit | 097b6e9ec4843acfc5c98c147d7aff418679f1db (patch) | |
tree | ded7f161bed3054aa5ea02feb8820f0b9b7dc270 | |
parent | 7381d4d15dd7fe108e1bf350c3f76035922c6ed4 (diff) | |
download | psutil-097b6e9ec4843acfc5c98c147d7aff418679f1db.tar.gz |
Windows / refactoring: utility functions for LoadLibraryA and GetProcAddress (#1417)
Windows / refactoring: add utility functions for LoadLibraryA and GetProcAddress. Centralize logic in one place.
-rw-r--r-- | make.bat | 7 | ||||
-rw-r--r-- | psutil/_psutil_windows.c | 157 | ||||
-rw-r--r-- | psutil/arch/windows/inet_ntop.h | 2 | ||||
-rw-r--r-- | psutil/arch/windows/process_handles.c | 14 | ||||
-rw-r--r-- | psutil/arch/windows/process_info.c | 148 | ||||
-rw-r--r-- | psutil/arch/windows/process_info.h | 6 |
6 files changed, 121 insertions, 213 deletions
@@ -20,8 +20,13 @@ rem set PYTHON=C:\Python34\python.exe & set TSCRIPT=foo.py & make.bat test rem ========================================================================== if "%PYTHON%" == "" ( - set PYTHON=C:\Python27\python.exe + if exist "C:\Python37\python.exe" ( + set PYTHON=C:\Python37\python.exe + ) else ( + set PYTHON=C:\Python27\python.exe + ) ) + if "%TSCRIPT%" == "" ( set TSCRIPT=psutil\tests\__main__.py ) diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c index b5b5a8f5..55a1cb22 100644 --- a/psutil/_psutil_windows.c +++ b/psutil/_psutil_windows.c @@ -205,15 +205,12 @@ psutil_get_num_cpus(int fail_on_err) { unsigned int ncpus = 0; SYSTEM_INFO sysinfo; static DWORD(CALLBACK *_GetActiveProcessorCount)(WORD) = NULL; - HINSTANCE hKernel32; // GetActiveProcessorCount is available only on 64 bit versions // of Windows from Windows 7 onward. - // Windows Vista 64 bit and Windows XP doesn't have it. - hKernel32 = GetModuleHandleW(L"KERNEL32"); - _GetActiveProcessorCount = (void*)GetProcAddress( - hKernel32, "GetActiveProcessorCount"); - + // Windows Vista 64 bit and Windows XP don't have it. + _GetActiveProcessorCount = \ + psutil_GetProcAddress("kernel32", "GetActiveProcessorCount"); if (_GetActiveProcessorCount != NULL) { ncpus = _GetActiveProcessorCount(ALL_PROCESSOR_GROUPS); if ((ncpus == 0) && (fail_on_err == 1)) { @@ -261,11 +258,9 @@ psutil_boot_time(PyObject *self, PyObject *args) { time_t pt; FILETIME fileTime; long long ll; - HINSTANCE hKernel32; - psutil_GetTickCount64 = NULL; + psutil_GetTickCount64; GetSystemTimeAsFileTime(&fileTime); - /* HUGE thanks to: http://johnstewien.spaces.live.com/blog/cns!E6885DB5CEBABBC8!831.entry @@ -288,12 +283,11 @@ psutil_boot_time(PyObject *self, PyObject *args) { pt = (time_t)((ll - 116444736000000000ull) / 10000000ull); // GetTickCount64() is Windows Vista+ only. Dinamically load - // GetTickCount64() at runtime. We may have used + // it at runtime. We may have used // "#if (_WIN32_WINNT >= 0x0600)" pre-processor but that way // the produced exe/wheels cannot be used on Windows XP, see: // https://github.com/giampaolo/psutil/issues/811#issuecomment-230639178 - hKernel32 = GetModuleHandleW(L"KERNEL32"); - psutil_GetTickCount64 = (void*)GetProcAddress(hKernel32, "GetTickCount64"); + psutil_GetTickCount64 = psutil_GetProcAddress("kernel32", "GetTickCount64"); if (psutil_GetTickCount64 != NULL) { // Windows >= Vista uptime = psutil_GetTickCount64() / (ULONGLONG)1000.00f; @@ -642,14 +636,10 @@ psutil_cpu_count_phys(PyObject *self, PyObject *args) { // it supports process groups, meaning this is able to report more // than 64 CPUs. See: // https://bugs.python.org/issue33166 - _GetLogicalProcessorInformationEx = \ - (PFN_GETLOGICALPROCESSORINFORMATIONEX)GetProcAddress( - GetModuleHandle(TEXT("kernel32")), - "GetLogicalProcessorInformationEx"); - if (_GetLogicalProcessorInformationEx == NULL) { - psutil_debug("failed loading GetLogicalProcessorInformationEx()"); - goto return_none; - } + _GetLogicalProcessorInformationEx = psutil_GetProcAddressFromLib( + "kernel32", "GetLogicalProcessorInformationEx"); + if (_GetLogicalProcessorInformationEx == NULL) + return NULL; while (1) { rc = _GetLogicalProcessorInformationEx( @@ -765,7 +755,7 @@ psutil_proc_exe(PyObject *self, PyObject *args) { HANDLE hProcess; wchar_t exe[MAX_PATH]; #if (_WIN32_WINNT >= 0x0600) // >= Vista - PDWORD size = MAX_PATH; + unsigned int size = sizeof(exe); #endif if (! PyArg_ParseTuple(args, "l", &pid)) @@ -1057,7 +1047,6 @@ psutil_per_cpu_times(PyObject *self, PyObject *args) { // NtQuerySystemInformation stuff typedef DWORD (_stdcall * NTQSI_PROC) (int, PVOID, ULONG, PULONG); NTQSI_PROC NtQuerySystemInformation; - HINSTANCE hNtDll; double idle, kernel, systemt, user, interrupt, dpc; NTSTATUS status; @@ -1069,19 +1058,10 @@ psutil_per_cpu_times(PyObject *self, PyObject *args) { if (py_retlist == NULL) return NULL; - - // obtain NtQuerySystemInformation - hNtDll = LoadLibrary(TEXT("ntdll.dll")); - if (hNtDll == NULL) { - PyErr_SetFromWindowsErr(0); + NtQuerySystemInformation = \ + psutil_GetProcAddressFromLib("ntdll.dll", "NtQuerySystemInformation"); + if (NtQuerySystemInformation == NULL) goto error; - } - NtQuerySystemInformation = (NTQSI_PROC)GetProcAddress( - hNtDll, "NtQuerySystemInformation"); - if (NtQuerySystemInformation == NULL) { - PyErr_SetFromWindowsErr(0); - goto error; - } // retrieves number of processors ncpus = psutil_get_num_cpus(1); @@ -1144,7 +1124,6 @@ psutil_per_cpu_times(PyObject *self, PyObject *args) { } free(sppi); - FreeLibrary(hNtDll); return py_retlist; error: @@ -1152,8 +1131,6 @@ error: Py_DECREF(py_retlist); if (sppi) free(sppi); - if (hNtDll) - FreeLibrary(hNtDll); return NULL; } @@ -1672,11 +1649,26 @@ psutil_net_connections(PyObject *self, PyObject *args) { PyObject *_SOCK_STREAM = PyLong_FromLong((long)SOCK_STREAM); PyObject *_SOCK_DGRAM = PyLong_FromLong((long)SOCK_DGRAM); + // Import some functions. + rtlIpv4AddressToStringA = psutil_GetProcAddressFromLib( + "ntdll.dll", "RtlIpv4AddressToStringA"); + if (rtlIpv4AddressToStringA == NULL) + goto error; + rtlIpv6AddressToStringA = psutil_GetProcAddressFromLib( + "ntdll.dll", "RtlIpv6AddressToStringA"); + if (rtlIpv6AddressToStringA == NULL) + goto error; + getExtendedTcpTable = psutil_GetProcAddressFromLib( + "iphlpapi.dll", "GetExtendedTcpTable"); + if (getExtendedTcpTable == NULL) + goto error; + getExtendedUdpTable = psutil_GetProcAddressFromLib( + "iphlpapi.dll", "GetExtendedUdpTable"); + if (getExtendedUdpTable == NULL) + goto error; + if (! PyArg_ParseTuple(args, "lOO", &pid, &py_af_filter, &py_type_filter)) - { - _psutil_conn_decref_objs(); - return NULL; - } + goto error; if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) { _psutil_conn_decref_objs(); @@ -1696,27 +1688,6 @@ psutil_net_connections(PyObject *self, PyObject *args) { } } - // Import some functions. - { - HMODULE ntdll; - HMODULE iphlpapi; - - ntdll = LoadLibrary(TEXT("ntdll.dll")); - rtlIpv4AddressToStringA = (_RtlIpv4AddressToStringA)GetProcAddress( - ntdll, "RtlIpv4AddressToStringA"); - rtlIpv6AddressToStringA = (_RtlIpv6AddressToStringA)GetProcAddress( - ntdll, "RtlIpv6AddressToStringA"); - /* TODO: Check these two function pointers */ - - iphlpapi = LoadLibrary(TEXT("iphlpapi.dll")); - getExtendedTcpTable = (_GetExtendedTcpTable)GetProcAddress(iphlpapi, - "GetExtendedTcpTable"); - getExtendedUdpTable = (_GetExtendedUdpTable)GetProcAddress(iphlpapi, - "GetExtendedUdpTable"); - FreeLibrary(ntdll); - FreeLibrary(iphlpapi); - } - if ((getExtendedTcpTable == NULL) || (getExtendedUdpTable == NULL)) { PyErr_SetString(PyExc_NotImplementedError, "feature not supported on this Windows version"); @@ -2152,11 +2123,12 @@ psutil_proc_io_priority_get(PyObject *self, PyObject *args) { long pid; HANDLE hProcess; DWORD IoPriority; + _NtQueryInformationProcess NtQueryInformationProcess; - _NtQueryInformationProcess NtQueryInformationProcess = - (_NtQueryInformationProcess)GetProcAddress( - GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess"); - + NtQueryInformationProcess = \ + psutil_GetProcAddress("ntdll.dll", "NtQueryInformationProcess"); + if (NtQueryInformationProcess == NULL) + return NULL; if (! PyArg_ParseTuple(args, "l", &pid)) return NULL; hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION); @@ -2184,17 +2156,12 @@ psutil_proc_io_priority_set(PyObject *self, PyObject *args) { DWORD prio; HANDLE hProcess; DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION; + _NtSetInformationProcess NtSetInformationProcess; - _NtSetInformationProcess NtSetInformationProcess = - (_NtSetInformationProcess)GetProcAddress( - GetModuleHandleA("ntdll.dll"), "NtSetInformationProcess"); - - if (NtSetInformationProcess == NULL) { - PyErr_SetString(PyExc_RuntimeError, - "couldn't get NtSetInformationProcess syscall"); + NtSetInformationProcess = \ + psutil_GetProcAddress("ntdll.dll", "NtSetInformationProcess"); + if (NtSetInformationProcess == NULL) return NULL; - } - if (! PyArg_ParseTuple(args, "li", &pid, &prio)) return NULL; hProcess = psutil_handle_from_pid(pid, access); @@ -2797,23 +2764,21 @@ psutil_users(PyObject *self, PyObject *args) { PWTS_CLIENT_ADDRESS address; char address_str[50]; long long unix_time; - PWINSTATIONQUERYINFORMATIONW WinStationQueryInformationW; WINSTATION_INFO station_info; - HINSTANCE hInstWinSta = NULL; ULONG returnLen; - - PyObject *py_retlist = PyList_New(0); PyObject *py_tuple = NULL; PyObject *py_address = NULL; PyObject *py_username = NULL; + PyObject *py_retlist = PyList_New(0); if (py_retlist == NULL) return NULL; - hInstWinSta = LoadLibraryA("winsta.dll"); - WinStationQueryInformationW = (PWINSTATIONQUERYINFORMATIONW) \ - GetProcAddress(hInstWinSta, "WinStationQueryInformationW"); + WinStationQueryInformationW = psutil_GetProcAddressFromLib( + "winsta.dll", "WinStationQueryInformationW"); + if (WinStationQueryInformationW == NULL) + goto error; if (WTSEnumerateSessions(hServer, 0, 1, &sessions, &count) == 0) { PyErr_SetFromWindowsErr(0); @@ -2902,7 +2867,6 @@ psutil_users(PyObject *self, PyObject *args) { WTSFreeMemory(sessions); WTSFreeMemory(buffer_user); WTSFreeMemory(buffer_addr); - FreeLibrary(hInstWinSta); return py_retlist; error: @@ -2911,8 +2875,6 @@ error: Py_XDECREF(py_address); Py_DECREF(py_retlist); - if (hInstWinSta != NULL) - FreeLibrary(hInstWinSta); if (sessions != NULL) WTSFreeMemory(sessions); if (buffer_user != NULL) @@ -3100,7 +3062,6 @@ psutil_proc_memory_maps(PyObject *self, PyObject *args) { GetSystemInfo(&system_info); maxAddr = system_info.lpMaximumApplicationAddress; baseAddress = NULL; - previousAllocationBase = NULL; while (VirtualQueryEx(hProcess, baseAddress, &basicInfo, sizeof(MEMORY_BASIC_INFORMATION))) @@ -3135,7 +3096,7 @@ psutil_proc_memory_maps(PyObject *self, PyObject *args) { Py_DECREF(py_tuple); Py_DECREF(py_str); } - previousAllocationBase = basicInfo.AllocationBase; + previousAllocationBase = (ULONGLONG)basicInfo.AllocationBase; baseAddress = (PCHAR)baseAddress + basicInfo.RegionSize; } @@ -3517,11 +3478,8 @@ error: */ static PyObject * psutil_cpu_stats(PyObject *self, PyObject *args) { - // NtQuerySystemInformation stuff typedef DWORD (_stdcall * NTQSI_PROC) (int, PVOID, ULONG, PULONG); NTQSI_PROC NtQuerySystemInformation; - HINSTANCE hNtDll; - NTSTATUS status; _SYSTEM_PERFORMANCE_INFORMATION *spi = NULL; _SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION *sppi = NULL; @@ -3531,18 +3489,10 @@ psutil_cpu_stats(PyObject *self, PyObject *args) { ULONG64 dpcs = 0; ULONG interrupts = 0; - // obtain NtQuerySystemInformation - hNtDll = LoadLibrary(TEXT("ntdll.dll")); - if (hNtDll == NULL) { - PyErr_SetFromWindowsErr(0); - goto error; - } - NtQuerySystemInformation = (NTQSI_PROC)GetProcAddress( - hNtDll, "NtQuerySystemInformation"); - if (NtQuerySystemInformation == NULL) { - PyErr_SetFromWindowsErr(0); - goto error; - } + NtQuerySystemInformation = \ + psutil_GetProcAddressFromLib("ntdll.dll", "NtQuerySystemInformation"); + if (NtQuerySystemInformation == NULL) + return NULL; // retrieves number of processors ncpus = psutil_get_num_cpus(1); @@ -3613,7 +3563,6 @@ psutil_cpu_stats(PyObject *self, PyObject *args) { free(spi); free(InterruptInformation); free(sppi); - FreeLibrary(hNtDll); return Py_BuildValue( "kkkk", spi->ContextSwitches, @@ -3629,8 +3578,6 @@ error: free(InterruptInformation); if (sppi) free(sppi); - if (hNtDll) - FreeLibrary(hNtDll); return NULL; } diff --git a/psutil/arch/windows/inet_ntop.h b/psutil/arch/windows/inet_ntop.h index 70573a36..0dbf28bf 100644 --- a/psutil/arch/windows/inet_ntop.h +++ b/psutil/arch/windows/inet_ntop.h @@ -9,7 +9,7 @@ PCSTR WSAAPI inet_ntop( __in INT Family, - __in PVOID pAddr, + __in const VOID * pAddr, __out_ecount(StringBufSize) PSTR pStringBuf, __in size_t StringBufSize ); diff --git a/psutil/arch/windows/process_handles.c b/psutil/arch/windows/process_handles.c index 356e2368..4ac02b91 100644 --- a/psutil/arch/windows/process_handles.c +++ b/psutil/arch/windows/process_handles.c @@ -5,6 +5,7 @@ * */ #include "process_handles.h" +#include "process_info.h" #include "../../_psutil_common.h" static _NtQuerySystemInformation __NtQuerySystemInformation = NULL; @@ -22,12 +23,6 @@ ULONG g_dwSize = 0; ULONG g_dwLength = 0; -PVOID -GetLibraryProcAddress(PSTR LibraryName, PSTR ProcName) { - return GetProcAddress(GetModuleHandleA(LibraryName), ProcName); -} - - PyObject * psutil_get_open_files(long dwPid, HANDLE hProcess) { OSVERSIONINFO osvi; @@ -50,9 +45,10 @@ psutil_get_open_files_init(BOOL threaded) { return; // Resolve the Windows API calls - __NtQuerySystemInformation = - GetLibraryProcAddress("ntdll.dll", "NtQuerySystemInformation"); - __NtQueryObject = GetLibraryProcAddress("ntdll.dll", "NtQueryObject"); + __NtQuerySystemInformation = psutil_GetProcAddressFromLib( + "ntdll.dll", "NtQuerySystemInformation"); + __NtQueryObject = psutil_GetProcAddressFromLib( + "ntdll.dll", "NtQueryObject"); // Create events for signalling work between threads if (threaded == TRUE) { diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c index 6f877342..9001ab14 100644 --- a/psutil/arch/windows/process_info.c +++ b/psutil/arch/windows/process_info.c @@ -161,70 +161,42 @@ typedef struct { const int STATUS_INFO_LENGTH_MISMATCH = 0xC0000004; const int STATUS_BUFFER_TOO_SMALL = 0xC0000023L; -#define WINDOWS_UNINITIALIZED 0 -#define WINDOWS_XP 51 -#define WINDOWS_VISTA 60 -#define WINDOWS_7 61 -#define WINDOWS_8 62 -#define WINDOWS_81 63 -#define WINDOWS_10 100 - - -int get_windows_version() { - OSVERSIONINFO ver_info; - BOOL result; - DWORD dwMajorVersion; - DWORD dwMinorVersion; - DWORD dwBuildNumber; - static int windows_version = WINDOWS_UNINITIALIZED; - // windows_version is static - // and equal to WINDOWS_UNINITIALIZED only on first call - if (windows_version == WINDOWS_UNINITIALIZED) { - memset(&ver_info, 0, sizeof(ver_info)); - ver_info.dwOSVersionInfoSize = sizeof(OSVERSIONINFO); - result = GetVersionEx(&ver_info); - if (result != FALSE) { - dwMajorVersion = ver_info.dwMajorVersion; - dwMinorVersion = ver_info.dwMinorVersion; - dwBuildNumber = ver_info.dwBuildNumber; - // Windows XP, Windows server 2003 - if (dwMajorVersion == 5 && dwMinorVersion == 1) { - windows_version = WINDOWS_XP; - } - // Windows Vista - else if (dwMajorVersion == 6 && dwMinorVersion == 0) { - windows_version = WINDOWS_VISTA; - } - // Windows 7, Windows Server 2008 R2 - else if (dwMajorVersion == 6 && dwMinorVersion == 1) { - windows_version = WINDOWS_7; - } - // Windows 8, Windows Server 2012 - else if (dwMajorVersion == 6 && dwMinorVersion == 2) { - windows_version = WINDOWS_8; - } - // Windows 8.1, Windows Server 2012 R2 - else if (dwMajorVersion == 6 && dwMinorVersion == 3) - { - windows_version = WINDOWS_81; - } - // Windows 10, Windows Server 2016 - else if (dwMajorVersion == 10) { - windows_version = WINDOWS_10; - } - } + +// A wrapper around GetModuleHandle and GetProcAddress. +PVOID +psutil_GetProcAddress(LPCSTR libname, LPCSTR procname) { + HMODULE mod; + FARPROC addr; + + if ((mod = GetModuleHandleA(libname)) == NULL) { + PyErr_SetFromWindowsErrWithFilename(0, libname); + return NULL; + } + if ((addr = GetProcAddress(mod, procname)) == NULL) { + PyErr_SetFromWindowsErrWithFilename(0, procname); + return NULL; } - return windows_version; + return addr; } -_NtQueryInformationProcess psutil_NtQueryInformationProcess() { - static _NtQueryInformationProcess NtQueryInformationProcess = NULL; - if (NtQueryInformationProcess == NULL) { - NtQueryInformationProcess = (_NtQueryInformationProcess)GetProcAddress( - GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess"); +// A wrapper around LoadLibrary and GetProcAddress. +PVOID +psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname) { + HMODULE mod; + FARPROC addr; + + if ((mod = LoadLibraryA(libname)) == NULL) { + PyErr_SetFromWindowsErrWithFilename(0, libname); + return NULL; + } + if ((addr = GetProcAddress(mod, procname)) == NULL) { + PyErr_SetFromWindowsErrWithFilename(0, procname); + FreeLibrary(mod); + return NULL; } - return NtQueryInformationProcess; + FreeLibrary(mod); + return addr; } @@ -523,11 +495,8 @@ psutil_get_process_region_size64(HANDLE hProcess, MEMORY_BASIC_INFORMATION64 info64; if (NtWow64QueryVirtualMemory64 == NULL) { - NtWow64QueryVirtualMemory64 = - (_NtWow64QueryVirtualMemory64)GetProcAddress( - GetModuleHandleA("ntdll.dll"), - "NtWow64QueryVirtualMemory64"); - + NtWow64QueryVirtualMemory64 = psutil_GetProcAddressFromLib( + "ntdll.dll", "NtWow64QueryVirtualMemory64"); if (NtWow64QueryVirtualMemory64 == NULL) { PyErr_SetString(PyExc_NotImplementedError, "NtWow64QueryVirtualMemory64 missing"); @@ -611,12 +580,15 @@ psutil_get_process_data(long pid, #endif DWORD access = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ; + NtQueryInformationProcess = \ + psutil_GetProcAddress("ntdll.dll", "NtQueryInformationProcess"); + if (NtQueryInformationProcess == NULL) + return -1; + hProcess = psutil_handle_from_pid(pid, access); if (hProcess == NULL) return -1; - NtQueryInformationProcess = psutil_NtQueryInformationProcess(); - #ifdef _WIN64 /* 64 bit case. Check if the target is a 32 bit process running in WoW64 * mode. */ @@ -679,11 +651,9 @@ psutil_get_process_data(long pid, RTL_USER_PROCESS_PARAMETERS64 procParameters64; if (NtWow64QueryInformationProcess64 == NULL) { - NtWow64QueryInformationProcess64 = - (_NtQueryInformationProcess)GetProcAddress( - GetModuleHandleA("ntdll.dll"), - "NtWow64QueryInformationProcess64"); - + NtWow64QueryInformationProcess64 = \ + psutil_GetProcAddressFromLib( + "ntdll.dll", "NtWow64QueryInformationProcess64"); if (NtWow64QueryInformationProcess64 == NULL) { PyErr_SetString(PyExc_NotImplementedError, "NtWow64QueryInformationProcess64 missing"); @@ -703,11 +673,9 @@ psutil_get_process_data(long pid, // read peb if (NtWow64ReadVirtualMemory64 == NULL) { - NtWow64ReadVirtualMemory64 = - (_NtWow64ReadVirtualMemory64)GetProcAddress( - GetModuleHandleA("ntdll.dll"), - "NtWow64ReadVirtualMemory64"); - + NtWow64ReadVirtualMemory64 = \ + psutil_GetProcAddressFromLib( + "ntdll.dll", "NtWow64ReadVirtualMemory64"); if (NtWow64ReadVirtualMemory64 == NULL) { PyErr_SetString(PyExc_NotImplementedError, "NtWow64ReadVirtualMemory64 missing"); @@ -870,13 +838,12 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) { WCHAR * cmdline_buffer_wchar = NULL; PUNICODE_STRING tmp = NULL; DWORD string_size; - _NtQueryInformationProcess NtQueryInformationProcess = NULL; + _NtQueryInformationProcess NtQueryInformationProcess; - NtQueryInformationProcess = psutil_NtQueryInformationProcess(); - if (NtQueryInformationProcess == NULL) { - PyErr_SetFromWindowsErr(0); + NtQueryInformationProcess = \ + psutil_GetProcAddress("ntdll.dll", "NtQueryInformationProcess"); + if (NtQueryInformationProcess == NULL) goto error; - } cmdline_buffer = calloc(ret_length, 1); if (cmdline_buffer == NULL) { @@ -885,10 +852,8 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) { } hProcess = psutil_handle_from_pid(pid, PROCESS_QUERY_LIMITED_INFORMATION); - if (hProcess == NULL) { - PyErr_SetFromWindowsErr(0); + if (hProcess == NULL) goto error; - } status = NtQueryInformationProcess( hProcess, 60, // ProcessCommandLineInformation @@ -939,10 +904,8 @@ psutil_get_cmdline(long pid, int use_peb) { PyObject *py_unicode = NULL; LPWSTR *szArglist = NULL; int nArgs, i; - int windows_version; int func_ret; - windows_version = get_windows_version(); /* By defaut, still use PEB (if command line params have been patched in the PEB, we will get the actual ones). Reading the PEB to get the @@ -1058,14 +1021,13 @@ psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess, PVOID buffer; ULONG bufferSize; PSYSTEM_PROCESS_INFORMATION process; - - // get NtQuerySystemInformation typedef DWORD (_stdcall * NTQSI_PROC) (int, PVOID, ULONG, PULONG); NTQSI_PROC NtQuerySystemInformation; - HINSTANCE hNtDll; - hNtDll = LoadLibrary(TEXT("ntdll.dll")); - NtQuerySystemInformation = (NTQSI_PROC)GetProcAddress( - hNtDll, "NtQuerySystemInformation"); + + NtQuerySystemInformation = \ + psutil_GetProcAddressFromLib("ntdll.dll", "NtQuerySystemInformation"); + if (NtQuerySystemInformation == NULL) + goto error; bufferSize = initialBufferSize; buffer = malloc(bufferSize); @@ -1077,7 +1039,6 @@ psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess, while (TRUE) { status = NtQuerySystemInformation(SystemProcessInformation, buffer, bufferSize, &bufferSize); - if (status == STATUS_BUFFER_TOO_SMALL || status == STATUS_INFO_LENGTH_MISMATCH) { @@ -1115,7 +1076,6 @@ psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess, goto error; error: - FreeLibrary(hNtDll); if (buffer != NULL) free(buffer); return 0; diff --git a/psutil/arch/windows/process_info.h b/psutil/arch/windows/process_info.h index f85c1efd..d2e60b7c 100644 --- a/psutil/arch/windows/process_info.h +++ b/psutil/arch/windows/process_info.h @@ -15,7 +15,6 @@ #define HANDLE_TO_PYNUM(handle) PyLong_FromUnsignedLong((unsigned long) handle) #define PYNUM_TO_HANDLE(obj) ((HANDLE)PyLong_AsUnsignedLong(obj)) - DWORD* psutil_get_pids(DWORD *numberOfReturnedPIDs); HANDLE psutil_handle_from_pid(DWORD pid, DWORD dwDesiredAccess); int psutil_pid_is_running(DWORD pid); @@ -24,9 +23,10 @@ int psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess, int psutil_assert_pid_exists(DWORD pid, char *err); int psutil_assert_pid_not_exists(DWORD pid, char *err); +PVOID psutil_GetProcAddress(LPCSTR libname, LPCSTR procname); +PVOID psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname); - -PyObject* psutil_get_cmdline(long pid); +PyObject* psutil_get_cmdline(long pid, int use_peb); PyObject* psutil_get_cwd(long pid); PyObject* psutil_get_environ(long pid); |