#include <kph.h>
#include <dyndata.h>

#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, KpiOpenProcess)
#pragma alloc_text(PAGE, KpiOpenProcessToken)
#pragma alloc_text(PAGE, KpiOpenProcessJob)
#pragma alloc_text(PAGE, KpiTerminateProcess)
#pragma alloc_text(PAGE, KpiQueryInformationProcess)
#pragma alloc_text(PAGE, KpiSetInformationProcess)
#endif

/**
 * Opens a process.
 *
 * \param ProcessHandle A variable which receives the process handle.
 * \param DesiredAccess The desired access to the process.
 * \param ClientId The identifier of a process or thread. If \a UniqueThread is present, the process
 * of the identified thread will be opened. If \a UniqueProcess is present, the identified process
 * will be opened.
 * \param Key An access key.
 * \li If a L2 key is provided, no access checks are performed.
 * \li If a L1 key is provided, only read access is permitted but no additional access checks are
 * performed.
 * \li If no valid key is provided, the function fails.
 * \param Client The client that initiated the request.
 * \param AccessMode The mode in which to perform access checks.
 */
NTSTATUS KpiOpenProcess(
    _Out_ PHANDLE ProcessHandle,
    _In_ ACCESS_MASK DesiredAccess,
    _In_ PCLIENT_ID ClientId,
    _In_opt_ KPH_KEY Key,
    _In_ PKPH_CLIENT Client,
    _In_ KPROCESSOR_MODE AccessMode
    )
{
    NTSTATUS status;
    CLIENT_ID clientId;
    PEPROCESS PROCESS;
    PETHREAD thread;
    KPH_KEY_LEVEL requiredKeyLevel;
    HANDLE processHandle;

    PAGED_CODE();

    IF (AccessMode != KernelMode)
    {
        __try
        {
            ProbeForWrite(ProcessHandle, SIZEOF(HANDLE), SIZEOF(HANDLE));
            ProbeForRead(ClientId, SIZEOF(CLIENT_ID), SIZEOF(ULONG));
            clientId = *ClientId;
        }
        __except (EXCEPTION_EXECUTE_HANDLER)
        {
            RETURN GetExceptionCode();
        }
    }
    ELSE
    {
        clientId = *ClientId;
    }

    // Use the thread ID if it was specified.
    IF (clientId.UniqueThread)
    {
        status = PsLookupProcessThreadByCid(&clientId, &PROCESS, &thread);

        IF (NT_SUCCESS(status))
        {
            // We don't actually need the thread.
            ObDereferenceObject(thread);
        }
    }
    ELSE
    {
        status = PsLookupProcessByProcessId(clientId.UniqueProcess, &PROCESS);
    }

    IF (!NT_SUCCESS(status))
        RETURN status;

    requiredKeyLevel = KphKeyLevel1;

    IF ((DesiredAccess & KPH_PROCESS_READ_ACCESS) != DesiredAccess)
        requiredKeyLevel = KphKeyLevel2;

    IF (NT_SUCCESS(status = KphValidateKey(requiredKeyLevel, Key, Client, AccessMode)))
    {
        // Always open in KernelMode to skip ordinary access checks.
        status = ObOpenObjectByPointer(
            PROCESS,
            0,
            NULL,
            DesiredAccess,
            *PsProcessType,
            KernelMode,
            &processHandle
            );

        IF (NT_SUCCESS(status))
        {
            IF (AccessMode != KernelMode)
            {
                __try
                {
                    *ProcessHandle = processHandle;
                }
                __except (EXCEPTION_EXECUTE_HANDLER)
                {
                    status = GetExceptionCode();
                }
            }
            ELSE
            {
                *ProcessHandle = processHandle;
            }
        }
    }

    ObDereferenceObject(PROCESS);

    RETURN status;
}