extern "C"
{

#include "hooked_proc.h"
#include "debug.h"
#include "func.h"
#include "undocnt.h"

NTSTATUS NewZwOpenProcess(PHANDLE ProcessHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes,PCLIENT_ID ClientId);
NTSTATUS NewZwOpenThread(PHANDLE ThreadHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes,PCLIENT_ID ClientId);
} //extern "C"


/*
 pointers to original functions
*/

ZW_OPEN_PROCESS OldZwOpenProcess=NULL;
ZW_OPEN_THREAD OldZwOpenThread=NULL;

/*
 our implementation of ZwOpenProcess
 at first check whether pid is protected
 if so return STATUS_ACCESS_DENIED
 otherwise call original function
*/

NTSTATUS NewZwOpenProcess(PHANDLE ProcessHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes,PCLIENT_ID ClientId)
{
  DbgMsg("hooked_proc.cpp: NewZwOpenProcess(ProcessHandle:0x%.8X,DesiredAccess:0x%.8X,ObjectAttributes:0x%.8X,ClientId:0x%.8X)",
         ProcessHandle,DesiredAccess,ObjectAttributes,ClientId);

  int cid_valid=func_is_good_read_ptr(ClientId,sizeof(CLIENT_ID));
  if (cid_valid)
  {
    DbgMsg("hooked_proc.cpp: NewZwOpenProcess: ClientId->UniqueProcess=0x%.8X",ClientId->UniqueProcess);
    DbgMsg("hooked_proc.cpp: NewZwOpenProcess: ClientId->UniqueThread=0x%.8X",ClientId->UniqueThread);
  }

  NTSTATUS status;

  int protect=cid_valid?func_check_process_protection((ULONG)ClientId->UniqueProcess):FALSE;

  if (!protect) status=OldZwOpenProcess(ProcessHandle,DesiredAccess,ObjectAttributes,ClientId);
  else status=STATUS_ACCESS_DENIED;

  DbgMsg("hooked_proc.cpp: NewZwOpenProcess(-):0x%.8X",status);
  return status;
}


/*
 our implementation of ZwOpenThread
 we call original function to get thread handle to be able to call 
 ZwQueryInformationThread on it to retrieve process id that owns the thread
 we can check then whether the pid of thread owning process is protected
 and if so then we close the handle and return STATUS_ACCESS_DENIED

 we use original function call because we want parameters to be checked
 that's why we need to use ThreadHandle for the original call because it is 
 located in usermode, if we call kernel function without checking 
 the parameters here we could have BSOD everytime user process calls this 
 function with bad parameters unless we make whole call twice - once for us 
 to retrieve handle for ZwQueryInformationThread to check the pid and then 
 once for original call
*/

NTSTATUS NewZwOpenThread(PHANDLE ThreadHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes,PCLIENT_ID ClientId)
{
  DbgMsg("hooked_proc.cpp: NewZwOpenThread(ThreadHandle:0x%.8X,DesiredAccess:0x%.8X,ObjectAttributes:0x%.8X,ClientId:0x%.8X)",
         ThreadHandle,DesiredAccess,ObjectAttributes,ClientId);

  int cid_valid=func_is_good_read_ptr(ClientId,sizeof(CLIENT_ID));
  int thr_valid=func_is_good_read_ptr(ThreadHandle,sizeof(HANDLE));
  if (cid_valid)
  {
    DbgMsg("hooked_proc.cpp: NewZwOpenThread: ClientId->UniqueProcess=0x%.8X",ClientId->UniqueProcess);
    DbgMsg("hooked_proc.cpp: NewZwOpenThread: ClientId->UniqueThread=0x%.8X",ClientId->UniqueThread);
  }

  //save original ThreadHandle value, we may restore it later
  HANDLE org_handle;
  if (thr_valid) org_handle=*ThreadHandle;

  NTSTATUS status=OldZwOpenThread(ThreadHandle,DesiredAccess,ObjectAttributes,ClientId);

  if NT_SUCCESS(status)
  {
    THREAD_BASIC_INFORMATION info;
    NTSTATUS status_priv=_ZwQueryInformationThread(*ThreadHandle,ThreadBasicInformation,&info,sizeof(info),NULL);
    if (NT_SUCCESS(status_priv))
    {
      //secondly check for process protection that owns this thread
      int protect=func_check_process_protection((ULONG)info.ClientId.UniqueProcess);
      if (protect)
      {
        status=STATUS_ACCESS_DENIED;

        ZwClose(*ThreadHandle);
        //restore original ThreadHandle value
        *ThreadHandle=org_handle; 
      }
    }
  }

  DbgMsg("hooked_proc.cpp: NewZwOpenThread(-):0x%.8X",status);
  return status;
}