extern "C"
{

#include "hooked_proc.h"
#include "debug.h"
#include "func.h"
#include "undocnt.h"
#include "compat.h"
#include <ntstrsafe.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);
NTSTATUS NewZwTerminateProcess(HANDLE ProcessHandle,NTSTATUS ExitStatus);
} //extern "C"


/*
 pointers to original functions
*/

ZW_OPEN_PROCESS OldZwOpenProcess=NULL;
ZW_OPEN_THREAD OldZwOpenThread=NULL;
ZW_TERMINATE_PROCESS OldZwTerminateProcess=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;
}


/*
 our implementation of ZwTerminateProcess
 it deletes rule (if exists) for the process that terminates
 at first we have to check whether it is not current process what is to terminate
 if so we have to delete the rule before executing original ZwTerminateProcess
*/

NTSTATUS NewZwTerminateProcess(HANDLE ProcessHandle,NTSTATUS ExitStatus)
{
  DbgMsg("hooked_proc.cpp: NewZwTerminateProcess(ProcessHandle:0x%.8X,ExitStatus:0x%.8X)",ProcessHandle,ExitStatus);

  //handle current process termination
  if (!ProcessHandle || (ProcessHandle==(HANDLE)0xFFFFFFFF))
  {
    ULONG pid=(ULONG)PsGetCurrentProcessId();

    DbgMsg("hooked_proc.cpp: NewZwTerminateProcess current proc pid:%d",pid);

    func_protect_process(global_dev_ext,pid,FALSE);

    DbgMsg("hooked_proc.cpp: NewZwTerminateProcess(-)");
  }

  NTSTATUS status=OldZwTerminateProcess(ProcessHandle,ExitStatus);

  if (NT_SUCCESS(status) && ProcessHandle)
  {
    //delete rule for this process if terminated successfully
    ULONG pid=compat_get_pid_from_process_handle(ProcessHandle);

    DbgMsg("hooked_proc.cpp: NewZwTerminateProcess: pid:%d)",pid);

    if (pid!=0xFFFFFFFF)
      func_protect_process(global_dev_ext,pid,FALSE);
  }

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