threads_util.cpp

Go to the documentation of this file.

#include "threads_util.h"
 
#include <peconv.h>
#include <tlhelp32.h>
#include "../utils/ntddk.h"
 
#ifdef _DEBUG
#include <iostream>
#endif
 
bool pesieve::util::fetch_threads_info(DWORD pid, std::vector<thread_info>& threads_info)
{
    BYTE* buffer = nullptr;
    ULONG buffer_size = 0;
    ULONG ret_len = 0;
 
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    while (status != STATUS_SUCCESS) {
        status = NtQuerySystemInformation(SystemProcessInformation, buffer, buffer_size, &ret_len);
        if (status == STATUS_INFO_LENGTH_MISMATCH) {
            free(buffer);
            buffer = nullptr;
            buffer_size = 0;
            buffer = (BYTE*)calloc(ret_len, 1);
            if (!buffer) {
                return false;
            }
            buffer_size = ret_len;
            continue; // try again
        }
        break; //other error, or success
    };
 
    if (status != STATUS_SUCCESS) {
        free(buffer);
        return false;
    }
 
    bool found = false;
    SYSTEM_PROCESS_INFORMATION* info = (SYSTEM_PROCESS_INFORMATION*)buffer;
    while (info) {
        if (info->UniqueProcessId == pid) {
            found = true;
            break;
        }
        if (!info->NextEntryOffset) {
            break;
        }
        size_t record_size = info->NextEntryOffset;
        if (record_size < sizeof(SYSTEM_PROCESS_INFORMATION)) {
            // Record size smaller than expected, probably it is an old system that doesn not support the new version of this API
#ifdef _DEBUG
            std::cout << "The new version of SYSTEM_PROCESS_INFORMATION is not supported!\n";
#endif
            break;
        }
        info = (SYSTEM_PROCESS_INFORMATION*)((ULONG_PTR)info + info->NextEntryOffset);
        if (!peconv::validate_ptr(buffer, buffer_size, info, sizeof(SYSTEM_PROCESS_INFORMATION))) {
            break;
        }
    }
 
    if (!found) {
        free(buffer);
        return false;
    }
 
    size_t thread_count = info->NumberOfThreads;
    for (size_t i = 0; i < thread_count; i++) {
        thread_info threadi;
 
        threadi.tid = MASK_TO_DWORD((ULONGLONG)info->Threads[i].ClientId.UniqueThread);
        threadi.is_extended = true;
        threadi.ext.start_addr = (ULONG_PTR)info->Threads[i].StartAddress;
        threadi.ext.state = info->Threads[i].ThreadState;
        threadi.ext.wait_reason = info->Threads[i].WaitReason;
        threadi.ext.wait_time  = info->Threads[i].WaitTime;
        threads_info.push_back(threadi);
    }
 
    free(buffer);
    return true;
}
 
bool pesieve::util::fetch_threads_by_snapshot(DWORD pid, std::vector<thread_info>& threads_info)
{
    HANDLE hThreadSnapShot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if (hThreadSnapShot == INVALID_HANDLE_VALUE) {
        const DWORD err = GetLastError();
#ifdef _DEBUG
        std::cerr << "[-] Could not create threads snapshot. Error: " << std::dec << err << std::endl;
#endif
        return false;
    }
    THREADENTRY32 th32 = { 0 };
    th32.dwSize = sizeof(THREADENTRY32);
 
    //check all threads in the process:
    if (!Thread32First(hThreadSnapShot, &th32)) {
        CloseHandle(hThreadSnapShot);
#ifdef _DEBUG
        std::cerr << "[-] Could not enumerate thread. Error: " << GetLastError() << std::endl;
#endif
        return false;
    }
    do {
        if (th32.th32OwnerProcessID != pid) {
            continue;
        }
 
        thread_info threadi;
        threadi.tid = th32.th32ThreadID;
        threadi.is_extended = false;
        threads_info.push_back(threadi);
 
    } while (Thread32Next(hThreadSnapShot, &th32));
 
    CloseHandle(hThreadSnapShot);
    return true;
}