extern "C"
{

#include <wdm.h>
#include "func.h"
#include "debug.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_check_process_protection(ULONG pid);

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;


/*
 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;
  while (item)
  {
    if ((item->type==type) && (item->process.pid==id))
    {
      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;
      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(-):%d",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;
}


/*
 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;
}