/*
 hook2 is basic hook driver, hooks ZwOpenProcess in SDT
 hook can be seen in DebugView
*/

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

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


NTSTATUS hook2_create(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook2_close(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS hook2_device_control(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID hook2_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 hook2_create(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
  DbgMsg("hook2.cpp: hook2_create(DeviceObject:0x%.8X,Irp:0x%.8X)",DeviceObject,Irp);

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

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

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

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

  PIO_STACK_LOCATION stack=IoGetCurrentIrpStackLocation(Irp);

  NTSTATUS status=STATUS_SUCCESS;
  Irp->IoStatus.Information=0;

  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("hook2.cpp: hook2_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_BUFFER comm_buf=(PDRVCOMM_BUFFER)buf_out;
 
  /*
   for every command we implement the functionality 
  */
  switch (code)
  {
    case IOCTL_HOOK_START:
      ret=hooking_hook();
      comm_buf->status=ret;
      Irp->IoStatus.Information=sizeof(DRVCOMM_BUFFER);
      break;

    case IOCTL_HOOK_STOP:
      ret=hooking_unhook();
      comm_buf->status=ret;
      Irp->IoStatus.Information=sizeof(DRVCOMM_BUFFER);
      break;

    /*
     unknown codes should also be handled
    */
    default:
      status=STATUS_INVALID_DEVICE_REQUEST;
  }

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

  DbgMsg("hook2.cpp: hook2_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 hook2_unload(PDRIVER_OBJECT DriverObject)
{    
  DbgMsg("hook2.cpp: hook2_unload(DriverObject:0x%.8X)",DriverObject);

  hooking_unhook();

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

  DbgMsg("hook2.cpp: hook2_unload(-)");
}


/*
 DriverEntry is common driver entry point
*/

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject,PUNICODE_STRING RegistryPath)
{
  DbgMsg("hook2.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,0,&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]          = hook2_create;
    DriverObject->MajorFunction[IRP_MJ_CLOSE]           = hook2_close;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]  = hook2_device_control;
    DriverObject->DriverUnload                          = hook2_unload;

    /*
     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("hook2.cpp: DriverEntry error: IoCreateSymbolicLink failed with status 0x%.8X",status);
      IoDeleteDevice(DriverObject->DeviceObject);
    }
  } else DbgMsg("hook2.cpp: DriverEntry error: IoCreateDevice failed with status 0x%.8X",status);

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