extended process enumeration

src/BlackBone/Include/FunctionTypes.h

@@ -225,6 +225,9 @@ typedef decltype(&RtlInitUnicodeString) fnRtlInitUnicodeString;
 // RtlFreeUnicodeString
 typedef decltype(&RtlFreeUnicodeString) fnRtlFreeUnicodeString;
 
+// NtQuerySystemInformation
+typedef decltype(&NtQuerySystemInformation) fnNtQuerySystemInformation;
+
 // NtQueryInformationProcess
 typedef decltype(&NtQueryInformationProcess) fnNtQueryInformationProcess;
 

src/BlackBone/Include/NativeStructures.h

@@ -35,15 +35,8 @@ namespace blackbone
     template <typename T>
     struct _UNICODE_STRING_T
     {
-        union
-        {
-            struct
-            {
-                WORD Length;
-                WORD MaximumLength;
-            };
-            T dummy;
-        };
+        WORD Length;
+        WORD MaximumLength;
         T Buffer;
     };
 
@@ -318,18 +311,78 @@ namespace blackbone
     template<typename T>
     struct _THREAD_BASIC_INFORMATION_T
     {
-        NTSTATUS  ExitStatus;
-        T         TebBaseAddress;
+        NTSTATUS ExitStatus;
+        T TebBaseAddress;
+        _CLIENT_ID_T<T> ClientID;
+        T AffinityMask;
+        LONG Priority;
+        LONG BasePriority;
+    };
 
-        struct
-        {
-            T  UniqueProcess;
-            T  UniqueThread;
-        } ClientId;
+    struct VM_COUNTERS
+    {
+        SIZE_T PeakVirtualSize;
+        SIZE_T VirtualSize;
+        ULONG PageFaultCount;
+        SIZE_T PeakWorkingSetSize;
+        SIZE_T WorkingSetSize;
+        SIZE_T QuotaPeakPagedPoolUsage;
+        SIZE_T QuotaPagedPoolUsage;
+        SIZE_T QuotaPeakNonPagedPoolUsage;
+        SIZE_T QuotaNonPagedPoolUsage;
+        SIZE_T PagefileUsage;
+        SIZE_T PeakPagefileUsage;
+    };
 
-        T      AffinityMask;
-        LONG   Priority;
-        LONG   BasePriority;
+    template<typename T>
+    struct _SYSTEM_THREAD_INFORMATION_T
+    {
+        LARGE_INTEGER KernelTime;
+        LARGE_INTEGER UserTime;
+        LARGE_INTEGER CreateTime;
+        ULONG WaitTime;
+        T StartAddress;
+        _CLIENT_ID_T<T> ClientId;
+        LONG Priority;
+        LONG BasePriority;
+        ULONG ContextSwitches;
+        ULONG ThreadState;
+        ULONG WaitReason;
+    };
+
+    template<typename T>
+    struct _SYSTEM_EXTENDED_THREAD_INFORMATION_T
+    {
+        _SYSTEM_THREAD_INFORMATION_T<T> ThreadInfo;
+        T StackBase;
+        T StackLimit;
+        T Win32StartAddress;
+        T Reserved[4];
+    };
+
+    template<typename T>
+    struct _SYSTEM_PROCESS_INFORMATION_T
+    {
+        ULONG NextEntryOffset;
+        ULONG NumberOfThreads;
+        LARGE_INTEGER WorkingSetPrivateSize;
+        ULONG HardFaultCount;
+        ULONG NumberOfThreadsHighWatermark;
+        ULONGLONG CycleTime;
+        LARGE_INTEGER CreateTime;
+        LARGE_INTEGER UserTime;
+        LARGE_INTEGER KernelTime;
+        _UNICODE_STRING_T<T> ImageName;
+        LONG BasePriority;
+        T UniqueProcessId;
+        T InheritedFromUniqueProcessId;
+        ULONG HandleCount;
+        ULONG SessionId;
+        ULONG_PTR UniqueProcessKey;
+        VM_COUNTERS VmCounters;
+        SIZE_T PrivatePageCount;
+        IO_COUNTERS IoCounters;
+        _SYSTEM_EXTENDED_THREAD_INFORMATION_T<T> Threads[1];
     };
 
     struct _XSAVE_FORMAT64

src/BlackBone/Misc/NameResolve.cpp

@@ -13,9 +13,12 @@ namespace blackbone
 
 NameResolve::NameResolve()
 {
-    DynImport::load( "RtlDosApplyFileIsolationRedirection_Ustr", L"ntdll.dll" );
-    DynImport::load( "RtlInitUnicodeString", L"ntdll.dll" );
-    DynImport::load( "RtlFreeUnicodeString", L"ntdll.dll" );
+    HMODULE ntdll = GetModuleHandleW( L"ntdll.dll" );
+
+    DynImport::load( "NtQuerySystemInformation", ntdll );
+    DynImport::load( "RtlDosApplyFileIsolationRedirection_Ustr", ntdll );
+    DynImport::load( "RtlInitUnicodeString", ntdll );
+    DynImport::load( "RtlFreeUnicodeString", ntdll );
 }
 
 NameResolve::~NameResolve()

src/BlackBone/Process/Process.cpp

@@ -1,5 +1,6 @@
 #include "Process.h"
 #include "../Misc/NameResolve.h"
+#include "../Misc/DynImport.h"
 
 #include <memory>
 
@@ -18,9 +19,8 @@ Process::Process()
 {
     GrantPriviledge( SE_DEBUG_NAME );
     GrantPriviledge( SE_LOAD_DRIVER_NAME );
-    GrantPriviledge( SE_LOCK_MEMORY_NAME );
 
-    NameResolve::Instance().Initialize();
+    NameResolve::Instance().Initialize(); 
 }
 
 Process::~Process(void)
@@ -235,5 +235,95 @@ void Process::EnumByName( const std::wstring& name, std::vector<DWORD>& found )
     }
 }
 
+/// <summary>
+/// Search for process by executable name or by process ID
+/// </summary>
+/// <param name="pid">Target process ID. rocess name. If empty - function will retrieve all existing processes</param>
+/// <param name="name">Process executable name. If empty - function will retrieve all existing processes</param>
+/// <param name="found">Found processses</param>
+/// <param name="includeThreads">If set to true, function will retrieve info ablout process threads</param>
+/// <returns>Status code</returns>
+NTSTATUS Process::EnumByNameOrPID(
+    uint32_t pid,
+    const std::wstring& name,
+    std::vector<ProcessInfo>& found,
+    bool includeThreads /*= false*/
+    )
+{
+    ULONG bufSize = 0x100;
+    uint8_t tmpbuf[0x100];
+    uint8_t* buffer = tmpbuf;
+    ULONG returnLength = 0;
+
+    found.clear();
+
+    // Query process info
+    NTSTATUS status = GET_IMPORT( NtQuerySystemInformation )((SYSTEM_INFORMATION_CLASS)57, buffer, bufSize, &returnLength);
+    if (!NT_SUCCESS( status ))
+    {
+        bufSize = returnLength;
+        buffer = (uint8_t*)VirtualAlloc( NULL, bufSize, MEM_COMMIT, PAGE_READWRITE );
+        status = GET_IMPORT( NtQuerySystemInformation )((SYSTEM_INFORMATION_CLASS)57, buffer, bufSize, &returnLength);
+        if (!NT_SUCCESS( status ))
+        {
+            VirtualFree( buffer, 0, MEM_RELEASE );
+            return status;
+        }
+    }
+
+    // Parse info
+    for (auto pInfo = reinterpret_cast<_SYSTEM_PROCESS_INFORMATION_T<DWORD_PTR>*>(buffer);;)
+    {
+        //  Skip idle process, compare name or compare pid
+        if (pInfo->UniqueProcessId != 0 && (
+            name.empty() && pid == 0 ||
+            _wcsicmp( name.c_str(), (wchar_t*)pInfo->ImageName.Buffer ) == 0 ||
+            pid == pInfo->UniqueProcessId))
+        {
+            ProcessInfo info;
+
+            info.pid = static_cast<uint32_t>(pInfo->UniqueProcessId);
+
+            if (pInfo->ImageName.Buffer)
+                info.imageName = reinterpret_cast<wchar_t*>(pInfo->ImageName.Buffer);
+
+            // Get threads info
+            if (includeThreads)
+            {
+                int64_t minTime = 0xFFFFFFFFFFFFFFFF;
+                ULONG mainIdx = 0;
+
+                for (ULONG i = 0; i < pInfo->NumberOfThreads; i++)
+                {
+                    ThreadInfo tinfo;
+                    auto& thd = pInfo->Threads[i];
+
+                    tinfo.tid = static_cast<uint32_t>(thd.ThreadInfo.ClientId.UniqueThread);
+                    tinfo.startAddress = static_cast<uintptr_t>(thd.ThreadInfo.StartAddress);
+
+                    // Check for main thread
+                    if (thd.ThreadInfo.CreateTime.QuadPart < minTime)
+                    {
+                        minTime = thd.ThreadInfo.CreateTime.QuadPart;
+                        mainIdx = i;
+                    }
+
+                    info.threads.emplace_back( tinfo );
+                    info.threads[mainIdx].mainThread = true;
+                }
+            }
+
+            found.emplace_back( info );
+        }
+
+        if (pInfo->NextEntryOffset)
+            pInfo = reinterpret_cast<_SYSTEM_PROCESS_INFORMATION_T<DWORD_PTR>*>((uint8_t*)pInfo + pInfo->NextEntryOffset);
+        else
+            break;        
+    }
+
+    VirtualFree( buffer, 0, MEM_RELEASE );
+    return status;
+}
 
 }

src/BlackBone/Process/Process.h

@@ -17,6 +17,21 @@
 namespace blackbone
 {
 
+struct ThreadInfo
+{
+    uint32_t tid = 0;
+    uintptr_t startAddress = 0;
+    bool mainThread = false;
+};
+
+struct ProcessInfo
+{
+    uint32_t pid = 0;
+    std::wstring imageName;
+    std::vector<ThreadInfo> threads;
+};
+
+
 #define DEFAULT_ACCESS_P  PROCESS_QUERY_INFORMATION | \
                           PROCESS_VM_READ           | \
                           PROCESS_VM_WRITE          | \
@@ -97,6 +112,21 @@ class Process
     /// <param name="found">Found processses</param>
     BLACKBONE_API static void EnumByName( const std::wstring& name, std::vector<DWORD>& found );
 
+    /// <summary>
+    /// Search for process by executable name or by process ID
+    /// </summary>
+    /// <param name="pid">Target process ID. rocess name. If empty - function will retrieve all existing processes</param>
+    /// <param name="name">Process executable name. If empty - function will retrieve all existing processes</param>
+    /// <param name="found">Found processses</param>
+    /// <param name="includeThreads">If set to true, function will retrieve info ablout process threads</param>
+    /// <returns>Status code</returns>
+    BLACKBONE_API static NTSTATUS EnumByNameOrPID( 
+        uint32_t pid,
+        const std::wstring& name, 
+        std::vector<ProcessInfo>& found, 
+        bool includeThreads = false
+        );
+
     //
     // Subroutines
     //

src/BlackBoneDrv/Inject.c

@@ -467,17 +467,21 @@ NTSTATUS BBLookupProcessThread( IN HANDLE pid, OUT PETHREAD* ppThread )
         return status;
     }
 
-    // Find a target thread
+    // Find target thread
     if (NT_SUCCESS( status ))
     {
         status = STATUS_NOT_FOUND;
-        for (; pInfo->NextEntryOffset; pInfo = (PSYSTEM_PROCESS_INFO)((PUCHAR)pInfo + pInfo->NextEntryOffset))
+        for (;;)
         {
             if (pInfo->UniqueProcessId == pid)
             {
                 status = STATUS_SUCCESS;
                 break;
             }
+            else if (pInfo->NextEntryOffset)
+                pInfo = (PSYSTEM_PROCESS_INFO)((PUCHAR)pInfo + pInfo->NextEntryOffset);
+            else
+                break;
         }
     }