source: nscp/trunk/include/EnumProcess.cpp @ d2394a2

Last change on this file since d2394a2 was 484be40, checked in by Michael Medin <michael@…>, 6 years ago
  • Restructured build and include to work from CLI + Started to add makefiles to build from CLI + Added finnish counters (#98)
    • Fixed performance counter issue with check_nt (#99)
  • Property mode set to 100644
File size: 10.3 KB
Line 
1/**************************************************************************
2*   Copyright (C) 2004-2007 by Michael Medin <michael@medin.name>         *
3*                                                                         *
4*   This code is part of NSClient++ - http://trac.nakednuns.org/nscp      *
5*                                                                         *
6*   This program is free software; you can redistribute it and/or modify  *
7*   it under the terms of the GNU General Public License as published by  *
8*   the Free Software Foundation; either version 2 of the License, or     *
9*   (at your option) any later version.                                   *
10*                                                                         *
11*   This program is distributed in the hope that it will be useful,       *
12*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
13*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
14*   GNU General Public License for more details.                          *
15*                                                                         *
16*   You should have received a copy of the GNU General Public License     *
17*   along with this program; if not, write to the                         *
18*   Free Software Foundation, Inc.,                                       *
19*   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
20***************************************************************************/
21#define WIN32_LEAN_AND_MEAN             // Exclude rarely-used stuff from Windows headers
22#include <windows.h>
23#include <tchar.h>
24
25#include "EnumProcess.h"
26
27//////////////////////////////////////////////////////////////////////
28// Construction/Destruction
29//////////////////////////////////////////////////////////////////////
30
31
32CEnumProcess::CEnumProcess() : m_pProcesses(NULL), m_pModules(NULL), m_pCurrentP(NULL), m_pCurrentM(NULL), lpString(NULL)
33{
34        lpString = new CHAR[MAX_FILENAME+1];
35        m_hProcessSnap = INVALID_HANDLE_VALUE;
36        m_hModuleSnap = INVALID_HANDLE_VALUE;
37
38        PSAPI = ::LoadLibrary(_TEXT("PSAPI"));
39        if (PSAPI) 
40        {
41                // Setup variables
42                m_MAX_COUNT = 256;
43                m_cProcesses = 0;
44                m_cModules   = 0;
45
46                // Find PSAPI functions
47                FEnumProcesses = (PFEnumProcesses)::GetProcAddress(PSAPI, _TEXT("EnumProcesses"));
48                FEnumProcessModules = (PFEnumProcessModules)::GetProcAddress(PSAPI, _TEXT("EnumProcessModules"));
49#ifdef UNICODE
50                FGetModuleFileNameEx = (PFGetModuleFileNameEx)::GetProcAddress(PSAPI, _TEXT("GetModuleFileNameExW"));
51#else
52                FGetModuleFileNameEx = (PFGetModuleFileNameEx)::GetProcAddress(PSAPI, _TEXT("GetModuleFileNameExA"));
53#endif
54        }
55
56        TOOLHELP = ::LoadLibrary(_TEXT("Kernel32"));
57        if (TOOLHELP) 
58        {
59                // Setup variables
60                m_pe.dwSize = sizeof(m_pe);
61                m_me.dwSize = sizeof(m_me);
62                // Find ToolHelp functions
63
64                FCreateToolhelp32Snapshot = (PFCreateToolhelp32Snapshot)::GetProcAddress(TOOLHELP, _TEXT("CreateToolhelp32Snapshot"));
65                FProcess32First = (PFProcess32First)::GetProcAddress(TOOLHELP, _TEXT("Process32First"));
66                FProcess32Next = (PFProcess32Next)::GetProcAddress(TOOLHELP, _TEXT("Process32Next"));
67                FModule32First = (PFModule32First)::GetProcAddress(TOOLHELP, _TEXT("Module32First"));
68                FModule32Next = (PFModule32Next)::GetProcAddress(TOOLHELP, _TEXT("Module32Next"));
69        }
70
71        // Find the preferred method of enumeration
72        m_method = ENUM_METHOD::NONE;
73        int method = GetAvailableMethods();
74        if (method == (method|ENUM_METHOD::PSAPI))    m_method = ENUM_METHOD::PSAPI;
75        if (method == (method|ENUM_METHOD::TOOLHELP)) m_method = ENUM_METHOD::TOOLHELP;
76        if (method == (method|ENUM_METHOD::PROC16))   m_method += ENUM_METHOD::PROC16;
77
78}
79
80CEnumProcess::~CEnumProcess()
81{
82        delete [] lpString;
83        if (m_pProcesses) {delete[] m_pProcesses;}
84        if (m_pModules)   {delete[] m_pModules;}
85        if (PSAPI) FreeLibrary(PSAPI);
86        if (TOOLHELP) FreeLibrary(TOOLHELP);
87        if (INVALID_HANDLE_VALUE != m_hProcessSnap) ::CloseHandle(m_hProcessSnap);
88        if (INVALID_HANDLE_VALUE != m_hModuleSnap)  ::CloseHandle(m_hModuleSnap);
89}
90
91
92
93int CEnumProcess::GetAvailableMethods()
94{
95        int res = 0;
96        // Does all psapi functions exist?
97        if (PSAPI&&FEnumProcesses&&FEnumProcessModules&&FGetModuleFileNameEx)
98                res += ENUM_METHOD::PSAPI;
99        // How about Toolhelp?
100        if (TOOLHELP&&FCreateToolhelp32Snapshot&&FProcess32Next&&FProcess32Next&&FModule32First&&FModule32Next)
101                res += ENUM_METHOD::TOOLHELP;
102
103        return res;
104}
105
106int CEnumProcess::SetMethod(int method)
107{
108        int avail = GetAvailableMethods();
109
110        if (method != ENUM_METHOD::PROC16 && avail == (method|avail))
111                m_method = method;
112
113        return m_method;
114}
115
116int CEnumProcess::GetSuggestedMethod()
117{
118        return m_method;
119}
120// Retrieves the first process in the enumeration. Should obviously be called before
121// GetProcessNext
122////////////////////////////////////////////////////////////////////////////////////
123BOOL CEnumProcess::GetProcessFirst(CEnumProcess::CProcessEntry *pEntry)
124{
125        if (ENUM_METHOD::NONE == m_method) return FALSE;
126
127
128        if ((ENUM_METHOD::TOOLHELP|m_method) == m_method)
129                // Use ToolHelp functions
130                // ----------------------
131        {
132                m_hProcessSnap = FCreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
133                if (INVALID_HANDLE_VALUE == m_hProcessSnap) return FALSE;
134                if (!FProcess32First(m_hProcessSnap, &m_pe)) return FALSE;
135                pEntry->dwPID = m_pe.th32ProcessID;
136                pEntry->sFilename, m_pe.szExeFile;
137        }
138        else
139                // Use PSAPI functions
140                // ----------------------
141        {
142                if (m_pProcesses) {delete[] m_pProcesses;}
143                m_pProcesses = new DWORD[m_MAX_COUNT];
144                m_pCurrentP = m_pProcesses;
145                DWORD cbNeeded = 0;
146                BOOL OK = FEnumProcesses(m_pProcesses, m_MAX_COUNT*sizeof(DWORD), &cbNeeded);
147
148                // We might need more memory here..
149                if (cbNeeded >= m_MAX_COUNT*sizeof(DWORD))
150                {
151                        m_MAX_COUNT += 256;
152                        return GetProcessFirst(pEntry); // Try again.
153                }
154
155                if (!OK) return FALSE;
156                m_cProcesses = cbNeeded/sizeof(DWORD);
157                return FillPStructPSAPI(*m_pProcesses, pEntry);
158        }
159
160        return TRUE;
161}
162
163// Returns the following process
164////////////////////////////////////////////////////////////////
165BOOL CEnumProcess::GetProcessNext(CEnumProcess::CProcessEntry *pEntry)
166{
167        if (ENUM_METHOD::NONE == m_method) return FALSE;
168        pEntry->hTask16 = 0;
169
170
171        // Use ToolHelp functions
172        // ----------------------
173        if ((ENUM_METHOD::TOOLHELP|m_method) == m_method)
174        {
175                if (!FProcess32Next(m_hProcessSnap, &m_pe)) return FALSE;
176                pEntry->dwPID = m_pe.th32ProcessID;
177                pEntry->sFilename = m_pe.szExeFile;
178        }
179        else
180                // Use PSAPI functions
181                // ----------------------
182        {
183                if (--m_cProcesses <= 0) return FALSE;
184                FillPStructPSAPI(*++m_pCurrentP, pEntry);
185        }
186
187        return TRUE;
188}
189
190
191BOOL CEnumProcess::GetModuleFirst(DWORD dwPID, CEnumProcess::CModuleEntry *pEntry)
192{
193        if (ENUM_METHOD::NONE == m_method) return FALSE;
194        // Use ToolHelp functions
195        // ----------------------
196        if ((ENUM_METHOD::TOOLHELP|m_method) == m_method)
197        {
198                if (INVALID_HANDLE_VALUE != m_hModuleSnap)  ::CloseHandle(m_hModuleSnap);
199                m_hModuleSnap = FCreateToolhelp32Snapshot(TH32CS_SNAPMODULE, dwPID);
200
201                if(!FModule32First(m_hModuleSnap, &m_me)) return FALSE;
202
203                pEntry->pLoadBase = m_me.modBaseAddr;
204                pEntry->sFilename = m_me.szExePath;
205                pEntry->pPreferredBase = GetModulePreferredBase(dwPID, m_me.modBaseAddr);
206                return TRUE;
207        }
208        else
209                // Use PSAPI functions
210                // ----------------------
211        {
212                if (m_pModules) {delete[] m_pModules;}
213                m_pModules = new HMODULE[m_MAX_COUNT];
214                m_pCurrentM = m_pModules;
215                DWORD cbNeeded = 0;
216                HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, dwPID);
217                if (hProc)
218                {
219                        BOOL OK = FEnumProcessModules(hProc, m_pModules, m_MAX_COUNT*sizeof(HMODULE), &cbNeeded);
220                        CloseHandle(hProc);
221
222                        // We might need more memory here..
223                        if (cbNeeded >= m_MAX_COUNT*sizeof(HMODULE))
224                        {
225                                m_MAX_COUNT += 256;
226                                return GetModuleFirst(dwPID, pEntry); // Try again.
227                        }
228
229                        if (!OK) return FALSE;
230
231                        m_cModules = cbNeeded/sizeof(HMODULE);
232                        return FillMStructPSAPI(dwPID, *m_pCurrentM, pEntry);
233                }
234                return FALSE;
235        }
236}
237
238
239BOOL CEnumProcess::GetModuleNext(DWORD dwPID, CEnumProcess::CModuleEntry *pEntry)
240{
241        if (ENUM_METHOD::NONE == m_method) return FALSE;
242
243        // Use ToolHelp functions
244        // ----------------------
245        if ((ENUM_METHOD::TOOLHELP|m_method) == m_method)
246        {
247                if(!FModule32Next(m_hModuleSnap, &m_me)) return FALSE;
248
249                pEntry->pLoadBase = m_me.modBaseAddr;
250                pEntry->sFilename = m_me.szExePath;
251                pEntry->pPreferredBase = GetModulePreferredBase(dwPID, m_me.modBaseAddr);
252                return TRUE;
253        }
254        else
255                // Use PSAPI functions
256                // ----------------------
257        {
258                if (--m_cModules <= 0) return FALSE;
259                return FillMStructPSAPI(dwPID, *++m_pCurrentM, pEntry);
260        }
261
262}
263
264
265
266BOOL CEnumProcess::FillPStructPSAPI(DWORD dwPID, CEnumProcess::CProcessEntry* pEntry)
267{
268        pEntry->dwPID = dwPID;
269
270        // Open process to get filename
271        HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, dwPID);
272        if (hProc)
273        {
274                HMODULE hMod;
275                DWORD size;
276                // Get the first module (the process itself)
277                if( FEnumProcessModules(hProc, &hMod, sizeof(hMod), &size) )
278                {
279                        //Get filename
280
281                        if( !FGetModuleFileNameEx( hProc, hMod, lpString, MAX_FILENAME) ) {
282                                pEntry->sFilename = "N/A (error)";
283                        } else {
284                                pEntry->sFilename = lpString;
285                        }
286                }
287                CloseHandle(hProc);
288        }
289        else
290                pEntry->sFilename = "N/A (security restriction)";
291
292        return TRUE;
293}
294
295
296BOOL CEnumProcess::FillMStructPSAPI(DWORD dwPID, HMODULE mMod, CEnumProcess::CModuleEntry *pEntry)
297{
298        HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, dwPID);
299        if (hProc)
300        {
301                if( !FGetModuleFileNameEx( hProc, mMod, lpString, MAX_FILENAME) )
302                {
303                        pEntry->sFilename = "N/A (error)";
304                } else {
305                        pEntry->sFilename = lpString;
306                }
307                pEntry->pLoadBase = (PVOID) mMod;
308                pEntry->pPreferredBase = GetModulePreferredBase(dwPID, (PVOID)mMod);
309                CloseHandle(hProc);
310                return TRUE;
311        }
312        return FALSE;
313}
314
315
316
317PVOID CEnumProcess::GetModulePreferredBase(DWORD dwPID, PVOID pModBase)
318{
319        if (ENUM_METHOD::NONE == m_method) return NULL;
320        HANDLE hProc = OpenProcess(PROCESS_VM_READ, FALSE, dwPID);
321        if (hProc)
322        {
323                IMAGE_DOS_HEADER idh;
324                IMAGE_NT_HEADERS inh;
325                //Read DOS header
326                ReadProcessMemory(hProc, pModBase, &idh, sizeof(idh), NULL);
327
328                if (IMAGE_DOS_SIGNATURE == idh.e_magic) // DOS header OK?
329                        // Read NT headers at offset e_lfanew
330                        ReadProcessMemory(hProc, (PBYTE)pModBase + idh.e_lfanew, &inh, sizeof(inh), NULL);
331
332                CloseHandle(hProc);
333
334                if (IMAGE_NT_SIGNATURE == inh.Signature) //NT signature OK?
335                        // Get the preferred base...
336                        return (PVOID) inh.OptionalHeader.ImageBase;
337
338        }
339
340        return NULL; //didn't find anything useful..
341}
342
343
Note: See TracBrowser for help on using the repository browser.