1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4# Copyright (c) 2009-2014, Mario Vilas
5# All rights reserved.
6#
7# Redistribution and use in source and binary forms, with or without
8# modification, are permitted provided that the following conditions are met:
9#
10#     * Redistributions of source code must retain the above copyright notice,
11#       this list of conditions and the following disclaimer.
12#     * Redistributions in binary form must reproduce the above copyright
13#       notice,this list of conditions and the following disclaimer in the
14#       documentation and/or other materials provided with the distribution.
15#     * Neither the name of the copyright holder nor the names of its
16#       contributors may be used to endorse or promote products derived from
17#       this software without specific prior written permission.
18#
19# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
23# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29# POSSIBILITY OF SUCH DAMAGE.
30
31"""
32Wrapper for psapi.dll in ctypes.
33"""
34
35__revision__ = "$Id$"
36
37from winappdbg.win32.defines import *
38
39#==============================================================================
40# This is used later on to calculate the list of exported symbols.
41_all = None
42_all = set(vars().keys())
43#==============================================================================
44
45#--- PSAPI structures and constants -------------------------------------------
46
47LIST_MODULES_DEFAULT    = 0x00
48LIST_MODULES_32BIT      = 0x01
49LIST_MODULES_64BIT      = 0x02
50LIST_MODULES_ALL        = 0x03
51
52# typedef struct _MODULEINFO {
53#   LPVOID lpBaseOfDll;
54#   DWORD  SizeOfImage;
55#   LPVOID EntryPoint;
56# } MODULEINFO, *LPMODULEINFO;
57class MODULEINFO(Structure):
58    _fields_ = [
59        ("lpBaseOfDll",     LPVOID),    # remote pointer
60        ("SizeOfImage",     DWORD),
61        ("EntryPoint",      LPVOID),    # remote pointer
62]
63LPMODULEINFO = POINTER(MODULEINFO)
64
65#--- psapi.dll ----------------------------------------------------------------
66
67# BOOL WINAPI EnumDeviceDrivers(
68#   __out  LPVOID *lpImageBase,
69#   __in   DWORD cb,
70#   __out  LPDWORD lpcbNeeded
71# );
72def EnumDeviceDrivers():
73    _EnumDeviceDrivers = windll.psapi.EnumDeviceDrivers
74    _EnumDeviceDrivers.argtypes = [LPVOID, DWORD, LPDWORD]
75    _EnumDeviceDrivers.restype = bool
76    _EnumDeviceDrivers.errcheck = RaiseIfZero
77
78    size       = 0x1000
79    lpcbNeeded = DWORD(size)
80    unit       = sizeof(LPVOID)
81    while 1:
82        lpImageBase = (LPVOID * (size // unit))()
83        _EnumDeviceDrivers(byref(lpImageBase), lpcbNeeded, byref(lpcbNeeded))
84        needed = lpcbNeeded.value
85        if needed <= size:
86            break
87        size = needed
88    return [ lpImageBase[index] for index in compat.xrange(0, (needed // unit)) ]
89
90# BOOL WINAPI EnumProcesses(
91#   __out  DWORD *pProcessIds,
92#   __in   DWORD cb,
93#   __out  DWORD *pBytesReturned
94# );
95def EnumProcesses():
96    _EnumProcesses = windll.psapi.EnumProcesses
97    _EnumProcesses.argtypes = [LPVOID, DWORD, LPDWORD]
98    _EnumProcesses.restype = bool
99    _EnumProcesses.errcheck = RaiseIfZero
100
101    size            = 0x1000
102    cbBytesReturned = DWORD()
103    unit            = sizeof(DWORD)
104    while 1:
105        ProcessIds = (DWORD * (size // unit))()
106        cbBytesReturned.value = size
107        _EnumProcesses(byref(ProcessIds), cbBytesReturned, byref(cbBytesReturned))
108        returned = cbBytesReturned.value
109        if returned < size:
110            break
111        size = size + 0x1000
112    ProcessIdList = list()
113    for ProcessId in ProcessIds:
114        if ProcessId is None:
115            break
116        ProcessIdList.append(ProcessId)
117    return ProcessIdList
118
119# BOOL WINAPI EnumProcessModules(
120#   __in   HANDLE hProcess,
121#   __out  HMODULE *lphModule,
122#   __in   DWORD cb,
123#   __out  LPDWORD lpcbNeeded
124# );
125def EnumProcessModules(hProcess):
126    _EnumProcessModules = windll.psapi.EnumProcessModules
127    _EnumProcessModules.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD]
128    _EnumProcessModules.restype = bool
129    _EnumProcessModules.errcheck = RaiseIfZero
130
131    size = 0x1000
132    lpcbNeeded = DWORD(size)
133    unit = sizeof(HMODULE)
134    while 1:
135        lphModule = (HMODULE * (size // unit))()
136        _EnumProcessModules(hProcess, byref(lphModule), lpcbNeeded, byref(lpcbNeeded))
137        needed = lpcbNeeded.value
138        if needed <= size:
139            break
140        size = needed
141    return [ lphModule[index] for index in compat.xrange(0, int(needed // unit)) ]
142
143# BOOL WINAPI EnumProcessModulesEx(
144#   __in   HANDLE hProcess,
145#   __out  HMODULE *lphModule,
146#   __in   DWORD cb,
147#   __out  LPDWORD lpcbNeeded,
148#   __in   DWORD dwFilterFlag
149# );
150def EnumProcessModulesEx(hProcess, dwFilterFlag = LIST_MODULES_DEFAULT):
151    _EnumProcessModulesEx = windll.psapi.EnumProcessModulesEx
152    _EnumProcessModulesEx.argtypes = [HANDLE, LPVOID, DWORD, LPDWORD, DWORD]
153    _EnumProcessModulesEx.restype = bool
154    _EnumProcessModulesEx.errcheck = RaiseIfZero
155
156    size = 0x1000
157    lpcbNeeded = DWORD(size)
158    unit = sizeof(HMODULE)
159    while 1:
160        lphModule = (HMODULE * (size // unit))()
161        _EnumProcessModulesEx(hProcess, byref(lphModule), lpcbNeeded, byref(lpcbNeeded), dwFilterFlag)
162        needed = lpcbNeeded.value
163        if needed <= size:
164            break
165        size = needed
166    return [ lphModule[index] for index in compat.xrange(0, (needed // unit)) ]
167
168# DWORD WINAPI GetDeviceDriverBaseName(
169#   __in   LPVOID ImageBase,
170#   __out  LPTSTR lpBaseName,
171#   __in   DWORD nSize
172# );
173def GetDeviceDriverBaseNameA(ImageBase):
174    _GetDeviceDriverBaseNameA = windll.psapi.GetDeviceDriverBaseNameA
175    _GetDeviceDriverBaseNameA.argtypes = [LPVOID, LPSTR, DWORD]
176    _GetDeviceDriverBaseNameA.restype = DWORD
177
178    nSize = MAX_PATH
179    while 1:
180        lpBaseName = ctypes.create_string_buffer("", nSize)
181        nCopied = _GetDeviceDriverBaseNameA(ImageBase, lpBaseName, nSize)
182        if nCopied == 0:
183            raise ctypes.WinError()
184        if nCopied < (nSize - 1):
185            break
186        nSize = nSize + MAX_PATH
187    return lpBaseName.value
188
189def GetDeviceDriverBaseNameW(ImageBase):
190    _GetDeviceDriverBaseNameW = windll.psapi.GetDeviceDriverBaseNameW
191    _GetDeviceDriverBaseNameW.argtypes = [LPVOID, LPWSTR, DWORD]
192    _GetDeviceDriverBaseNameW.restype = DWORD
193
194    nSize = MAX_PATH
195    while 1:
196        lpBaseName = ctypes.create_unicode_buffer(u"", nSize)
197        nCopied = _GetDeviceDriverBaseNameW(ImageBase, lpBaseName, nSize)
198        if nCopied == 0:
199            raise ctypes.WinError()
200        if nCopied < (nSize - 1):
201            break
202        nSize = nSize + MAX_PATH
203    return lpBaseName.value
204
205GetDeviceDriverBaseName = GuessStringType(GetDeviceDriverBaseNameA, GetDeviceDriverBaseNameW)
206
207# DWORD WINAPI GetDeviceDriverFileName(
208#   __in   LPVOID ImageBase,
209#   __out  LPTSTR lpFilename,
210#   __in   DWORD nSize
211# );
212def GetDeviceDriverFileNameA(ImageBase):
213    _GetDeviceDriverFileNameA = windll.psapi.GetDeviceDriverFileNameA
214    _GetDeviceDriverFileNameA.argtypes = [LPVOID, LPSTR, DWORD]
215    _GetDeviceDriverFileNameA.restype = DWORD
216
217    nSize = MAX_PATH
218    while 1:
219        lpFilename = ctypes.create_string_buffer("", nSize)
220        nCopied = ctypes.windll.psapi.GetDeviceDriverFileNameA(ImageBase, lpFilename, nSize)
221        if nCopied == 0:
222            raise ctypes.WinError()
223        if nCopied < (nSize - 1):
224            break
225        nSize = nSize + MAX_PATH
226    return lpFilename.value
227
228def GetDeviceDriverFileNameW(ImageBase):
229    _GetDeviceDriverFileNameW = windll.psapi.GetDeviceDriverFileNameW
230    _GetDeviceDriverFileNameW.argtypes = [LPVOID, LPWSTR, DWORD]
231    _GetDeviceDriverFileNameW.restype = DWORD
232
233    nSize = MAX_PATH
234    while 1:
235        lpFilename = ctypes.create_unicode_buffer(u"", nSize)
236        nCopied = ctypes.windll.psapi.GetDeviceDriverFileNameW(ImageBase, lpFilename, nSize)
237        if nCopied == 0:
238            raise ctypes.WinError()
239        if nCopied < (nSize - 1):
240            break
241        nSize = nSize + MAX_PATH
242    return lpFilename.value
243
244GetDeviceDriverFileName = GuessStringType(GetDeviceDriverFileNameA, GetDeviceDriverFileNameW)
245
246# DWORD WINAPI GetMappedFileName(
247#   __in   HANDLE hProcess,
248#   __in   LPVOID lpv,
249#   __out  LPTSTR lpFilename,
250#   __in   DWORD nSize
251# );
252def GetMappedFileNameA(hProcess, lpv):
253    _GetMappedFileNameA = ctypes.windll.psapi.GetMappedFileNameA
254    _GetMappedFileNameA.argtypes = [HANDLE, LPVOID, LPSTR, DWORD]
255    _GetMappedFileNameA.restype = DWORD
256
257    nSize = MAX_PATH
258    while 1:
259        lpFilename = ctypes.create_string_buffer("", nSize)
260        nCopied = _GetMappedFileNameA(hProcess, lpv, lpFilename, nSize)
261        if nCopied == 0:
262            raise ctypes.WinError()
263        if nCopied < (nSize - 1):
264            break
265        nSize = nSize + MAX_PATH
266    return lpFilename.value
267
268def GetMappedFileNameW(hProcess, lpv):
269    _GetMappedFileNameW = ctypes.windll.psapi.GetMappedFileNameW
270    _GetMappedFileNameW.argtypes = [HANDLE, LPVOID, LPWSTR, DWORD]
271    _GetMappedFileNameW.restype = DWORD
272
273    nSize = MAX_PATH
274    while 1:
275        lpFilename = ctypes.create_unicode_buffer(u"", nSize)
276        nCopied = _GetMappedFileNameW(hProcess, lpv, lpFilename, nSize)
277        if nCopied == 0:
278            raise ctypes.WinError()
279        if nCopied < (nSize - 1):
280            break
281        nSize = nSize + MAX_PATH
282    return lpFilename.value
283
284GetMappedFileName = GuessStringType(GetMappedFileNameA, GetMappedFileNameW)
285
286# DWORD WINAPI GetModuleFileNameEx(
287#   __in      HANDLE hProcess,
288#   __in_opt  HMODULE hModule,
289#   __out     LPTSTR lpFilename,
290#   __in      DWORD nSize
291# );
292def GetModuleFileNameExA(hProcess, hModule = None):
293    _GetModuleFileNameExA = ctypes.windll.psapi.GetModuleFileNameExA
294    _GetModuleFileNameExA.argtypes = [HANDLE, HMODULE, LPSTR, DWORD]
295    _GetModuleFileNameExA.restype = DWORD
296
297    nSize = MAX_PATH
298    while 1:
299        lpFilename = ctypes.create_string_buffer("", nSize)
300        nCopied = _GetModuleFileNameExA(hProcess, hModule, lpFilename, nSize)
301        if nCopied == 0:
302            raise ctypes.WinError()
303        if nCopied < (nSize - 1):
304            break
305        nSize = nSize + MAX_PATH
306    return lpFilename.value
307
308def GetModuleFileNameExW(hProcess, hModule = None):
309    _GetModuleFileNameExW = ctypes.windll.psapi.GetModuleFileNameExW
310    _GetModuleFileNameExW.argtypes = [HANDLE, HMODULE, LPWSTR, DWORD]
311    _GetModuleFileNameExW.restype = DWORD
312
313    nSize = MAX_PATH
314    while 1:
315        lpFilename = ctypes.create_unicode_buffer(u"", nSize)
316        nCopied = _GetModuleFileNameExW(hProcess, hModule, lpFilename, nSize)
317        if nCopied == 0:
318            raise ctypes.WinError()
319        if nCopied < (nSize - 1):
320            break
321        nSize = nSize + MAX_PATH
322    return lpFilename.value
323
324GetModuleFileNameEx = GuessStringType(GetModuleFileNameExA, GetModuleFileNameExW)
325
326# BOOL WINAPI GetModuleInformation(
327#   __in   HANDLE hProcess,
328#   __in   HMODULE hModule,
329#   __out  LPMODULEINFO lpmodinfo,
330#   __in   DWORD cb
331# );
332def GetModuleInformation(hProcess, hModule, lpmodinfo = None):
333    _GetModuleInformation = windll.psapi.GetModuleInformation
334    _GetModuleInformation.argtypes = [HANDLE, HMODULE, LPMODULEINFO, DWORD]
335    _GetModuleInformation.restype = bool
336    _GetModuleInformation.errcheck = RaiseIfZero
337
338    if lpmodinfo is None:
339        lpmodinfo = MODULEINFO()
340    _GetModuleInformation(hProcess, hModule, byref(lpmodinfo), sizeof(lpmodinfo))
341    return lpmodinfo
342
343# DWORD WINAPI GetProcessImageFileName(
344#   __in   HANDLE hProcess,
345#   __out  LPTSTR lpImageFileName,
346#   __in   DWORD nSize
347# );
348def GetProcessImageFileNameA(hProcess):
349    _GetProcessImageFileNameA = windll.psapi.GetProcessImageFileNameA
350    _GetProcessImageFileNameA.argtypes = [HANDLE, LPSTR, DWORD]
351    _GetProcessImageFileNameA.restype = DWORD
352
353    nSize = MAX_PATH
354    while 1:
355        lpFilename = ctypes.create_string_buffer("", nSize)
356        nCopied = _GetProcessImageFileNameA(hProcess, lpFilename, nSize)
357        if nCopied == 0:
358            raise ctypes.WinError()
359        if nCopied < (nSize - 1):
360            break
361        nSize = nSize + MAX_PATH
362    return lpFilename.value
363
364def GetProcessImageFileNameW(hProcess):
365    _GetProcessImageFileNameW = windll.psapi.GetProcessImageFileNameW
366    _GetProcessImageFileNameW.argtypes = [HANDLE, LPWSTR, DWORD]
367    _GetProcessImageFileNameW.restype = DWORD
368
369    nSize = MAX_PATH
370    while 1:
371        lpFilename = ctypes.create_unicode_buffer(u"", nSize)
372        nCopied = _GetProcessImageFileNameW(hProcess, lpFilename, nSize)
373        if nCopied == 0:
374            raise ctypes.WinError()
375        if nCopied < (nSize - 1):
376            break
377        nSize = nSize + MAX_PATH
378    return lpFilename.value
379
380GetProcessImageFileName = GuessStringType(GetProcessImageFileNameA, GetProcessImageFileNameW)
381
382#==============================================================================
383# This calculates the list of exported symbols.
384_all = set(vars().keys()).difference(_all)
385__all__ = [_x for _x in _all if not _x.startswith('_')]
386__all__.sort()
387#==============================================================================
388