/*
 hook3 is driver with sample device extension and communication with user mode 
 application using DeviceIoControl, it maintain a list of rules for pids 
 in device extension but does nothing with the list yet
*/

extern "C" //check part 1 - hook1.cpp for comment on this
{

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

NTSTATUS hook3_create(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook3_close(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook3_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID hook3_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 hook3_create(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook3.cpp: hook3_create(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

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

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

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

  DbgMsg("hook3.cpp: hook3_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 hook3_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook3.cpp: hook3_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("hook3.cpp: hook3_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:
      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("hook3.cpp: hook3_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 hook3_unload(PDRIVER_OBJECT DriverObject)
{    
  DbgMsg("hook3.cpp: hook3_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("hook3.cpp: hook3_unload(-)");
}


/*
 DriverEntry is common driver entry point
*/

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject,PUNICODE_STRING RegistryPath)
{
  DbgMsg("hook3.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]          = hook3_create;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]           = hook3_close;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = hook3_device_control;
    DriverObject->DriverUnload                          = hook3_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
    */
    PDEVICE_EXTENSION dev_ext=(PDEVICE_EXTENSION)dev_obj->DeviceExtension;

    if (dev_ext)
    {
      /*
       we use mutex as a synchronization mechanism to maintain our list
      */
      KeInitializeMutex(&dev_ext->rules_mutex,0);
      dev_ext->first_rule=NULL;
      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("hook3.cpp: DriverEntry error: IoCreateSymbolicLink failed with status 0x%.8X",status);
        IoDeleteDevice(DriverObject->DeviceObject);
      }
    } else 
    {
      DbgMsg("hook3.cpp: DriverEntry error: no device extension");
      IoDeleteDevice(DriverObject->DeviceObject);
      status=STATUS_NO_SUCH_DEVICE;
    }
  } else DbgMsg("hook3.cpp: DriverEntry error: IoCreateDevice failed with status 0x%.8X",status);

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