common.cpp
上传用户:zhenhuadz
上传日期:2007-01-07
资源大小:1511k
文件大小:16k
源码类别:

远程控制编程

开发平台:

Visual C++

  1. /******************************************************************************
  2. *       This is a part of the Microsoft Source Code Samples. 
  3. *       Copyright (C) 1994-1996 Microsoft Corporation.
  4. *       All rights reserved.
  5. *       This source code is only intended as a supplement to 
  6. *       Microsoft Development Tools and/or WinHelp documentation.
  7. *       See these sources for detailed information regarding the 
  8. *       Microsoft samples programs.
  9. ******************************************************************************/
  10. //SAMPLE: hacked to use MFC by mikeblas@nwlink.com
  11. #include "stdafx.h"
  12. /*++
  13. Module Name:
  14.     common.c
  15. Abstract:
  16.     This module contains common apis used by tlist & kill.
  17. --*/
  18. #include <windows.h>
  19. #include <winperf.h>   // for Windows NT
  20. #include <stdio.h>
  21. #include <stdlib.h>
  22. #include <string.h>
  23. #include <tlhelp32.h>  // for Windows 95
  24. #include "common.h"
  25. //
  26. // manafest constants
  27. //
  28. #define INITIAL_SIZE        51200
  29. #define EXTEND_SIZE         25600
  30. #define REGKEY_PERF         "software\microsoft\windows nt\currentversion\perflib"
  31. #define REGSUBKEY_COUNTERS  "Counters"
  32. #define PROCESS_COUNTER     "process"
  33. #define PROCESSID_COUNTER   "id process"
  34. #define UNKNOWN_TASK        "unknown"
  35. //
  36. // Function pointer types for accessing Toolhelp32 functions dynamically.
  37. // By dynamically accessing these functions, we can do so only on Windows
  38. // 95 and still run on Windows NT, which does not have these functions.
  39. //
  40. typedef BOOL (WINAPI *PROCESSWALK)(HANDLE hSnapshot, LPPROCESSENTRY32 lppe);
  41. typedef HANDLE (WINAPI *CREATESNAPSHOT)(DWORD dwFlags, DWORD th32ProcessID);
  42. //
  43. // prototypes
  44. //
  45. BOOL CALLBACK
  46. EnumWindowsProc(
  47.     HWND    hwnd,
  48.     LPARAM lParam
  49.     );
  50. DWORD GetTaskList95(CTaskList& refList)
  51. /*++
  52. Routine Description:
  53.     Provides an API for getting a list of tasks running at the time of the
  54.     API call.  This function uses Toolhelp32to get the task list and is 
  55.     therefore straight WIN32 calls that anyone can call.
  56. Arguments:
  57.     dwNumTasks       - maximum number of tasks that the pTask array can hold
  58. Return Value:
  59.     Number of tasks placed into the pTask array.
  60. --*/
  61. {
  62.    CREATESNAPSHOT pCreateToolhelp32Snapshot = NULL;
  63.    PROCESSWALK    pProcess32First           = NULL;
  64.    PROCESSWALK    pProcess32Next            = NULL;
  65.    HINSTANCE      hKernel        = NULL;
  66.    HINSTANCE      hProcessSnap   = NULL;
  67.    PROCESSENTRY32 pe32           = {0};
  68.    DWORD          dwTaskCount    = 0;
  69.     // Obtain a module handle to KERNEL so that we can get the addresses of
  70.     // the 32-bit Toolhelp functions we need.
  71.     hKernel = GetModuleHandle("KERNEL32.DLL");
  72.     if (hKernel)
  73.     {
  74.         pCreateToolhelp32Snapshot =
  75.           (CREATESNAPSHOT)GetProcAddress(hKernel, "CreateToolhelp32Snapshot");
  76.         pProcess32First = (PROCESSWALK)GetProcAddress(hKernel,
  77.                                                       "Process32First");
  78.         pProcess32Next  = (PROCESSWALK)GetProcAddress(hKernel,
  79.                                                       "Process32Next");
  80.     }
  81.     
  82.     // make sure we got addresses of all needed Toolhelp functions.
  83.     if (!(pProcess32First && pProcess32Next && pCreateToolhelp32Snapshot))
  84.        return 0;
  85.        
  86.     // Take a snapshot of all processes currently in the system.
  87.     hProcessSnap = (HINSTANCE) pCreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
  88.     if (hProcessSnap == (HANDLE)-1)
  89.         return 0;
  90.     // Walk the snapshot of processes and for each process, get information
  91.     // to display.
  92.     dwTaskCount = 0;
  93.     pe32.dwSize = sizeof(PROCESSENTRY32);   // must be filled out before use
  94.     if (pProcess32First(hProcessSnap, &pe32))
  95.     {
  96.         do
  97.         {
  98.             LPSTR pCurChar;
  99. //SAMPLE: use our class instead of raw structs
  100. CTaskListEntry* pEntry = new CTaskListEntry;
  101.             // strip path and leave executabe filename splitpath
  102.             for (pCurChar = (pe32.szExeFile + lstrlen (pe32.szExeFile));
  103.                  *pCurChar != '\' && pCurChar != pe32.szExeFile; 
  104.                  --pCurChar)
  105. //  ; // nothing
  106.             lstrcpy(pEntry->ProcessName, pCurChar);
  107.             pEntry->flags = 0;
  108.             pEntry->dwProcessId = pe32.th32ProcessID;
  109.             ++dwTaskCount;   // keep track of how many tasks we've got so far
  110. refList.Add(pEntry);
  111.         }
  112.         while (pProcess32Next(hProcessSnap, &pe32));
  113.     }
  114.     else
  115.         dwTaskCount = 0;    // Couldn't walk the list of processes.
  116.     // Don't forget to clean up the snapshot object...
  117.     CloseHandle (hProcessSnap);
  118.     return dwTaskCount;
  119. }
  120. DWORD
  121. GetTaskListNT(
  122. CTaskList& refList
  123.     )
  124. /*++
  125. Routine Description:
  126.     Provides an API for getting a list of tasks running at the time of the
  127.     API call.  This function uses the registry performance data to get the
  128.     task list and is therefore straight WIN32 calls that anyone can call.
  129. Arguments:
  130.     dwNumTasks       - maximum number of tasks that the pTask array can hold
  131. Return Value:
  132.     Number of tasks placed into the pTask array.
  133. --*/
  134. {
  135.     DWORD                        rc;
  136.     HKEY                         hKeyNames;
  137.     DWORD                        dwType;
  138.     DWORD                        dwSize;
  139.     LPSTR                        buf = NULL;
  140.     CHAR                         szSubKey[1024];
  141.     LANGID                       lid;
  142.     LPSTR                        p;
  143.     LPSTR                        p2;
  144.     PPERF_DATA_BLOCK             pPerf;
  145.     PPERF_OBJECT_TYPE            pObj;
  146.     PPERF_INSTANCE_DEFINITION    pInst;
  147.     PPERF_COUNTER_BLOCK          pCounter;
  148.     PPERF_COUNTER_DEFINITION     pCounterDef;
  149.     DWORD                        i;
  150.     DWORD                        dwProcessIdTitle;
  151.     DWORD                        dwProcessIdCounter;
  152.     CHAR                         szProcessName[MAX_PATH];
  153.     //
  154.     // Look for the list of counters.  Always use the neutral
  155.     // English version, regardless of the local language.  We
  156.     // are looking for some particular keys, and we are always
  157.     // going to do our looking in English.  We are not going
  158.     // to show the user the counter names, so there is no need
  159.     // to go find the corresponding name in the local language.
  160.     //
  161.     lid = MAKELANGID( LANG_ENGLISH, SUBLANG_NEUTRAL );
  162.     sprintf( szSubKey, "%s\%03x", REGKEY_PERF, lid );
  163.     rc = RegOpenKeyEx( HKEY_LOCAL_MACHINE,
  164.                        szSubKey,
  165.                        0,
  166.                        KEY_READ,
  167.                        &hKeyNames
  168.                      );
  169.     if (rc != ERROR_SUCCESS) {
  170.         goto exit;
  171.     }
  172.     //
  173.     // get the buffer size for the counter names
  174.     //
  175.     rc = RegQueryValueEx( hKeyNames,
  176.                           REGSUBKEY_COUNTERS,
  177.                           NULL,
  178.                           &dwType,
  179.                           NULL,
  180.                           &dwSize
  181.                         );
  182.     if (rc != ERROR_SUCCESS) {
  183.         goto exit;
  184.     }
  185.     //
  186.     // allocate the counter names buffer
  187.     //
  188.     buf = (LPSTR) malloc( dwSize );
  189.     if (buf == NULL)
  190.         goto exit;
  191.     memset( buf, 0, dwSize );
  192.     //
  193.     // read the counter names from the registry
  194.     //
  195.     rc = RegQueryValueEx( hKeyNames,
  196.                           REGSUBKEY_COUNTERS,
  197.                           NULL,
  198.                           &dwType,
  199.                           (LPBYTE) buf,
  200.                           &dwSize
  201.                         );
  202.     if (rc != ERROR_SUCCESS) {
  203.         goto exit;
  204.     }
  205.     //
  206.     // now loop thru the counter names looking for the following counters:
  207.     //
  208.     //      1.  "Process"           process name
  209.     //      2.  "ID Process"        process id
  210.     //
  211.     // the buffer contains multiple null terminated strings and then
  212.     // finally null terminated at the end.  the strings are in pairs of
  213.     // counter number and counter name.
  214.     //
  215.     p = buf;
  216.     while (*p) {
  217.         if (p > buf) {
  218.             for( p2=p-2; isdigit(*p2); p2--) ;
  219.             }
  220.         if (stricmp(p, PROCESS_COUNTER) == 0) {
  221.             //
  222.             // look backwards for the counter number
  223.             //
  224.             for( p2=p-2; isdigit(*p2); p2--) ;
  225.             strcpy( szSubKey, p2+1 );
  226.         }
  227.         else
  228.         if (stricmp(p, PROCESSID_COUNTER) == 0) {
  229.             //
  230.             // look backwards for the counter number
  231.             //
  232.             for( p2=p-2; isdigit(*p2); p2--) ;
  233.             dwProcessIdTitle = atol( p2+1 );
  234.         }
  235.         //
  236.         // next string
  237.         //
  238.         p += (strlen(p) + 1);
  239.     }
  240.     //
  241.     // free the counter names buffer
  242.     //
  243.     free( buf );
  244.     //
  245.     // allocate the initial buffer for the performance data
  246.     //
  247.     dwSize = INITIAL_SIZE;
  248.     buf = (LPSTR) malloc( dwSize );
  249.     if (buf == NULL) {
  250.         goto exit;
  251.     }
  252.     memset( buf, 0, dwSize );
  253.     while (TRUE) {
  254.         rc = RegQueryValueEx( HKEY_PERFORMANCE_DATA,
  255.                               szSubKey,
  256.                               NULL,
  257.                               &dwType,
  258.                               (LPBYTE) buf,
  259.                               &dwSize
  260.                             );
  261.         pPerf = (PPERF_DATA_BLOCK) buf;
  262.         //
  263.         // check for success and valid perf data block signature
  264.         //
  265.         if ((rc == ERROR_SUCCESS) &&
  266.             (dwSize > 0) &&
  267.             (pPerf)->Signature[0] == (WCHAR)'P' &&
  268.             (pPerf)->Signature[1] == (WCHAR)'E' &&
  269.             (pPerf)->Signature[2] == (WCHAR)'R' &&
  270.             (pPerf)->Signature[3] == (WCHAR)'F' ) {
  271.             break;
  272.         }
  273.         //
  274.         // if buffer is not big enough, reallocate and try again
  275.         //
  276.         if (rc == ERROR_MORE_DATA) {
  277.             dwSize += EXTEND_SIZE;
  278.             buf = (LPSTR) realloc( buf, dwSize );
  279.             memset( buf, 0, dwSize );
  280.         }
  281.         else {
  282.             goto exit;
  283.         }
  284.     }
  285.     //
  286.     // set the perf_object_type pointer
  287.     //
  288.     pObj = (PPERF_OBJECT_TYPE) ((DWORD)pPerf + pPerf->HeaderLength);
  289.     //
  290.     // loop thru the performance counter definition records looking
  291.     // for the process id counter and then save its offset
  292.     //
  293.     pCounterDef = (PPERF_COUNTER_DEFINITION) ((DWORD)pObj + pObj->HeaderLength);
  294.     for (i=0; i<(DWORD)pObj->NumCounters; i++) {
  295.         if (pCounterDef->CounterNameTitleIndex == dwProcessIdTitle) {
  296.             dwProcessIdCounter = pCounterDef->CounterOffset;
  297.             break;
  298.         }
  299.         pCounterDef++;
  300.     }
  301.     pInst = (PPERF_INSTANCE_DEFINITION) ((DWORD)pObj + pObj->DefinitionLength);
  302.     //
  303.     // loop thru the performance instance data extracting each process name
  304.     // and process id
  305.     //
  306.     for (i=0; i < (DWORD)pObj->NumInstances; i++)
  307. {
  308. CTaskListEntry* pEntry = new CTaskListEntry;
  309.         //
  310.         // pointer to the process name
  311.         //
  312.         p = (LPSTR) ((DWORD)pInst + pInst->NameOffset);
  313.         //
  314.         // convert it to ascii
  315.         //
  316.         rc = WideCharToMultiByte( CP_ACP,
  317.                                   0,
  318.                                   (LPCWSTR)p,
  319.                                   -1,
  320.                                   szProcessName,
  321.                                   sizeof(szProcessName),
  322.                                   NULL,
  323.                                   NULL
  324.                                 );
  325.         if (!rc) {
  326.             //
  327.     // if we cant convert the string then use a default value
  328.             //
  329.             strcpy( pEntry->ProcessName, UNKNOWN_TASK );
  330.         }
  331.         if (strlen(szProcessName)+4 <= sizeof(pEntry->ProcessName)) {
  332.             strcpy( pEntry->ProcessName, szProcessName );
  333.             strcat( pEntry->ProcessName, ".exe" );
  334.         }
  335.         //
  336.         // get the process id
  337.         //
  338.         pCounter = (PPERF_COUNTER_BLOCK) ((DWORD)pInst + pInst->ByteLength);
  339.         pEntry->flags = 0;
  340.         pEntry->dwProcessId = *((LPDWORD) ((DWORD)pCounter + dwProcessIdCounter));
  341.         if (pEntry->dwProcessId == 0) {
  342.             pEntry->dwProcessId = (DWORD)-2;
  343.         }
  344.         //
  345.         // next process
  346.         //
  347. refList.Add(pEntry);
  348.         pInst = (PPERF_INSTANCE_DEFINITION) ((DWORD)pCounter + pCounter->ByteLength);
  349.     }
  350. exit:
  351.     if (buf) {
  352.         free( buf );
  353.     }
  354.     RegCloseKey( hKeyNames );
  355.     RegCloseKey( HKEY_PERFORMANCE_DATA );
  356.     return (DWORD)pObj->NumInstances;
  357. }
  358. BOOL
  359. EnableDebugPriv95(
  360.     VOID
  361.     )
  362. /*++
  363. Routine Description:
  364.     Changes the process's privilege so that kill works properly.
  365. Arguments:
  366. Return Value:
  367.     TRUE             - success
  368.     FALSE            - failure
  369. Comments: 
  370.     Always returns TRUE
  371. --*/
  372. {
  373.    return TRUE;
  374. }
  375. BOOL
  376. EnableDebugPrivNT(
  377.     VOID
  378.     )
  379. /*++
  380. Routine Description:
  381.     Changes the process's privilege so that kill works properly.
  382. Arguments:
  383. Return Value:
  384.     TRUE             - success
  385.     FALSE            - failure
  386. --*/
  387. {
  388.     HANDLE hToken;
  389.     LUID DebugValue;
  390.     TOKEN_PRIVILEGES tkp;
  391.     //
  392.     // Retrieve a handle of the access token
  393.     //
  394.     if (!OpenProcessToken(GetCurrentProcess(),
  395.             TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
  396.             &hToken)) {
  397.         printf("OpenProcessToken failed with %dn", GetLastError());
  398.         return FALSE;
  399.     }
  400.     //
  401.     // Enable the SE_DEBUG_NAME privilege
  402.     //
  403.     if (!LookupPrivilegeValue((LPSTR) NULL,
  404.             SE_DEBUG_NAME,
  405.             &DebugValue)) {
  406.         printf("LookupPrivilegeValue failed with %dn", GetLastError());
  407.         return FALSE;
  408.     }
  409.     tkp.PrivilegeCount = 1;
  410.     tkp.Privileges[0].Luid = DebugValue;
  411.     tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
  412.     AdjustTokenPrivileges(hToken,
  413.         FALSE,
  414.         &tkp,
  415.         sizeof(TOKEN_PRIVILEGES),
  416.         (PTOKEN_PRIVILEGES) NULL,
  417.         (PDWORD) NULL);
  418.     //
  419.     // The return value of AdjustTokenPrivileges can't be tested
  420.     //
  421.     if (GetLastError() != ERROR_SUCCESS) {
  422.         printf("AdjustTokenPrivileges failed with %dn", GetLastError());
  423.         return FALSE;
  424.     }
  425.     return TRUE;
  426. }
  427. BOOL KillProcess(CTaskListEntry* pEntry, BOOL fForce)
  428. {
  429.     HANDLE            hProcess;
  430.     if (fForce || !pEntry->hwnd)
  431. {
  432.         hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pEntry->dwProcessId );
  433.         if (hProcess)
  434. {
  435.             hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, pEntry->dwProcessId );
  436.             if (hProcess == NULL)
  437. {
  438.                 return FALSE;
  439.             }
  440.             if (!TerminateProcess( hProcess, 1 )) 
  441. {
  442.                 CloseHandle( hProcess );
  443.                 return FALSE;
  444.             }
  445.             CloseHandle( hProcess );
  446.             return TRUE;
  447.         }
  448.     }
  449.     //
  450.     // kill the process
  451.     //
  452.     PostMessage(pEntry->hwnd, WM_CLOSE, 0, 0);
  453.     return TRUE;
  454. }
  455. VOID GetWindowTitles(CTaskList& refList)
  456. {
  457.     //
  458.     // enumerate all windows
  459.     //
  460.     EnumWindows( EnumWindowsProc, (LPARAM) &refList );
  461. }
  462. BOOL CALLBACK
  463. EnumWindowsProc(
  464.     HWND    hwnd,
  465.     LPARAM   lParam
  466.     )
  467. {
  468.     DWORD             pid = 0;
  469.     int               i;
  470.     CHAR              buf[TITLE_SIZE];
  471. //SAMPLE: superfunky cast! don't do this yourself
  472. // unless you know what it's for: stupid casts cause bugs.
  473. CTaskList&   refList = *((CTaskList*) lParam);
  474.     //
  475.     // get the processid for this window
  476.     //
  477.     if (!GetWindowThreadProcessId( hwnd, &pid )) {
  478.         return TRUE;
  479.     }
  480.     //
  481.     // look for the task in the task list for this window
  482.     //
  483.     for (i=0; i < refList.GetSize(); i++)
  484. {
  485.         if (refList[i]->dwProcessId == pid)
  486. {
  487.             refList[i]->hwnd = hwnd;
  488.             //
  489. // we found the task so lets try to get the
  490.             // window text
  491.             //
  492.             if (GetWindowText( refList[i]->hwnd, buf, sizeof(buf) ))
  493. {
  494.                 //
  495. // got it, so lets save it
  496.                 //
  497.                 strcpy( refList[i]->WindowTitle, buf);
  498.             }
  499.             break;
  500.         }
  501.     }
  502.     //
  503.     // continue the enumeration
  504.     //
  505.     return TRUE;
  506. }