killprocess.cpp

teamviewer source code vc++

// Kill a process by name
// by Eugene Polonsky  
// NT4 requires separate handling, because it doesn't support the ToolHelp32 API 
// for more details see http://www.codeproject.com/threads/killprocess.asp

#include "stdhdrs.h"
#include "KillProcess.h"

CKillProcess::CKillProcess() :
	FCreateToolhelp32Snapshot(NULL), 
    FProcess32First(NULL), FProcess32Next(NULL), 
    m_hKernelLib(NULL),
    m_hNTLib(NULL)
{
	FHeapFree = NULL;
	FHeapAlloc = NULL;
	FQuerySysInfo = NULL;
	FGetProcessHeap = NULL;

    m_hKernelLib = ::LoadLibraryA("Kernel32");
    if (m_hKernelLib)  
    {
        // Find ToolHelp functions
        FCreateToolhelp32Snapshot = 
            (PFCreateToolhelp32Snapshot)
            ::GetProcAddress(m_hKernelLib,
            _TEXT("CreateToolhelp32Snapshot"));
        FProcess32First = (PFProcess32First)
            ::GetProcAddress(m_hKernelLib, 
            _TEXT("Process32First"));
        FProcess32Next = (PFProcess32Next)
            ::GetProcAddress(m_hKernelLib, 
            _TEXT("Process32Next"));
    }
    if(!FCreateToolhelp32Snapshot || 
        !FProcess32First || !FProcess32Next)
    { // i.e. we couldn't find the ToolHelp functions, 
        //so we must be on NT4. Let's load the
        // undocumented one instead.
        if(!m_hKernelLib)
            return; // can't do anything at all without 
        //the kernel.

        m_hNTLib = ::LoadLibraryA("ntdll.dll");
        if(m_hNTLib)
        {
            FQuerySysInfo = 
                (PFZwQuerySystemInformation)
                ::GetProcAddress(m_hNTLib, 
                _TEXT("ZwQuerySystemInformation"));
            // load some support funcs from the kernel
            FGetProcessHeap = (PFGetProcessHeap)
                ::GetProcAddress(m_hKernelLib, 
                _TEXT("GetProcessHeap"));
            FHeapAlloc = (PFHeapAlloc)
                ::GetProcAddress(m_hKernelLib, 
                _TEXT("HeapAlloc"));
            FHeapFree = (PFHeapFree)
                ::GetProcAddress(m_hKernelLib, 
                _TEXT("HeapFree"));
        }
    }
}
       
CKillProcess::~CKillProcess()
{
    if(m_hKernelLib)
        FreeLibrary(m_hKernelLib);
    if(m_hNTLib)
        FreeLibrary(m_hNTLib);
}

bool CKillProcess::KillProcess(IN const char* pstrProcessName)
{
    DWORD dwId;
    HANDLE hProcess = FindProcess(pstrProcessName, 
        dwId);
    BOOL bResult;
    if(!hProcess)
        return false;

    // TerminateAppEnum() posts WM_CLOSE to all windows whose PID
    // matches your process's.
    ::EnumWindows((WNDENUMPROC)
        CKillProcess::TerminateAppEnum, 
        (LPARAM) dwId);
    // Wait on the handle. If it signals, great. 
    //If it times out, then you kill it.
    if(WaitForSingleObject(hProcess, 1000)
        !=WAIT_OBJECT_0)
        bResult = TerminateProcess(hProcess,0);
    else
        bResult = TRUE; 
    CloseHandle(hProcess);
    return bResult == TRUE;
}

HANDLE CKillProcess::FindProcess(IN const char* pstrProcessName, OUT DWORD& dwId)
{
    if(!m_hKernelLib)
        return NULL;

    if(FCreateToolhelp32Snapshot && FProcess32First && 
        FProcess32Next) // use toolhelpapi
        return THFindProcess(pstrProcessName, dwId);
    if(FQuerySysInfo && FHeapAlloc && 
        FGetProcessHeap && FHeapFree) // use NT api
        return NTFindProcess(pstrProcessName, dwId);
    // neither one got loaded. Strange.
    return NULL;
}

HANDLE CKillProcess::THFindProcess(IN const char* pstrProcessName, OUT DWORD& dwId)
{
    HANDLE            hSnapShot=NULL;
    HANDLE            hResult = NULL;
    PROCESSENTRY32    processInfo;
    char*            pstrExeName;

    bool bFirst = true;
    ::ZeroMemory(&processInfo, sizeof(PROCESSENTRY32));
    processInfo.dwSize = sizeof(PROCESSENTRY32);
    hSnapShot = FCreateToolhelp32Snapshot(
        TH32CS_SNAPPROCESS, 0);
    if(hSnapShot == INVALID_HANDLE_VALUE)
        return NULL; 

    // ok now let's iterate with Process32Next until we 
    // match up the name of our process
    while((bFirst ? FProcess32First(hSnapShot, 
        &processInfo) : FProcess32Next(hSnapShot, 
        &processInfo)))
    {
        bFirst = false;
        // we need to check for path... and extract 
        // just the exe name
        pstrExeName = strrchr(processInfo.szExeFile, 
            '\\');
        if(!pstrExeName)
            pstrExeName = processInfo.szExeFile;
        else
            pstrExeName++; // skip the \
        // ok now compare against our process name
        if(stricmp(pstrExeName, pstrProcessName) == 0) 
            // wee weee we found it
        {
            // let's get a HANDLE on it
            hResult=OpenProcess(
                SYNCHRONIZE|PROCESS_TERMINATE, TRUE, 
                processInfo.th32ProcessID);
            dwId = processInfo.th32ProcessID;
            break;
        }
    } // while(Process32Next(hSnapShot, &processInfo){
    if(hSnapShot)
        CloseHandle(hSnapShot);
    return hResult;
}

HANDLE CKillProcess::NTFindProcess(IN const char* pstrProcessName, OUT DWORD& dwId)
{
    HANDLE hHeap = FGetProcessHeap();
    NTSTATUS Status;
    ULONG cbBuffer = 0x8000;
    PVOID pBuffer = NULL;
    HANDLE hResult = NULL;
    // it is difficult to say a priory which size of 
    // the buffer will be enough to retrieve all 
    // information, so we startwith 32K buffer and 
    // increase its size until we get the
    // information successfully
    do
    {
        pBuffer = HeapAlloc(hHeap, 0, cbBuffer);
        if (pBuffer == NULL)
            return SetLastError(
            ERROR_NOT_ENOUGH_MEMORY), NULL;

        Status = FQuerySysInfo(
            SystemProcessesAndThreadsInformation,
            pBuffer, cbBuffer, NULL);

        if (Status == STATUS_INFO_LENGTH_MISMATCH)
        {
            HeapFree(hHeap, 0, pBuffer);
            cbBuffer *= 2;
        }
        else if (!NT_SUCCESS(Status))
        {
            HeapFree(hHeap, 0, pBuffer);
            return SetLastError(Status), NULL;
        }
    }
    while (Status == STATUS_INFO_LENGTH_MISMATCH);

    PSYSTEM_PROCESSES pProcesses = 
        (PSYSTEM_PROCESSES)pBuffer;

    for (;;)
    {
        PCWSTR pszProcessName = 
            pProcesses->ProcessName.Buffer;
        if (pszProcessName == NULL)
            pszProcessName = L"Idle";

        CHAR szProcessName[MAX_PATH];
        WideCharToMultiByte(CP_ACP, 0, pszProcessName, 
            -1,szProcessName, MAX_PATH, NULL, NULL);

        if(stricmp(szProcessName, pstrProcessName) 
            == 0) // found it
        {
            hResult=OpenProcess(
                SYNCHRONIZE|PROCESS_TERMINATE, TRUE, 
                pProcesses->ProcessId);
            dwId = pProcesses->ProcessId;
            break;
        }

        if (pProcesses->NextEntryDelta == 0)
            break;

        // find the address of the next process 
        // structure
        pProcesses = (PSYSTEM_PROCESSES)(
            ((LPBYTE)pProcesses)
            + pProcesses->NextEntryDelta);
    }

    HeapFree(hHeap, 0, pBuffer);
    return hResult;
}

// callback function for window enumeration
BOOL CALLBACK CKillProcess::TerminateAppEnum( HWND hwnd, LPARAM lParam )
{
    DWORD dwID ;

    GetWindowThreadProcessId(hwnd, &dwID) ;

    if(dwID == (DWORD)lParam)
    {
        PostMessage(hwnd, WM_CLOSE, 0, 0) ;
    }

    return TRUE ;
}