extern "C"
{

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

PFUNC_RULE func_find_rule(PDEVICE_EXTENSION dev_ext,ULONG type,ULONG id);
int func_add_rule(PDEVICE_EXTENSION dev_ext,PFUNC_RULE rule);
int func_del_rule(PDEVICE_EXTENSION dev_ext,PFUNC_RULE rule);
int func_free_list(PDEVICE_EXTENSION dev_ext);
#ifdef DEBUG
void func_print_rules(PDEVICE_EXTENSION dev_ext);
#endif
int func_protect_process(PDEVICE_EXTENSION dev_ext,ULONG pid,int enable);
int func_protect_file(PDEVICE_EXTENSION dev_ext,wchar_t *name,int enable);
int func_check_process_protection(ULONG pid);
int func_check_file_protection(PWSTR name);
int func_notify_ask_user(PDEVICE_EXTENSION dev_ext,ULONG caller_pid,ULONG type,ULONG id,ULONG access);
int func_init(void);
int func_get_object_name_by_handle(HANDLE handle,UCHAR *buffer,int *length);
int func_is_good_read_ptr(PVOID buf,ULONG size);
} //extern "C"


/*
 global device extension is global variable for device extension
 we need it to be able to access device extenstion data - list of rules 
 from hooked function
*/

PDEVICE_EXTENSION global_dev_ext=NULL;


/*
 address of protected section object \Device\PhysicalMemory
*/

PVOID func_protected_device_physical_memory=NULL;


/*
 find rule returns pointer to rule object in the list of rules if appropriate
 rule with given type and id is found
*/

PFUNC_RULE func_find_rule(PDEVICE_EXTENSION dev_ext,ULONG type,ULONG id)
{
  DbgMsg("func.cpp: func_find_rule(dev_ext:0x%.8X,type:0x%.8X,id:0x%.8X)",dev_ext,type,id);

  PFUNC_RULE item=dev_ext->first_rule,res=NULL;
  int done=0,file_len;

  //to have faster file compare we use length of its name too
  if (type==RULE_TYPE_FILE_PROTECTION)
    file_len=(wcslen((wchar_t *)id)+1)*sizeof(wchar_t);

  while (item)
  {
    if (item->type==type)
    {
      done=((type==RULE_TYPE_PROCESS_PROTECTION) && (item->process.pid==id)
         || (type==RULE_TYPE_FILE_PROTECTION)    && (item->file.length==file_len) && (!wcsicmp(item->file.name,(wchar_t *)id)));
      if (done)
      {
        res=item;
        break;
      }
    }
    item=item->next;
  }

  DbgMsg("func.cpp: func_find_rule(-):0x%.8X",res);
  return res;
}


/*
 add rule adds rule to the end of the list
*/

int func_add_rule(PDEVICE_EXTENSION dev_ext,PFUNC_RULE rule)
{
  DbgMsg("func.cpp: func_add_rule(dev_ext:0x%.8X,rule:0x%.8X)",dev_ext,rule);

  if (dev_ext->last_rule) dev_ext->last_rule->next=rule;
  else dev_ext->first_rule=rule;

  rule->prev=dev_ext->last_rule;
  rule->next=NULL;
  dev_ext->last_rule=rule;

  DbgMsg("func.cpp: func_add_rule(-):TRUE");
  return TRUE;
}


/*
 del rule deletes rule from the list
*/

int func_del_rule(PDEVICE_EXTENSION dev_ext,PFUNC_RULE rule)
{
  DbgMsg("func.cpp: func_del_rule(dev_ext:0x%.8X,rule:0x%.8X)",dev_ext,rule);
  if (rule->next) rule->next->prev=rule->prev;
  else dev_ext->last_rule=rule->prev;

  if (rule->prev) rule->prev->next=rule->next;
  else dev_ext->first_rule=rule->next;

  ExFreePool(rule);

  DbgMsg("func.cpp: func_del_rule(-):TRUE");
  return TRUE;
}


/*
 free list deletes whole rule list and frees memory
*/

int func_free_list(PDEVICE_EXTENSION dev_ext)
{
  DbgMsg("func.cpp: func_free_list(dev_ext:0x%.8X)",dev_ext);

  /*
   write access to the shared memory always requires protection 
   we will use mutex to make such synchronization here
  */
  NTSTATUS status=KeWaitForMutexObject(&dev_ext->rules_mutex,Executive,KernelMode,FALSE,NULL);

  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_free_list error: KeWaitForMutexObject failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_free_list(-):FALSE");
    return FALSE;
  }

  while (dev_ext->first_rule)
    func_del_rule(dev_ext,dev_ext->first_rule);

  KeReleaseMutex(&dev_ext->rules_mutex,FALSE);

  DbgMsg("func.cpp: func_free_list(-):TRUE");
  return TRUE;
}


#ifdef DEBUG
/*
 print rules prints all rules as debug messages 
*/

void func_print_rules(PDEVICE_EXTENSION dev_ext)
{
  DbgMsg("func.cpp: func_print_rules(dev_ext:0x%.8X)",dev_ext);

  PFUNC_RULE rule=dev_ext->first_rule;
  int i=0;
  while (rule)
  {
    switch (rule->type)
    {
      case RULE_TYPE_PROCESS_PROTECTION:DbgMsg("%.3d) addr:0x%.8X, type:PROC PROT, pid:%d",i,rule,rule->process.pid); break;
      case RULE_TYPE_FILE_PROTECTION:DbgMsg("%.3d) addr:0x%.8X, type:PROC PROT, length:%d, name:%S",i,rule,rule->file.length,rule->file.name); break;
      default:DbgMsg("%.3d) unknown rule",i);
    }
    rule=rule->next;
    i++;
  }
  DbgMsg("func.cpp: func_print_rules(-)");
  return;
}
#endif


/*
 protect process thread manages list of rules 
 it enables/disables protection rule for specific process id
*/

int func_protect_process(PDEVICE_EXTENSION dev_ext,ULONG pid,int enable)
{
  DbgMsg("func.cpp: func_protect_process(dev_ext:0x%.8X,id:0x%.8X,enable:%d)",
         dev_ext,pid,enable);

  NTSTATUS status=KeWaitForMutexObject(&dev_ext->rules_mutex,Executive,KernelMode,FALSE,NULL);

  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_protect_process error: KeWaitForMutexObject failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_protect_process(-):FALSE");
    return FALSE;
  }

  int res=FALSE;
  ULONG type=RULE_TYPE_PROCESS_PROTECTION;

  PFUNC_RULE rule=func_find_rule(dev_ext,type,pid);

  if (enable)
  {
    if (!rule)
    {
      rule=(PFUNC_RULE)ExAllocatePool(PagedPool,sizeof(FUNC_RULE));
      if (rule)
      {
        rule->type=type;
        rule->process.pid=pid;
        res=func_add_rule(dev_ext,rule);
      } else DbgMsg("func.cpp: func_protect_process error: ExAllocatePool failed");
    } else
    {
      DbgMsg("func.cpp: func_protect_process error: rule exists");
      res=ERROR_RULE_EXISTS;
    }
  } else
  {
    if (!rule)
    {
      DbgMsg("func.cpp: func_protect_process error: rule does not exist");
      res=ERROR_RULE_DOES_NOT_EXIST;
    } else res=func_del_rule(dev_ext,rule);
  } 

#ifdef DEBUG
  func_print_rules(dev_ext);
#endif

  KeReleaseMutex(&dev_ext->rules_mutex,FALSE);
  DbgMsg("func.cpp: func_protect_process(-):0x%.8X",res);

  return res;
}


/*
 protect file manages list of rules 
 it enables/disables protection rule for specific file
*/

int func_protect_file(PDEVICE_EXTENSION dev_ext,wchar_t *name,int enable)
{
  DbgMsg("func.cpp: func_protect_file(dev_ext:0x%.8X,name:%S,enable:%d)",
         dev_ext,name,enable);

  NTSTATUS status=KeWaitForMutexObject(&dev_ext->rules_mutex,Executive,KernelMode,FALSE,NULL);

  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_protect_file error: KeWaitForMutexObject failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_protect_file(-):FALSE");
    return FALSE;
  }

  int res=FALSE,len;
  ULONG type=RULE_TYPE_FILE_PROTECTION;
  len=(wcslen(name)+1)*sizeof(wchar_t);

  PFUNC_RULE rule=func_find_rule(dev_ext,type,(ULONG)name);

  if (enable)
  {
    if (!rule)
    {
      rule=(PFUNC_RULE)ExAllocatePool(PagedPool,sizeof(FUNC_RULE));
      if (rule)
      {
        rule->type=type;
        RtlCopyMemory(rule->file.name,name,len);
        rule->file.length=len;
        res=func_add_rule(dev_ext,rule);
      } else DbgMsg("func.cpp: func_protect_file error: ExAllocatePool failed");
    } else
    {
      DbgMsg("func.cpp: func_protect_file error: rule exists");
      res=ERROR_RULE_EXISTS;
    }
  } else
  {
    if (!rule)
    {
      DbgMsg("func.cpp: func_protect_file error: rule does not exist");
      res=ERROR_RULE_DOES_NOT_EXIST;
    } else res=func_del_rule(dev_ext,rule);
  } 

#ifdef DEBUG
  func_print_rules(dev_ext);
#endif

  KeReleaseMutex(&dev_ext->rules_mutex,FALSE);
  DbgMsg("func.cpp: func_protect_file(-):0x%.8X",res);

  return res;
}


/*
 check for process protection, 
 returns TRUE if pid is protected
*/

int func_check_process_protection(ULONG pid)
{
  DbgMsg("func.cpp: func_check_process_protection(pid:%d)",pid);

  NTSTATUS status=KeWaitForMutexObject(&global_dev_ext->rules_mutex,Executive,KernelMode,FALSE,NULL);
  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_check_process_protection error: KeWaitForMutexObject failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_check_process_protection(-):FALSE");
    return FALSE;
  }

  PFUNC_RULE rule=func_find_rule(global_dev_ext,RULE_TYPE_PROCESS_PROTECTION,pid);
  KeReleaseMutex(&global_dev_ext->rules_mutex,FALSE);

  int res=rule!=NULL;
  DbgMsg("func.cpp: func_check_process_protection(-):%d",res);
  return res;
}


/*
 check for file protection
 returns TRUE if file is protected
*/

int func_check_file_protection(PWSTR name)
{
  DbgMsg("func.cpp: func_check_file_protection(name:%S)",name);

  NTSTATUS status=KeWaitForMutexObject(&global_dev_ext->rules_mutex,Executive,KernelMode,FALSE,NULL);
  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_check_file_protection error: KeWaitForMutexObject failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_check_file_protection(-):FALSE");
    return FALSE;
  }

  PFUNC_RULE rule=func_find_rule(global_dev_ext,RULE_TYPE_FILE_PROTECTION,(ULONG)name);
  KeReleaseMutex(&global_dev_ext->rules_mutex,FALSE);

  int res=rule!=NULL;
  DbgMsg("func.cpp: func_check_file_protection(-):%d",res);
  return res;
}


/*
 asks user mode application for permission for protected object 
 using driver notification implemented as async IO
*/

int func_notify_ask_user(PDEVICE_EXTENSION dev_ext,ULONG caller_pid,ULONG type,ULONG id,ULONG access)
{
  DbgMsg("func.cpp: func_notify_ask_user: dev_ext:0x%.8X,caller_pid:0x%.8X,type:0x%.8X,id:0x%.8X,access:0x%.8X",
         dev_ext,caller_pid,type,id,access);

  /*
   we have to protect this part with mutex because our implementation can handle 
   only one pending IRP at the moment
  */
  NTSTATUS status=KeWaitForMutexObject(&dev_ext->notify_mutex,Executive,KernelMode,FALSE,NULL);
  if (!NT_SUCCESS(status))
  {
    DbgMsg("func.cpp: func_notify_ask_user error: KeWaitForMutexObject wait for notify_mutex failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_notify_ask_user(-):FALSE");
    return FALSE;
  }

  /*
   second protection with event here is because we have to wait for IRP from user mode application
   user mode application sends IOCTL_NOTIFY and receives status pending
   this event protects one IRP to be processed twice 
   using the event caller of this function waits for our app to complete its stuff 
   with old results and send new IOCTL_NOTIFY
  */
  status=KeWaitForSingleObject(&dev_ext->notify_irp_event,Executive,KernelMode,FALSE,NULL);
  if (!NT_SUCCESS(status))
  {
    KeReleaseMutex(&dev_ext->notify_mutex,FALSE);
    DbgMsg("func.cpp: func_notify_ask_user error: KeWaitForSingleObject wait for notify_irp_event failed with status 0x%.8X",status);
    DbgMsg("func.cpp: func_notify_ask_user(-):FALSE");
    return FALSE;
  }

  /*
   we have to clear irp event before we leave this mutex-protected code
   then we read currently pending irp and clear it in device context
   to better handle possible errors
  */
  KeClearEvent(&dev_ext->notify_irp_event);

  PIRP irp=dev_ext->notify_irp;

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

  UCHAR *buf_out=(UCHAR *)irp->AssociatedIrp.SystemBuffer;
  PDRVCOMM_RESPONSE_BUFFER buf_res=(PDRVCOMM_RESPONSE_BUFFER)buf_out;
  RtlZeroMemory(buf_res,sizeof(DRVCOMM_RESPONSE_BUFFER));
  buf_res->status=1;

  int len;
  PUNICODE_STRING regpath;
  switch (type)
  {
    case RULE_TYPE_LOAD_DRIVER_PROTECTION:
      buf_res->type=RESPONSE_NOTIFY_LOAD_DRIVER_CHECK;
      buf_res->parameters.load_driver_check.caller_pid=caller_pid;
                    
      regpath=(PUNICODE_STRING)id;
      len=min(sizeof(buf_res->parameters.load_driver_check.regpath)-2,regpath->Length);
      RtlCopyMemory(buf_res->parameters.load_driver_check.regpath,regpath->Buffer,len);
      buf_res->parameters.load_driver_check.regpath[len/2]=0x0000;
      break;

    case RULE_TYPE_PHYSICAL_MEMORY_PROTECTION:
      buf_res->type=RESPONSE_NOTIFY_PHYSICAL_MEMORY_CHECK;
      buf_res->parameters.physical_memory_check.caller_pid=caller_pid;
      buf_res->parameters.physical_memory_check.access=access;


    case RULE_TYPE_PROCESS_PROTECTION:
      buf_res->type=RESPONSE_NOTIFY_PROCESS_CHECK;
      buf_res->parameters.process_check.caller_pid=caller_pid;
      buf_res->parameters.process_check.pid=id;
      buf_res->parameters.process_check.access=access;
      break;

    case RULE_TYPE_FILE_PROTECTION:
      buf_res->type=RESPONSE_NOTIFY_FILE_CHECK;
      len=min(sizeof(buf_res->parameters.file_check.name),wcslen((wchar_t *)id)*sizeof(wchar_t));

      RtlCopyMemory(buf_res->parameters.file_check.name,(PVOID)id,len);
      buf_res->parameters.file_check.name[len/2]=0x0000;

      buf_res->parameters.file_check.caller_pid=caller_pid;
      buf_res->parameters.file_check.access=access;
      break;
  }
 
  //complete the request 
  irp->IoStatus.Information=sizeof(DRVCOMM_RESPONSE_BUFFER);
  irp->IoStatus.Status=status;
  IoCompleteRequest(irp,IO_NO_INCREMENT);

  //and wait for answer 
  KeWaitForSingleObject(&dev_ext->notify_event,Executive,KernelMode,FALSE,NULL);
  
  int res=dev_ext->notify_answer;

  KeReleaseMutex(&dev_ext->notify_mutex,FALSE);

  DbgMsg("func.cpp: func_notify_ask_user(-):%d",res);
  return res;
}


/*
 initialization for protected objects
 returns TRUE on success
*/

int func_init(void)
{
  DbgMsg("func.cpp: func_init()");
  int res=FALSE;

  OBJECT_ATTRIBUTES oa;
  HANDLE section;
  UNICODE_STRING sec_name;
  RtlInitUnicodeString(&sec_name,FUNC_DEVICE_PHYSICAL_MEMORY);
  InitializeObjectAttributes(&oa,&sec_name,0,NULL,NULL);

  NTSTATUS status=ZwOpenSection(&section,SECTION_ALL_ACCESS,&oa);
  if (NT_SUCCESS(status))
  {
    PVOID obj;
    status=ObReferenceObjectByHandle(section,0,0,KernelMode,&obj,NULL);

    if (NT_SUCCESS(status))
    {
      func_protected_device_physical_memory=obj;
      res=func_protected_device_physical_memory!=NULL;
      DbgMsg("func.cpp: func_init: func_protected_device_physical_memory=0x%.8X",func_protected_device_physical_memory);
      ObDereferenceObject(obj);
    } else DbgMsg("func.cpp: func_init error: ObReferenceObjectByHandle failed with status 0x%.8X",status);

    ZwClose(section);
  } else DbgMsg("func.cpp: func_init error: ZwOpenSection failed with status 0x%.8X",status);

  DbgMsg("func.cpp: func_init(-):%d",res);
  return res;
}



/*
 this function returns name of object from its handle 
 if the function succeed the return value is TRUE and length is set to number of bytes written 
 to buffer including null termination
 if the buffer is not long enough or other error occur the return value is FALSE
*/

int func_get_object_name_by_handle(HANDLE handle,UCHAR *buffer,int *length)
{
  DbgMsg("func.cpp: func_get_object_name_by_handle(handle:0x%.8X;buffer:0x%.8X;length:0x%.8X)",handle,buffer,length);

  int res=FALSE;

  PVOID obj;
  NTSTATUS status=ObReferenceObjectByHandle(handle,0,0,KernelMode,&obj,NULL);
  if (NT_SUCCESS(status))
  {
     PUNICODE_STRING uni_name=(PUNICODE_STRING)buffer;
     POBJECT_NAME_INFORMATION oni=(POBJECT_NAME_INFORMATION)uni_name;

     ULONG ret_len=0;

     status=ObQueryNameString(obj,oni,*length,&ret_len);
     if (NT_SUCCESS(status))
     {
       if (uni_name->Buffer)
       {
         int slen=wcslen(uni_name->Buffer);
         *length=(slen+1)*sizeof(wchar_t);
         RtlCopyMemory(buffer,uni_name->Buffer,*length);
       } else 
       {
         //noname objects
         buffer[0]=0x00;
         buffer[1]=0x00;
         *length=2;
       }
       res=TRUE;
     }

     ObDereferenceObject(obj);
  }

  DbgMsg("func.cpp: func_get_object_name_by_handle(-):%d",res);
  return res;
}


/*
 this function checks user buffer for read access 
 returns true if the buffer is ok
*/

int func_is_good_read_ptr(PVOID buf,ULONG size)
{
  DbgMsg("func.cpp: func_is_good_read_ptr(buf:0x%.8X;size:0x%.8X)",buf,size);

  int res=TRUE;

  __try
  {
    ProbeForRead(buf,size,sizeof(char));
    ULONG sum=0;
    PULONG p=(PULONG)buf;
    int i;
    for (i=0;i<(int)(size/sizeof(ULONG));i++) sum+=p[i];
    for (int j=0;j<(int)(size%sizeof(ULONG));j++) sum+=*((UCHAR*)&p[i]+j);
  } __except(EXCEPTION_EXECUTE_HANDLER)
  {
    DbgPrint("func.cpp: func_is_good_read_ptr error: exception occurred");
    res=FALSE;
  }

  DbgMsg("func.cpp: func_is_good_read_ptr(-):%d",res);
  return res;
}