common.cpp

一个进程管理的源程序.zip

#include "stdafx.h"




#include <windows.h>
#include <winperf.h>   // for Windows NT
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tlhelp32.h>  // for Windows 95

#include "common.h"


//
// manafest constants
//
#define INITIAL_SIZE        51200
#define EXTEND_SIZE         25600
#define REGKEY_PERF         "software\\microsoft\\windows nt\\currentversion\\perflib"
#define REGSUBKEY_COUNTERS  "Counters"
#define PROCESS_COUNTER     "process"
#define PROCESSID_COUNTER   "id process"
#define UNKNOWN_TASK        "unknown"


typedef BOOL (WINAPI *PROCESSWALK)(HANDLE hSnapshot, LPPROCESSENTRY32 lppe);
typedef HANDLE (WINAPI *CREATESNAPSHOT)(DWORD dwFlags, DWORD th32ProcessID);


//
// prototypes
//
BOOL CALLBACK
EnumWindowsProc(
    HWND    hwnd,
    LPARAM lParam
    );



DWORD GetTaskList95(CTaskList& refList)


{
   CREATESNAPSHOT pCreateToolhelp32Snapshot = NULL;
   PROCESSWALK    pProcess32First           = NULL;
   PROCESSWALK    pProcess32Next            = NULL;

   HINSTANCE      hKernel        = NULL;
   HINSTANCE      hProcessSnap   = NULL;
   PROCESSENTRY32 pe32           = {0};
   DWORD          dwTaskCount    = 0;
    hKernel = GetModuleHandle("KERNEL32.DLL");

    if (hKernel)
    {
        pCreateToolhelp32Snapshot =
          (CREATESNAPSHOT)GetProcAddress(hKernel, "CreateToolhelp32Snapshot");

        pProcess32First = (PROCESSWALK)GetProcAddress(hKernel,
                                                      "Process32First");
        pProcess32Next  = (PROCESSWALK)GetProcAddress(hKernel,
                                                      "Process32Next");
    }
    
    // make sure we got addresses of all needed Toolhelp functions.
    if (!(pProcess32First && pProcess32Next && pCreateToolhelp32Snapshot))
       return 0;
       

    // Take a snapshot of all processes currently in the system.
    hProcessSnap = (HINSTANCE) pCreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hProcessSnap == (HANDLE)-1)
        return 0;

    // Walk the snapshot of processes and for each process, get information
    // to display.
    dwTaskCount = 0;
    pe32.dwSize = sizeof(PROCESSENTRY32);   // must be filled out before use
    if (pProcess32First(hProcessSnap, &pe32))
    {
        do
        {
            LPSTR pCurChar;

			//SAMPLE: use our class instead of raw structs
			CTaskListEntry* pEntry = new CTaskListEntry;

            // strip path and leave executabe filename splitpath
            for (pCurChar = (pe32.szExeFile + lstrlen (pe32.szExeFile));
                 *pCurChar != '\\' && pCurChar != pe32.szExeFile; --pCurChar) 
					 lstrcpy(pEntry->ProcessName, pCurChar);
			lstrcpy(pEntry->szExeFile,pe32.szExeFile);
            pEntry->flags =pe32.dwFlags;
			HANDLE hProcess;
           hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pe32.th32ProcessID );

		   pEntry->processPriority=GetPriorityClass(hProcess);
	
			pEntry->pcPriClassBase=pe32.pcPriClassBase;
            pEntry->dwProcessId = pe32.th32ProcessID;
            pEntry->th32ParentProcessID=pe32.th32ParentProcessID;    
			pEntry->cntThreads= pe32.cntThreads;
            ++dwTaskCount;   // keep track of how many tasks we've got so far
			refList.Add(pEntry);
        }
        while (pProcess32Next(hProcessSnap, &pe32));
    }
    else
        dwTaskCount = 0;    // Couldn't walk the list of processes.

    // Don't forget to clean up the snapshot object...
    CloseHandle (hProcessSnap);

    return dwTaskCount;
}


DWORD GetTaskListNT(CTaskList& refList)


{
    DWORD                        rc;
    HKEY                         hKeyNames;
    DWORD                        dwType;
    DWORD                        dwSize;
    LPSTR                        buf = NULL;
    CHAR                         szSubKey[1024];
    LANGID                       lid;
    LPSTR                        p;
    LPSTR                        p2;
    PPERF_DATA_BLOCK             pPerf;
    PPERF_OBJECT_TYPE            pObj;
    PPERF_INSTANCE_DEFINITION    pInst;
    PPERF_COUNTER_BLOCK          pCounter;
    PPERF_COUNTER_DEFINITION     pCounterDef;
    DWORD                        i;
    DWORD                        dwProcessIdTitle;
    DWORD                        dwProcessIdCounter;
    CHAR                         szProcessName[MAX_PATH];

    lid = MAKELANGID( LANG_ENGLISH, SUBLANG_NEUTRAL );
    sprintf( szSubKey, "%s\\%03x", REGKEY_PERF, lid );
    rc = RegOpenKeyEx( HKEY_LOCAL_MACHINE,
                       szSubKey,
                       0,
                       KEY_READ,
                       &hKeyNames
                     );
    if (rc != ERROR_SUCCESS) {
        goto exit;
    }

    //
    // get the buffer size for the counter names
    //
    rc = RegQueryValueEx( hKeyNames,
                          REGSUBKEY_COUNTERS,
                          NULL,
                          &dwType,
                          NULL,
                          &dwSize
                        );

    if (rc != ERROR_SUCCESS) {
        goto exit;
    }

    //
    // allocate the counter names buffer
    //
    buf = (LPSTR) malloc( dwSize );
    if (buf == NULL)
        goto exit;
    memset( buf, 0, dwSize );

    //
    // read the counter names from the registry
    //
    rc = RegQueryValueEx( hKeyNames,
                          REGSUBKEY_COUNTERS,
                          NULL,
                          &dwType,
                          (LPBYTE) buf,
                          &dwSize
                        );

    if (rc != ERROR_SUCCESS) {
        goto exit;
    }


    p = buf;
    while (*p) {
        if (p > buf) {
            for( p2=p-2; isdigit(*p2); p2--) ;
            }
        if (stricmp(p, PROCESS_COUNTER) == 0) {
            //
            // look backwards for the counter number
            //
            for( p2=p-2; isdigit(*p2); p2--) ;
            strcpy( szSubKey, p2+1 );
        }
        else
        if (stricmp(p, PROCESSID_COUNTER) == 0) {
            //
            // look backwards for the counter number
            //
            for( p2=p-2; isdigit(*p2); p2--) ;
            dwProcessIdTitle = atol( p2+1 );
        }
        //
        // next string
        //
        p += (strlen(p) + 1);
    }

    //
    // free the counter names buffer
    //
    free( buf );


    //
    // allocate the initial buffer for the performance data
    //
    dwSize = INITIAL_SIZE;
    buf = (LPSTR) malloc( dwSize );
    if (buf == NULL) {
        goto exit;
    }
    memset( buf, 0, dwSize );


    while (TRUE) {

        rc = RegQueryValueEx( HKEY_PERFORMANCE_DATA,
                              szSubKey,
                              NULL,
                              &dwType,
                              (LPBYTE) buf,
                              &dwSize
                            );

        pPerf = (PPERF_DATA_BLOCK) buf;

        //
        // check for success and valid perf data block signature
        //
        if ((rc == ERROR_SUCCESS) &&
            (dwSize > 0) &&
            (pPerf)->Signature[0] == (WCHAR)'P' &&
            (pPerf)->Signature[1] == (WCHAR)'E' &&
            (pPerf)->Signature[2] == (WCHAR)'R' &&
            (pPerf)->Signature[3] == (WCHAR)'F' ) {
            break;
        }

        //
        // if buffer is not big enough, reallocate and try again
        //
        if (rc == ERROR_MORE_DATA) {
            dwSize += EXTEND_SIZE;
            buf = (LPSTR) realloc( buf, dwSize );
            memset( buf, 0, dwSize );
        }
        else {
            goto exit;
        }
    }

    //
    // set the perf_object_type pointer
    //
    pObj = (PPERF_OBJECT_TYPE) ((DWORD)pPerf + pPerf->HeaderLength);

    //
    // loop thru the performance counter definition records looking
    // for the process id counter and then save its offset
    //
    pCounterDef = (PPERF_COUNTER_DEFINITION) ((DWORD)pObj + pObj->HeaderLength);
    for (i=0; i<(DWORD)pObj->NumCounters; i++) {
        if (pCounterDef->CounterNameTitleIndex == dwProcessIdTitle) {
            dwProcessIdCounter = pCounterDef->CounterOffset;
            break;
        }
        pCounterDef++;
    }

    pInst = (PPERF_INSTANCE_DEFINITION) ((DWORD)pObj + pObj->DefinitionLength);

    //
    // loop thru the performance instance data extracting each process name
    // and process id
    //
    for (i=0; i < (DWORD)pObj->NumInstances; i++)
	{
		CTaskListEntry* pEntry = new CTaskListEntry;

        //
        // pointer to the process name
        //
        p = (LPSTR) ((DWORD)pInst + pInst->NameOffset);

        //
        // convert it to ascii
        //
        rc = WideCharToMultiByte( CP_ACP,
                                  0,
                                  (LPCWSTR)p,
                                  -1,
                                  szProcessName,
                                  sizeof(szProcessName),
                                  NULL,
                                  NULL
                                );

        if (!rc) {
            //
		    // if we cant convert the string then use a default value
            //
            strcpy( pEntry->ProcessName, UNKNOWN_TASK );
        }

        if (strlen(szProcessName)+4 <= sizeof(pEntry->ProcessName)) {
            strcpy( pEntry->ProcessName, szProcessName );
            strcat( pEntry->ProcessName, ".exe" );
        }

        //
        // get the process id
        //
        pCounter = (PPERF_COUNTER_BLOCK) ((DWORD)pInst + pInst->ByteLength);
        pEntry->flags = 0;
        pEntry->dwProcessId = *((LPDWORD) ((DWORD)pCounter + dwProcessIdCounter));
        if (pEntry->dwProcessId == 0) {
            pEntry->dwProcessId = (DWORD)-2;
        }

        //
        // next process
        //
		refList.Add(pEntry);
        pInst = (PPERF_INSTANCE_DEFINITION) ((DWORD)pCounter + pCounter->ByteLength);
    }

exit:
    if (buf) {
        free( buf );
    }

    RegCloseKey( hKeyNames );
    RegCloseKey( HKEY_PERFORMANCE_DATA );

    return (DWORD)pObj->NumInstances;
}



BOOL
EnableDebugPriv95(
    VOID
    )


{
   return TRUE;
}



BOOL EnableDebugPrivNT( VOID  )


{
    HANDLE hToken;
    LUID DebugValue;
    TOKEN_PRIVILEGES tkp;


    //
    // Retrieve a handle of the access token
    //
    if (!OpenProcessToken(GetCurrentProcess(),
            TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
            &hToken)) {
        printf("OpenProcessToken failed with %d\n", GetLastError());
        return FALSE;
    }

    //
    // Enable the SE_DEBUG_NAME privilege
    //
    if (!LookupPrivilegeValue((LPSTR) NULL,
            SE_DEBUG_NAME,
            &DebugValue)) {
        printf("LookupPrivilegeValue failed with %d\n", GetLastError());
        return FALSE;
    }

    tkp.PrivilegeCount = 1;
    tkp.Privileges[0].Luid = DebugValue;
    tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    AdjustTokenPrivileges(hToken,
        FALSE,
        &tkp,
        sizeof(TOKEN_PRIVILEGES),
        (PTOKEN_PRIVILEGES) NULL,
        (PDWORD) NULL);

    //
    // The return value of AdjustTokenPrivileges can't be tested
    //
    if (GetLastError() != ERROR_SUCCESS) {
        printf("AdjustTokenPrivileges failed with %d\n", GetLastError());
        return FALSE;
    }

    return TRUE;
}

BOOL KillProcess(CTaskListEntry* pEntry, BOOL fForce)
{
    HANDLE  hProcess;
    if (fForce || !pEntry->hwnd)
	{
        hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pEntry->dwProcessId );
        if (hProcess)
		{
            hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pEntry->dwProcessId );
            if (hProcess == NULL)
			{
                return FALSE;
            }
          
            if (!TerminateProcess( hProcess, 1)) 
			{
                CloseHandle( hProcess );
                return FALSE;
            }
            
			CloseHandle( hProcess );
            return TRUE;
        }
    }

    //
    // kill the process
    //
    PostMessage(pEntry->hwnd, WM_CLOSE, 0, 0);

    return TRUE;
}


VOID GetWindowTitles(CTaskList& refList)
{
    //
    // enumerate all windows
    //
    EnumWindows( EnumWindowsProc, (LPARAM) &refList );
}

BOOL CALLBACK EnumWindowsProc(HWND    hwnd,  LPARAM   lParam )
{
    DWORD             pid = 0;
    int               i;
    CHAR              buf[TITLE_SIZE];
	CTaskList&		  refList = *((CTaskList*) lParam);

    //
    // get the processid for this window
    //
    if (!GetWindowThreadProcessId( hwnd, &pid )) {
        return TRUE;
    }

    //
    // look for the task in the task list for this window
    //
    for (i=0; i < refList.GetSize(); i++)
	{
        if (refList[i]->dwProcessId == pid)
		{
            refList[i]->hwnd = hwnd;
            //
			// we found the task so lets try to get the
            // window text
            //
            if (GetWindowText( refList[i]->hwnd, buf, sizeof(buf) ))
			{
                //
				// got it, so lets save it
                //
                strcpy( refList[i]->WindowTitle, buf);
            }
            break;
        }
    }

    //
    // continue the enumeration
    //
    return TRUE;
}