extern "C"
{

#include "hooked_proc.h"
#include "debug.h"
#include "func.h"
#include "undocnt.h"
#include "compat.h"
#include <ntstrsafe.h>
#include <minmax.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);
NTSTATUS NewZwLoadDriver(PUNICODE_STRING DriverServiceName);
NTSTATUS NewZwOpenSection(PHANDLE SectionHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes);
} //extern "C"


/*
 pointers to original functions
*/

ZW_OPEN_PROCESS OldZwOpenProcess=NULL;
ZW_OPEN_THREAD OldZwOpenThread=NULL;
ZW_TERMINATE_PROCESS OldZwTerminateProcess=NULL;
ZW_LOAD_DRIVER OldZwLoadDriver=NULL;
ZW_OPEN_SECTION OldZwOpenSection=NULL;


/*
 our implementation of ZwOpenProcess
 at first check whether pid is protected
 if so ask user what to do
 then return STATUS_ACCESS_DENIED or call original function using the user answer
*/

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;
  ULONG pid=(ULONG)PsGetCurrentProcessId();

  if (protect)
    protect=func_notify_ask_user(global_dev_ext,pid,RULE_TYPE_PROCESS_PROTECTION,(ULONG)ClientId->UniqueProcess,DesiredAccess);

  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 ask user what to do and do nothing or close the handle 
 and return STATUS_ACCESS_DENIED if user denied the access

 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);

      ULONG pid=(ULONG)PsGetCurrentProcessId();

      if (protect)
       protect=func_notify_ask_user(global_dev_ext,pid,RULE_TYPE_PROCESS_PROTECTION,(ULONG)info.ClientId.UniqueProcess,DesiredAccess);

      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;
}


/*
 our implementation of ZwLoadDriver
 everytime we ask user what to do
 then return STATUS_ACCESS_DENIED or call original function using the user answer
*/

NTSTATUS NewZwLoadDriver(PUNICODE_STRING DriverServiceName)
{
  DbgMsg("hooked_proc.cpp: NewZwLoadDriver(DriverServiceName:0x%.8X)",DriverServiceName);

  int protect=0;
  NTSTATUS status;
  ULONG pid=(ULONG)PsGetCurrentProcessId();

  int buf_valid=func_is_good_read_ptr(DriverServiceName,sizeof(UNICODE_STRING))
             && func_is_good_read_ptr(DriverServiceName->Buffer,DriverServiceName->Length);

  if (buf_valid)
  {
    wchar_t name[512];
    int len=min(sizeof(name)-2,DriverServiceName->Length);
    RtlCopyMemory(name,DriverServiceName->Buffer,len);
    name[len/2]=0x0000;
    DbgMsg("hooked_proc.cpp: NewZwLoadDriver: DriverServiceName->Buffer=%S",name);

    protect=func_notify_ask_user(global_dev_ext,pid,RULE_TYPE_LOAD_DRIVER_PROTECTION,(ULONG)DriverServiceName,0);
  }

  if (!protect) status=OldZwLoadDriver(DriverServiceName);
  else status=STATUS_ACCESS_DENIED;

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



/*
 our implementation of ZwOpenSection
 we dereference handle of successful call and compare object with our protected \Device\PhysicalMemory 
 section object address
 if the object is protected we ask user and return STATUS_ACCESS_DENIED or pass original handle
*/

NTSTATUS NewZwOpenSection(PHANDLE SectionHandle,ACCESS_MASK DesiredAccess,POBJECT_ATTRIBUTES ObjectAttributes)
{
  DbgMsg("hooked_proc.cpp: NewZwOpenSection(SectionHandle:0x%.8X;DesiredAccess:0x%.8X;ObjectAttributes:0x%.8X)",
          SectionHandle,DesiredAccess,ObjectAttributes);

  int sec_valid=func_is_good_read_ptr(SectionHandle,sizeof(HANDLE));

#ifdef DEBUG
  int oa_valid=func_is_good_read_ptr(ObjectAttributes,sizeof(OBJECT_ATTRIBUTES))
            && func_is_good_read_ptr(ObjectAttributes->ObjectName,sizeof(UNICODE_STRING))
            && func_is_good_read_ptr(ObjectAttributes->ObjectName->Buffer,ObjectAttributes->ObjectName->Length);

  if (oa_valid)
  {
    wchar_t name[512];
    int len=min(sizeof(name)-2,ObjectAttributes->ObjectName->Length);
    RtlCopyMemory(name,ObjectAttributes->ObjectName->Buffer,len);
    name[len/2]=0x0000;
    DbgMsg("hooked_proc.cpp: NewZwOpenSection: ObjectAttributes->ObjectName->Buffer=%S",name);
  }
#endif
  
  //save original value of section handle
  HANDLE org_handle;
  if (sec_valid) org_handle=*SectionHandle;

  NTSTATUS status=OldZwOpenSection(SectionHandle,DesiredAccess,ObjectAttributes);
  if (NT_SUCCESS(status))
  {
    PVOID obj;
    NTSTATUS status2=ObReferenceObjectByHandle(*SectionHandle,0,0,KernelMode,&obj,NULL);

    if (NT_SUCCESS(status2))
    {
      DbgMsg("hooked_proc.cpp: NewZwOpenSection: obj=0x%.8X",obj);

      //compare with protected \Device\PhysicalMemory section object address
      if (obj==func_protected_device_physical_memory)
      {
        ULONG pid=(ULONG)PsGetCurrentProcessId();
        int protect=func_notify_ask_user(global_dev_ext,pid,RULE_TYPE_PHYSICAL_MEMORY_PROTECTION,0,DesiredAccess);

        if (protect)
        {
          //if access was denied we return STATUS_ACCESS_DENIED and refresh original handle value
          status=STATUS_ACCESS_DENIED;
          ZwClose(*SectionHandle);
          *SectionHandle=org_handle;
        }
      }
      ObDereferenceObject(obj);
    }
  }

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