// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

#include "precomp.h"

#include "NtPrivApi.hpp"

[[nodiscard]] NTSTATUS NtPrivApi::s_GetProcessParentId(_Inout_ PULONG ProcessId)
{
    // TODO: Get Parent current not really available without winternl + NtQueryInformationProcess. http://osgvsowi/8394495
    OBJECT_ATTRIBUTES oa;
    InitializeObjectAttributes(&oa, nullptr, 0, nullptr, nullptr);

    CLIENT_ID ClientId;
    ClientId.UniqueProcess = UlongToHandle(*ProcessId);
    ClientId.UniqueThread = 0;

    HANDLE ProcessHandle;
    NTSTATUS Status = s_NtOpenProcess(&ProcessHandle, PROCESS_QUERY_LIMITED_INFORMATION, &oa, &ClientId);

    PROCESS_BASIC_INFORMATION BasicInfo = { 0 };
    if (NT_SUCCESS(Status))
    {
        Status = s_NtQueryInformationProcess(ProcessHandle, ProcessBasicInformation, &BasicInfo, sizeof(BasicInfo), nullptr);
        LOG_IF_FAILED(s_NtClose(ProcessHandle));
    }

    if (!NT_SUCCESS(Status))
    {
        *ProcessId = 0;
        return Status;
    }

    // This is the actual field name, but in the public SDK, it's named Reserved3. We need to pursue publishing the real name.
    //*ProcessId = (ULONG)BasicInfo.InheritedFromUniqueProcessId;
    *ProcessId = (ULONG)BasicInfo.Reserved3;
    return STATUS_SUCCESS;
}

[[nodiscard]] NTSTATUS NtPrivApi::s_NtOpenProcess(_Out_ PHANDLE ProcessHandle,
                                                  _In_ ACCESS_MASK DesiredAccess,
                                                  _In_ POBJECT_ATTRIBUTES ObjectAttributes,
                                                  _In_opt_ PCLIENT_ID ClientId)
{
    HMODULE hNtDll = _Instance()._hNtDll;

    if (hNtDll != nullptr)
    {
        typedef NTSTATUS (*PfnNtOpenProcess)(HANDLE ProcessHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES ObjectAttributes, PCLIENT_ID ClientId);

        static PfnNtOpenProcess pfn = (PfnNtOpenProcess)GetProcAddress(hNtDll, "NtOpenProcess");

        if (pfn != nullptr)
        {
            return pfn(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);
        }
    }

    return STATUS_UNSUCCESSFUL;
}

[[nodiscard]] NTSTATUS NtPrivApi::s_NtQueryInformationProcess(_In_ HANDLE ProcessHandle,
                                                              _In_ PROCESSINFOCLASS ProcessInformationClass,
                                                              _Out_ PVOID ProcessInformation,
                                                              _In_ ULONG ProcessInformationLength,
                                                              _Out_opt_ PULONG ReturnLength)
{
    HMODULE hNtDll = _Instance()._hNtDll;

    if (hNtDll != nullptr)
    {
        typedef NTSTATUS (*PfnNtQueryInformationProcess)(HANDLE ProcessHandle, PROCESSINFOCLASS ProcessInformationClass, PVOID ProcessInformation, ULONG ProcessInformationLength, PULONG ReturnLength);

        static PfnNtQueryInformationProcess pfn = (PfnNtQueryInformationProcess)GetProcAddress(hNtDll, "NtQueryInformationProcess");

        if (pfn != nullptr)
        {
            return pfn(ProcessHandle, ProcessInformationClass, ProcessInformation, ProcessInformationLength, ReturnLength);
        }
    }

    return STATUS_UNSUCCESSFUL;
}

[[nodiscard]] NTSTATUS NtPrivApi::s_NtClose(_In_ HANDLE Handle)
{
    HMODULE hNtDll = _Instance()._hNtDll;

    if (hNtDll != nullptr)
    {
        typedef NTSTATUS (*PfnNtClose)(HANDLE Handle);

        static PfnNtClose pfn = (PfnNtClose)GetProcAddress(hNtDll, "NtClose");

        if (pfn != nullptr)
        {
            return pfn(Handle);
        }
    }

    return STATUS_UNSUCCESSFUL;
}

NtPrivApi::NtPrivApi()
{
    // NOTE: Use LoadLibraryExW with LOAD_LIBRARY_SEARCH_SYSTEM32 flag below to avoid unneeded directory traversal.
    //       This has triggered CPG boot IO warnings in the past.
    _hNtDll = LoadLibraryExW(L"ntdll.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32);
}

NtPrivApi::~NtPrivApi()
{
    if (_hNtDll != nullptr)
    {
        FreeLibrary(_hNtDll);
        _hNtDll = nullptr;
    }
}

NtPrivApi& NtPrivApi::_Instance()
{
    static NtPrivApi ntapi;
    return ntapi;
}