/*
 hook6 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
*/

extern "C"
{

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

NTSTATUS hook6_create(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook6_close(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook6_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID hook6_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;


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

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

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

  DbgMsg("hook6.cpp: hook6_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 hook6_close(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook6.cpp: hook6_close(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

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

  DbgMsg("hook6.cpp: hook6_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 hook6_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook6.cpp: hook6_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("hook6.cpp: hook6_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;
 

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

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

    case IOCTL_HOOK_STOP:
      ret=hooking_unhook();
      buf_res->status=ret;
      break;

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

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

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

  DbgMsg("hook6.cpp: hook6_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 hook6_unload(PDRIVER_OBJECT DriverObject)
{    
  DbgMsg("hook6.cpp: hook6_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("hook6.cpp: hook6_unload(-)");
}


/*
 DriverEntry is common driver entry point
*/

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject,PUNICODE_STRING RegistryPath)
{
  DbgMsg("hook6.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]          = hook6_create;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]           = hook6_close;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = hook6_device_control;
    DriverObject->DriverUnload                          = hook6_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
      */
      KeInitializeMutex(&global_dev_ext->rules_mutex,0);
      global_dev_ext->first_rule=NULL;
      global_dev_ext->last_rule=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("hook6.cpp: DriverEntry error: IoCreateSymbolicLink failed with status 0x%.8X",status);
        IoDeleteDevice(DriverObject->DeviceObject);
      }
    } else 
    {
      DbgMsg("hook6.cpp: DriverEntry error: no device extension");
      IoDeleteDevice(DriverObject->DeviceObject);
      status=STATUS_NO_SUCH_DEVICE;
    }
  } else DbgMsg("hook6.cpp: DriverEntry error: IoCreateDevice failed with status 0x%.8X",status);

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