NtLib.hpp

#include <ntstatus.h>

#define NOMINMAX
#define WIN32_NO_STATUS
#define WIN32_LEAN_AND_MEAN
#include <Windows.h>

#include <winternl.h>
#pragma comment(lib, "ntdll.lib")

namespace Nt
{
    template <typename Allocator>
    concept AllocatorTrait = requires (Allocator & allocator)
    {
        { allocator.resize(size_t{}) };
        { allocator.size() };
        { allocator.clear() };
        { allocator.data() };
        { allocator.empty() };
    };

    class DefaultAllocator
    {
    private:
        void* m_buf;
        size_t m_size;
        size_t m_reserved;

    public:
        DefaultAllocator() : m_buf(nullptr), m_size(0), m_reserved(0)
        {
        }

        ~DefaultAllocator()
        {
            clear();
        }

        DefaultAllocator(const DefaultAllocator& allocator) : DefaultAllocator()
        {
            resize(allocator.size());
            if (size() == allocator.size())
            {
                memcpy(m_buf, allocator.data(), allocator.size());
            }
        }

        DefaultAllocator(DefaultAllocator&& allocator) noexcept : m_buf(allocator.m_buf), m_size(allocator.m_size), m_reserved(allocator.m_reserved)
        {
            allocator.m_buf = nullptr;
            allocator.m_size = 0;
            allocator.m_reserved = 0;
        }

        DefaultAllocator& operator = (const DefaultAllocator& allocator)
        {
            if (&allocator == this)
            {
                return *this;
            }

            resize(allocator.size());
            if (size() == allocator.size())
            {
                memcpy(m_buf, allocator.data(), allocator.size());
            }

            return *this;
        }

        DefaultAllocator& operator = (DefaultAllocator&& allocator) noexcept
        {
            if (&allocator == this)
            {
                return *this;
            }

            clear();

            m_buf = allocator.m_buf;
            m_size = allocator.m_size;
            m_reserved = allocator.m_reserved;

            allocator.m_buf = nullptr;
            allocator.m_size = 0;
            allocator.m_reserved = 0;

            return *this;
        }

        void resize(const size_t size)
        {
            if (!size)
            {
                clear();
                return;
            }

            if (m_buf)
            {
                if (size <= m_reserved)
                {
                    m_size = size;
                    return;
                }

                clear();
            }

            m_buf = new unsigned char[size];
            if (!m_buf)
            {
                return;
            }

            m_size = size;
            m_reserved = size;
        }

        size_t size() const
        {
            return m_size;
        }

        void clear()
        {
            if (m_buf)
            {
                delete[] static_cast<unsigned char*>(m_buf);
                m_buf = nullptr;
                m_size = 0;
                m_reserved = 0;
            }
        }

        void* data()
        {
            return m_buf;
        }

        const void* data() const
        {
            return m_buf;
        }

        bool empty() const
        {
            return !m_buf || !m_size;
        }
    };

    namespace Api
    {
        enum class MEMORY_INFORMATION_CLASS
        {
            MemoryBasicInformation
        };

        extern "C" NTSYSAPI NTSTATUS NTAPI ZwQueryVirtualMemory(
            IN HANDLE hProcess,
            IN PVOID baseAddress,
            IN MEMORY_INFORMATION_CLASS infoClass,
            OUT PVOID buf,
            IN SIZE_T length,
            OUT OPTIONAL PSIZE_T resultLength
        );



        enum class KTHREAD_STATE : ULONG
        {
            Initialized,
            Ready,
            Running,
            Standby,
            Terminated,
            Waiting,
            Transition,
            DeferredReady,
            GateWaitObsolete,
            WaitingForProcessInSwap,
            MaximumThreadState
        };

        enum class KWAIT_REASON : ULONG
        {
            Executive,
            FreePage,
            PageIn,
            PoolAllocation,
            DelayExecution,
            Suspended,
            UserRequest,
            WrExecutive,
            WrFreePage,
            WrPageIn,
            WrPoolAllocation,
            WrDelayExecution,
            WrSuspended,
            WrUserRequest,
            WrEventPair,
            WrQueue,
            WrLpcReceive,
            WrLpcReply,
            WrVirtualMemory,
            WrPageOut,
            WrRendezvous,
            WrKeyedEvent,
            WrTerminated,
            WrProcessInSwap,
            WrCpuRateControl,
            WrCalloutStack,
            WrKernel,
            WrResource,
            WrPushLock,
            WrMutex,
            WrQuantumEnd,
            WrDispatchInt,
            WrPreempted,
            WrYieldExecution,
            WrFastMutex,
            WrGuardedMutex,
            WrRundown,
            WrAlertByThreadId,
            WrDeferredPreempt,
            MaximumWaitReason
        };

        struct SYSTEM_THREAD_INFORMATION
        {
            LARGE_INTEGER KernelTime;
            LARGE_INTEGER UserTime;
            LARGE_INTEGER CreateTime;
            ULONG WaitTime;
            PVOID StartAddress;
            CLIENT_ID ClientId;
            KPRIORITY Priority;
            LONG BasePriority;
            ULONG ContextSwitches;
            KTHREAD_STATE ThreadState;
            KWAIT_REASON WaitReason;
        };

        struct SYSTEM_PROCESS_INFORMATION
        {
            ULONG NextEntryOffset;
            ULONG NumberOfThreads;
            LARGE_INTEGER SpareLi1;
            LARGE_INTEGER SpareLi2;
            LARGE_INTEGER SpareLi3;
            LARGE_INTEGER CreateTime;
            LARGE_INTEGER UserTime;
            LARGE_INTEGER KernelTime;
            UNICODE_STRING ImageName;
            KPRIORITY BasePriority;
            HANDLE UniqueProcessId;
            HANDLE InheritedFromUniqueProcessId;
            ULONG HandleCount;
            ULONG SessionId;
            ULONG_PTR PageDirectoryBase;
            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;
            SIZE_T PrivatePageCount;
            LARGE_INTEGER ReadOperationCount;
            LARGE_INTEGER WriteOperationCount;
            LARGE_INTEGER OtherOperationCount;
            LARGE_INTEGER ReadTransferCount;
            LARGE_INTEGER WriteTransferCount;
            LARGE_INTEGER OtherTransferCount;
            SYSTEM_THREAD_INFORMATION Threads[1];
        };

        enum class SYSTEM_INFORMATION_CLASS
        {
            SystemBasicInformation = 0,
            SystemPerformanceInformation = 2,
            SystemTimeOfDayInformation = 3,
            SystemProcessInformation = 5,
            SystemProcessorPerformanceInformation = 8,
            SystemInterruptInformation = 23,
            SystemExceptionInformation = 33,
            SystemRegistryQuotaInformation = 37,
            SystemLookasideInformation = 45,
            SystemCodeIntegrityInformation = 103,
            SystemPolicyInformation = 134,
        };

        extern "C" NTSYSAPI NTSTATUS NTAPI ZwQuerySystemInformation(
            IN SYSTEM_INFORMATION_CLASS infoClass,
            OUT PVOID buf,
            IN ULONG len,
            OUT OPTIONAL PULONG returned
        );
    }


    template <typename Provider>
    concept ZwMemoryInfoProvider = requires (const Provider& provider)
    {
        {
            provider.ZwQueryVirtualMemory(
                HANDLE{} /*hProcess*/,
                PVOID{} /*baseAddress*/,
                Api::MEMORY_INFORMATION_CLASS{} /*infoClass*/,
                PVOID{} /*buf*/,
                SIZE_T{} /*size*/,
                PSIZE_T{} /*returned*/
            )
        };
    };

    template <typename Provider>
    concept ZwSystemInfoProvider = requires (const Provider& provider)
    {
        {
            provider.ZwQuerySystemInformation(
                Api::SYSTEM_INFORMATION_CLASS{} /*infoClass*/,
                PVOID{} /*buf*/,
                ULONG{} /*len*/,
                PULONG{} /*returned*/
            )
        };
    };

    struct DefaultProvider
    {
        static NTSTATUS ZwQueryVirtualMemory(HANDLE hProcess, PVOID baseAddress, Api::MEMORY_INFORMATION_CLASS info, PVOID buf, SIZE_T size, SIZE_T* returned)
        {
            return Api::ZwQueryVirtualMemory(hProcess, baseAddress, info, buf, size, returned);
        }

        static NTSTATUS ZwQuerySystemInformation(Api::SYSTEM_INFORMATION_CLASS infoClass, PVOID buf, ULONG len, PULONG returned)
        {
            return Api::ZwQuerySystemInformation(infoClass, buf, len, returned);
        }
    };

    template <ZwMemoryInfoProvider Provider>
    class AddressSpace
    {
    public:
        class Region
        {
        private:
            const AddressSpace* m_owner;
            const void* m_addr;
            MEMORY_BASIC_INFORMATION m_info;
            NTSTATUS m_status;

        private:
            static const void* invalidAddress()
            {
                return reinterpret_cast<const void*>(-1);
            }

        public:
            static Region makeInvalid(const AddressSpace& owner)
            {
                return Region(owner, invalidAddress());
            }

            explicit Region(const AddressSpace& owner, const void* const addr)
                : m_owner(&owner)
                , m_addr(addr)
                , m_info{}
                , m_status{ -1 }
            {
                if (addr != invalidAddress())
                {
                    update();
                }
            }

            bool update()
            {
                m_info = {};
                SIZE_T returned{};

                m_status = m_owner->provider().ZwQueryVirtualMemory(
                    m_owner->processHandle(),
                    const_cast<void*>(m_addr),
                    Api::MEMORY_INFORMATION_CLASS::MemoryBasicInformation,
                    &m_info,
                    sizeof(m_info),
                    &returned
                );

                if (!valid())
                {
                    m_addr = invalidAddress();
                    return false;
                }

                return true;
            }

            bool valid() const
            {
                return NT_SUCCESS(m_status);
            }

            NTSTATUS status() const
            {
                return m_status;
            }

            Region& operator ++ ()
            {
                if (!valid())
                {
                    return *this;
                }

                m_addr = static_cast<const unsigned char*>(m_info.BaseAddress) + m_info.RegionSize;
                update();
                return *this;
            }

            Region operator ++ (int)
            {
                Region prev = *this;
                ++(*this);
                return prev;
            }

            bool operator == (const Region& region) const
            {
                return (m_owner == region.m_owner) && (m_addr == region.m_addr);
            }

            const MEMORY_BASIC_INFORMATION* operator -> () const
            {
                return &m_info;
            }

            const MEMORY_BASIC_INFORMATION& operator * () const
            {
                return m_info;
            }

            explicit operator bool () const
            {
                return valid();
            }
        };

    private:
        const Provider& m_provider;
        const HANDLE m_hProcess;

    public:
        explicit AddressSpace(const Provider& provider, const HANDLE hProcess)
            : m_provider(provider)
            , m_hProcess(hProcess)
        {
        }

        Region query(const void* const addr) const
        {
            return Region(*this, addr);
        }

        Region begin() const
        {
            return Region(*this, 0);
        }

        Region end() const
        {
            return Region::makeInvalid(*this);
        }

        const Provider& provider() const
        {
            return m_provider;
        }

        HANDLE processHandle() const
        {
            return m_hProcess;
        }
    };

    
    enum class StorageType
    {
        fixed,
        variable
    };

    template <typename DataType, AllocatorTrait...>
    struct Storage;

    template <typename DataType>
    class Storage<DataType>
    {
    public:
        using Type = DataType;
        static constexpr auto k_type = StorageType::fixed;

    private:
        Type m_data{};

    public:
        Type* data()
        {
            return &m_data;
        }

        const Type* data() const
        {
            return &m_data;
        }

        size_t size() const
        {
            return sizeof(m_data);
        }
    };

    template <typename DataType, AllocatorTrait Allocator>
    class Storage<DataType, Allocator> : public Allocator
    {
    public:
        using Type = DataType;
        static constexpr auto k_type = StorageType::variable;

    public:
        Type* data()
        {
            if (Allocator::empty())
            {
                return nullptr;
            }

            void* const raw = Allocator::data();
            return static_cast<Type*>(raw);
        }

        const Type* data() const
        {
            if (Allocator::empty())
            {
                return nullptr;
            }

            const void* const raw = Allocator::data();
            return static_cast<const Type*>(raw);
        }
    };

    template <typename Storage>
    concept StorageTrait = requires (const Storage& storage)
    {
        { storage.size() };
        { storage.data() };
        { Storage::k_type };
        typename Storage::Type;
    };

    class SystemInfo
    {
    public:
        template <Api::SYSTEM_INFORMATION_CLASS infoClass, typename InfoType, AllocatorTrait... OptionalAllocator>
        struct Info : public Storage<InfoType, OptionalAllocator...>
        {
            static constexpr auto k_infoClass = infoClass;
            NTSTATUS status{};
            ULONG returned{};
        };

        template <AllocatorTrait Allocator>
        struct Processes : public Info<Api::SYSTEM_INFORMATION_CLASS::SystemProcessInformation, Api::SYSTEM_PROCESS_INFORMATION, Allocator>
        {
            using Super = Info<Api::SYSTEM_INFORMATION_CLASS::SystemProcessInformation, Api::SYSTEM_PROCESS_INFORMATION, Allocator>;

            class Process
            {
            public:
                struct ProcessEntry : public Api::SYSTEM_PROCESS_INFORMATION
                {
                    class Thread
                    {
                    public:
                        static constexpr auto k_invalidIndex = 0xFFFFFFFFu;

                    private:
                        const ProcessEntry* m_process;
                        unsigned int m_index;

                    public:
                        explicit Thread(const ProcessEntry* process, const unsigned int index) : m_process(process), m_index(index)
                        {
                        }

                        bool operator == (const Thread& thread) const
                        {
                            return (m_index == thread.m_index) && (m_process == thread.m_process);
                        }

                        Thread& operator ++ ()
                        {
                            if (m_index == k_invalidIndex)
                            {
                                return *this;
                            }

                            ++m_index;
                            if (m_index >= m_process->NumberOfThreads)
                            {
                                m_index = k_invalidIndex;
                            }

                            return *this;
                        }

                        Thread operator ++ (int)
                        {
                            const auto prev = *this;
                            ++(*this);
                            return prev;
                        }

                        const Api::SYSTEM_THREAD_INFORMATION* operator -> () const
                        {
                            return &(m_process->Threads[m_index]);
                        }

                        const Api::SYSTEM_THREAD_INFORMATION& operator * () const
                        {
                            return m_process->Threads[m_index];
                        }
                    };

                    Thread begin() const
                    {
                        if (!this->NumberOfThreads)
                        {
                            return end();
                        }

                        return Thread(this, 0);
                    }

                    Thread end() const
                    {
                        return Thread(this, Thread::k_invalidIndex);
                    }
                };

            private:
                const ProcessEntry* m_entry;

            public:
                explicit Process(const Api::SYSTEM_PROCESS_INFORMATION* const entry) : m_entry(static_cast<const ProcessEntry*>(entry))
                {
                }

                bool operator == (const Process& process) const
                {
                    return m_entry == process.m_entry;
                }

                Process& operator ++ ()
                {
                    if (!m_entry)
                    {
                        return *this;
                    }

                    if (m_entry->NextEntryOffset == 0)
                    {
                        m_entry = nullptr;
                        return *this;
                    }

                    m_entry = reinterpret_cast<const ProcessEntry*>(reinterpret_cast<const unsigned char*>(m_entry) + m_entry->NextEntryOffset);

                    return *this;
                }

                Process operator ++ (int)
                {
                    const auto prev = *this;
                    ++(*this);
                    return prev;
                }

                const ProcessEntry* operator -> () const
                {
                    return m_entry;
                }

                const ProcessEntry& operator * () const
                {
                    return *m_entry;
                }
            };

            Process begin() const
            {
                if (!Super::size())
                {
                    return end();
                }

                return ++Process(Super::data());
            }

            Process end() const
            {
                return Process(nullptr);
            }
        };

        template <typename Info, ZwSystemInfoProvider Provider = DefaultProvider>
        static Info query(const Provider& provider)
        {
            Info info{};
            if constexpr (Info::k_type == StorageType::fixed)
            {
                info.status = provider.ZwQuerySystemInformation(Info::k_infoClass, info.data(), static_cast<ULONG>(info.size()), &info.returned);
            }
            else if constexpr (Info::k_type == StorageType::variable)
            {
                info.status = provider.ZwQuerySystemInformation(Info::k_infoClass, nullptr, 0, &info.returned);
                while (info.status == STATUS_INFO_LENGTH_MISMATCH)
                {
                    info.resize(info.returned);
                    if (info.size() != info.returned)
                    {
                        info.status = STATUS_UNSUCCESSFUL;
                        break;
                    }
                    info.status = provider.ZwQuerySystemInformation(Info::k_infoClass, info.data(), static_cast<ULONG>(info.size()), &info.returned);
                }

                if (!NT_SUCCESS(info.status))
                {
                    info.clear();
                    info.returned = 0;
                    return info;
                }
            }
            else
            {
                static_assert("Unknown storage type");
            }

            return info;
        }
    };
}