/*
 hook9 is driver that protects processes using hooks for ZwOpenProcess 
 and ZwOpenThread
 when a process protection is requested all its threads are also protected

 if there is an active rule for some process and this process terminates 
 with ZwTerminateProcess, the rule is deleted

 generic PID from EPROCESS was implemented in compat.cpp

 for protected processes the access is not automatically denied
 instead of returning access denied driver sends notification to user mode application
 and it asks user whether or not to permit the access
 implementation of the notification is using asynchronous IO
 at first user mode application sends IOCTL_NOTIFY to our driver 
 driver stores the IRP and return STATUS_PENDING, it also sets cancel routine for the case 
 that user mode application closes the handle
 driver completes stored IRP when another app calls one of hooked services to access 
 protected process, after then it waits for a signal with user answer 
 when IRP completes user mode application receives notification as a asynchronous IO completion
 it asks user for the answer and sends it to the driver, driver signals the event 
 for the second application and it completes its stuff using the answer from user
 user mode application then sends new IOCTL_NOTIFY to be prepared for the next notification

 hook for ZwLoadDriver was added to protect loading kernel mode drivers 
 this is just example, there exist other ways how to load driver
 another problem is that if app uses Service Control Manager to install driver
 you won't be asked for permission for that app but for services.exe to load driver
 you can see this when you try to load older hookX when hook for ZwLoadDriver is installed
 you can watch the other results when you run e.g. Registry Monitor from Sysinternals
 if you do not permit loading driver you'll receive error message
 if you permit regmon starts without problems
 but once you permit loading its driver regmon won't unload it and that's why you'll not be asked 
 if you run regmon later again 

 another protection mechanism is implemented via ZwOpenSection hook to control access of user mode 
 application to \Device\PhysicalMemory section object, our implementation is resistant against 
 symbolic link cheating because we work with the object address rather than with its name 
*/

extern "C"
{

#include <wdm.h>
#include "debug.h"
#include "hooking.h"
#include "drvcomm.h"
#include "func.h"
#include "undocnt.h"
#include "compat.h"

NTSTATUS hook9_create(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook9_close(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook9_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID hook9_unload(PDRIVER_OBJECT DriverObject);
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject,PUNICODE_STRING RegistryPath);
} //extern "C"


/*
 dos device name is global variable because we use it in unload too,
 we can always make another RtlInitUnicodeString if we don't like global vars
*/

UNICODE_STRING dos_dev_name;


/*
 this is very simple implementation of cancel routine
 in our driver we know that there can be only one IRP that is pending at the moment
 in such case we would have to implement queue for such IRPs and manage this queue using spinlock
*/

VOID hook9_irp_cancel_routine(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook9.cpp: hook9_irp_cancel_routine(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);
  PDEVICE_EXTENSION dev_ext=(PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
  KIRQL oirql;

  //release the global cancel spin lock
  IoReleaseCancelSpinLock(Irp->CancelIrql);

  //clear notify_irp in device extension
  InterlockedExchange((LONG *)&dev_ext->notify_irp,NULL);

  //complete IRP
  Irp->IoStatus.Status=STATUS_CANCELLED;
  Irp->IoStatus.Information=0;

  IoCompleteRequest(Irp,IO_NO_INCREMENT);
  DbgMsg("hook9.cpp: hook9_irp_cancel_routine(-)");
  return;
}


/*
 create function is called everytime CreateFile is called on our device 
*/

NTSTATUS hook9_create(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook9.cpp: hook9_create(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

  NTSTATUS status=STATUS_SUCCESS;
  Irp->IoStatus.Status=status;
  IoCompleteRequest(Irp,IO_NO_INCREMENT);

  DbgMsg("hook9.cpp: hook9_create(-):0x%.8X)",status);

  return status;
}


/*
 close function is called everytime CloseHandle is called on our device
 close is associated with IRP_MJ_CLOSE and it is NOT executed in the context 
 of the CloseHandle caller, if we want to make some cleanup in that context 
 we rather associate cleanup function with IRP_MJ_CLEANUP
*/

NTSTATUS hook9_close(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook9.cpp: hook9_close(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

  NTSTATUS status=STATUS_SUCCESS;
  Irp->IoStatus.Status=status;
  IoCompleteRequest(Irp,IO_NO_INCREMENT);

  DbgMsg("hook9.cpp: hook9_close(-):0x%.8X)",status);

  return status;
}


/*
 device control function is called everytime DeviceIoControl is called on our
 device, it is common way how user mode app communicate with driver
*/

NTSTATUS hook9_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook9.cpp: hook9_device_control(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

  PIO_STACK_LOCATION stack=IoGetCurrentIrpStackLocation(Irp);

  NTSTATUS status=STATUS_SUCCESS;

  UCHAR *buf_in,*buf_out;
  ULONG buf_in_len,buf_out_len,code,ret;

  code=stack->Parameters.DeviceIoControl.IoControlCode;

  /*
   for Buffered IO both input and output buffer are the same 
   Irp->AssociatedIrp.SystemBuffer
  */
  buf_in=buf_out=(UCHAR *)Irp->AssociatedIrp.SystemBuffer;
  buf_in_len=stack->Parameters.DeviceIoControl.InputBufferLength;
  buf_out_len=stack->Parameters.DeviceIoControl.OutputBufferLength;

  DbgMsg("hook9.cpp: hook9_device_control: code:0x%.8X,buf_in:0x%.8X,buf_in_len:0x%.8X,buf_out:0x%.8X,buf_out_len:0x%.8X",
         code,buf_in,buf_in_len,buf_out,buf_out_len);


  PDRVCOMM_REQUEST_BUFFER buf_req=(PDRVCOMM_REQUEST_BUFFER)buf_in;
  PDRVCOMM_RESPONSE_BUFFER buf_res=(PDRVCOMM_RESPONSE_BUFFER)buf_out;
  PDEVICE_EXTENSION dev_ext=(PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
 
  PIRP notify_irp;

  Irp->IoStatus.Information=sizeof(DRVCOMM_RESPONSE_BUFFER);

  /*
   for every command we implement the functionality 
  */
  switch (code)
  {
    case IOCTL_HOOK_START:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_HOOK_START");
      SDT_index_ZwQueryInformationThread=buf_req->parameters.init.SDT_index_ZwQueryInformationThread;
      compat_init_offsets();
      func_init();
      ret=hooking_hook();
      buf_res->status=ret;
      break;


    case IOCTL_HOOK_STOP:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_HOOK_STOP");

      notify_irp=dev_ext->notify_irp;
      InterlockedExchange((LONG *)&dev_ext->notify_irp,NULL);
      if (notify_irp) IoCancelIrp(notify_irp);

      ret=hooking_unhook();
      buf_res->status=ret;
      break;


    case IOCTL_PROTECT_PROCESS:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_PROTECT_PROCESS");
      ret=func_protect_process(dev_ext,buf_req->parameters.protect_process.pid,
                               buf_req->parameters.protect_process.enable);
      buf_res->status=ret;
      break;


    case IOCTL_NOTIFY:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_NOTIFY");
      dev_ext->notify_irp=Irp;

      IoSetCancelRoutine(Irp,hook9_irp_cancel_routine);
      if (Irp->Cancel && IoSetCancelRoutine(Irp,NULL))
      {
        status=STATUS_CANCELLED;
        break;
      }

      IoMarkIrpPending(Irp);

      KeSetEvent(&dev_ext->notify_irp_event,0,FALSE);

      status=STATUS_PENDING;
      break;


    case IOCTL_NOTIFY_PROCESS_CHECK:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_NOTIFY_PROCESS_CHECK");
      InterlockedExchange((LONG *)&dev_ext->notify_answer,buf_req->parameters.process_check.permit);
      KeSetEvent(&dev_ext->notify_event,0,FALSE);
      break;

    case IOCTL_NOTIFY_LOAD_DRIVER_CHECK:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_NOTIFY_LOAD_DRIVER_CHECK");
      InterlockedExchange((LONG *)&dev_ext->notify_answer,buf_req->parameters.load_driver_check.permit);
      KeSetEvent(&dev_ext->notify_event,0,FALSE);
      break;

    case IOCTL_NOTIFY_PHYSICAL_MEMORY_CHECK:
      DbgMsg("hook9.cpp: hook9_device_control IOCTL_NOTIFY_PHYSICAL_MEMORY_CHECK");
      InterlockedExchange((LONG *)&dev_ext->notify_answer,buf_req->parameters.physical_memory_check.permit);
      KeSetEvent(&dev_ext->notify_event,0,FALSE);
      break;


    /*
     unknown codes should also be handled
    */
    default:
      status=STATUS_INVALID_DEVICE_REQUEST;
      Irp->IoStatus.Information=0;
  }

  if (status!=STATUS_PENDING)
  {
    Irp->IoStatus.Status=status;
    IoCompleteRequest(Irp,IO_NO_INCREMENT);
  }

  DbgMsg("hook9.cpp: hook9_device_control(-):0x%.8X)",status);

  return status;
}


/*
 unload is called when driver is being unloaded, if we do not implement unload
 function them our driver can't be unloaded dynamically
*/

VOID hook9_unload(PDRIVER_OBJECT DriverObject)
{    
  DbgMsg("hook9.cpp: hook9_unload(DriverObject:0x%.8X)",DriverObject);

  hooking_unhook();

  //free memory we've allocated for rules
  func_free_list((PDEVICE_EXTENSION)DriverObject->DeviceObject->DeviceExtension);

  IoDeleteSymbolicLink(&dos_dev_name);
  IoDeleteDevice(DriverObject->DeviceObject);

  DbgMsg("hook9.cpp: hook9_unload(-)");
}


/*
 DriverEntry is common driver entry point
*/

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject,PUNICODE_STRING RegistryPath)
{
  DbgMsg("hook9.cpp: DriverEntry(DriverObject:0x%.8X,RegistryPath:0x%.8X)",DriverObject,RegistryPath);

  UNICODE_STRING dev_name;
  RtlInitUnicodeString(&dev_name,DEVICE_NAME);
  RtlInitUnicodeString(&dos_dev_name,DOS_DEVICE_NAME); 

  /*
   if we want our driver to be accessible we need to create device for it, 
   one driver can have more devices 
  */
  PDEVICE_OBJECT dev_obj;
  NTSTATUS status=IoCreateDevice(DriverObject,sizeof(DEVICE_EXTENSION),&dev_name,FILE_DEVICE_UNKNOWN,FILE_DEVICE_SECURE_OPEN,FALSE,&dev_obj);

  if (NT_SUCCESS(status))
  {
    /*
     for some selected major functions we set handlers
    */
    DriverObject->MajorFunction[IRP_MJ_CREATE]          = hook9_create;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]           = hook9_close;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = hook9_device_control;
    DriverObject->DriverUnload                          = hook9_unload;

    /*
     there are few types of contexts, driver context, device context, 
     instance context etc. 
     second parameter of IoCreateDevice specifies size of device context 
     we will use this memory to store our global list of rules

     we use global variable for device context
    */
    global_dev_ext=(PDEVICE_EXTENSION)dev_obj->DeviceExtension;

    if (global_dev_ext)
    {
      /*
       we use mutex as a synchronization mechanism to maintain our list 
       and another one for driver notifications together with synchronization event
       irp event is needed to synchronize user mode and driver irp handling
      */
      KeInitializeMutex(&global_dev_ext->rules_mutex,0);
      KeInitializeMutex(&global_dev_ext->notify_mutex,0);
      KeInitializeEvent(&global_dev_ext->notify_event,SynchronizationEvent,FALSE);
      KeInitializeEvent(&global_dev_ext->notify_irp_event,NotificationEvent,FALSE);
      /*
       initialize other context variables
      */
      global_dev_ext->first_rule=NULL;
      global_dev_ext->last_rule=NULL;
      global_dev_ext->notify_irp=NULL;


      /*
       this selects the method of IO, we use buffered IO as it is comfortable 
       and effective for smaller packets
      */
      dev_obj->Flags|=DO_BUFFERED_IO;

      /*
       if we want user mode application to communicate our driver we need to make 
       a dos device link
      */
      status=IoCreateSymbolicLink(&dos_dev_name,&dev_name);
      if (!NT_SUCCESS(status))
      {
        DbgMsg("hook9.cpp: DriverEntry error: IoCreateSymbolicLink failed with status 0x%.8X",status);
        IoDeleteDevice(DriverObject->DeviceObject);
      }
    } else 
    {
      DbgMsg("hook9.cpp: DriverEntry error: no device extension");
      IoDeleteDevice(DriverObject->DeviceObject);
      status=STATUS_NO_SUCH_DEVICE;
    }
  } else DbgMsg("hook9.cpp: DriverEntry error: IoCreateDevice failed with status 0x%.8X",status);

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