/*
 driver loader is used to install the driver via service manager 
 it also sends some commands to driver
 usage is simple: load driver_name driver_display_name path
 you can use loadhook.bat 

 second thread was added to receive events from driver 
 it uses asynchronous IO to manage this
 see description in the main driver file for more
*/

#include <windows.h>
#include <stdio.h>
#include <ddk/ntddk.h>
#include "../drvcomm.h"


int work=0;                                     //indicator of application status
HANDLE notify_event;                            //sync event between notify thread and main thread
int notify_type;                                //type of driver notification
DWORD notify_data;                              //driver notification data
HANDLE notify_dev;                              //notify thread device file

/*
 second thread manages notifications from driver
 we do not take any action if something in this thread failed
 however, for the functionality of our program this thread is important 
 so in real world program this should be handled too
 also synchronization is missing in this example
*/

DWORD WINAPI notify_thread(LPVOID args)
{
  printf("Notify thread runs\n");
  HANDLE notify_dev=CreateFile(args,GENERIC_READ | GENERIC_WRITE,FILE_SHARE_READ | FILE_SHARE_WRITE,
                               NULL,OPEN_EXISTING,FILE_FLAG_OVERLAPPED,NULL);
  if (notify_dev==INVALID_HANDLE_VALUE) notify_dev=0;
  if (notify_dev)
  {
    OVERLAPPED overlapped;
    overlapped.hEvent=CreateEvent(NULL,TRUE,FALSE,NULL);

    if (overlapped.hEvent)
    {
      while (work)
      {
        DRVCOMM_REQUEST_BUFFER buf_req;
        DRVCOMM_RESPONSE_BUFFER buf_res;
        memset(&buf_req,0,sizeof(buf_req));
        memset(&buf_res,0,sizeof(buf_res));

        DWORD bytes;
        int ret=DeviceIoControl(notify_dev,IOCTL_NOTIFY,NULL,0,&buf_res,sizeof(buf_res),&bytes,&overlapped);

        int err=ret?0:GetLastError();

        if (err==ERROR_IO_PENDING)
        {
          if (GetOverlappedResult(notify_dev,&overlapped,&bytes,TRUE)) err=0;
          else err=GetLastError();
        }

        if (err)
        {
          if (err==ERROR_OPERATION_ABORTED) printf("Notify thread received exit signal\n");
          else printf("Notify thread error: unexpected DeviceIoControl error %d\n",err);
          break;
        }

        if (!bytes || !buf_res.status)
        {
          printf("Notify thread communication error: bytes=%d; buf_res.status=0x%.8X\n",bytes,buf_res.status);
          break;
        }

        /*
         we've received notification from driver
         we'll contact second thread for user input 
        */
        HKEY key;
        char name[512],*svc_name,imagepath[512],svc_path[512];
        DWORD imagepath_len=sizeof(imagepath);
        memset(imagepath,0,imagepath_len);
        err=0;

        switch (buf_res.type)
        {
          case RESPONSE_NOTIFY_LOAD_DRIVER_CHECK:
            /*
             we've got a system registry path to Services in HKLM, we open it and read ImagePath from there
            */
            notify_type=RESPONSE_NOTIFY_LOAD_DRIVER_CHECK;
            notify_data=(ULONG)&buf_res.parameters.load_driver_check.regpath;

            err=3;
            sprintf(name,"%S\0",&buf_res.parameters.load_driver_check.regpath);
            svc_name=strrchr(name,'\\');
            svc_name++;

            if (svc_name && *svc_name)
            {
              err--;
              sprintf(svc_path,"SYSTEM\\CurrentControlSet\\Services\\%s\0",svc_name);

              if (ret=RegOpenKeyEx(HKEY_LOCAL_MACHINE,svc_path,0,KEY_QUERY_VALUE,&key)==ERROR_SUCCESS)
              {
                err--;
                printf("Notify thread: RESPONSE_NOTIFY_LOAD_DRIVER_CHECK received, you should make a decision:\n");

                if (ret=RegQueryValueEx(key,"ImagePath",NULL,NULL,imagepath,&imagepath_len)==ERROR_SUCCESS)
                {
                  err--;
                  printf("permit loading driver ServiceName=%s ImagePath=%s for process pid=%d?\n",
                         svc_name,imagepath,buf_res.parameters.load_driver_check.caller_pid);

                  PulseEvent(notify_event);
                  WaitForSingleObject(notify_event,INFINITE);
                  printf("Notify thread wait done, answer = %s\n",notify_data?"deny":"permit");
                }
              }
            }

            if (err)
            {
              printf("Notify thread: RESPONSE_NOTIFY_LOAD_DRIVER_CHECK received, but error occurred: ");
              switch (err)
              {
                case 1:printf("RegQueryValueEx for ImagePath failed with code %d",ret); break;
                case 2:printf("RegOpenKeyEx failed with code %d",ret); break;
                case 3:printf("invalid service name"); break;
              }

              //something went wrong -> deny access
              printf(" -> denying access");
              notify_data=0;
            }
            buf_req.parameters.load_driver_check.permit=notify_data;
            ret=DeviceIoControl(notify_dev,IOCTL_NOTIFY_LOAD_DRIVER_CHECK,&buf_req,sizeof(buf_req),NULL,0,&bytes,&overlapped);

            if (!ret && (GetLastError()==ERROR_IO_PENDING)) GetOverlappedResult(notify_dev,&overlapped,&bytes,TRUE);
            break;


          case RESPONSE_NOTIFY_PHYSICAL_MEMORY_CHECK:
            printf("Notify thread: RESPONSE_NOTIFY_PROCESS_CHECK received, you should make a decision:\n");
            notify_type=RESPONSE_NOTIFY_PHYSICAL_MEMORY_CHECK;
            notify_data=0;

            printf("permit 0x%.8X access for process pid=%d to \\Device\\PhysicalMemory section object\n",
                   buf_res.parameters.physical_memory_check.access,buf_res.parameters.physical_memory_check.caller_pid);

            PulseEvent(notify_event);
            WaitForSingleObject(notify_event,INFINITE);
            printf("Notify thread wait done, answer = %s\n",notify_data?"deny":"permit");

            buf_req.parameters.physical_memory_check.permit=notify_data;
            ret=DeviceIoControl(notify_dev,IOCTL_NOTIFY_PHYSICAL_MEMORY_CHECK,&buf_req,sizeof(buf_req),NULL,0,&bytes,&overlapped);

            if (!ret && (GetLastError()==ERROR_IO_PENDING)) GetOverlappedResult(notify_dev,&overlapped,&bytes,TRUE);
            break;


          case RESPONSE_NOTIFY_PROCESS_CHECK:
            printf("Notify thread: RESPONSE_NOTIFY_PROCESS_CHECK received, you should make a decision:\n");
            notify_type=RESPONSE_NOTIFY_PROCESS_CHECK;
            notify_data=buf_res.parameters.process_check.pid;

            printf("permit 0x%.8X access for process pid=%d to process pid=%d?\n",
                   buf_res.parameters.process_check.access,buf_res.parameters.process_check.caller_pid,notify_data);

            PulseEvent(notify_event);
            WaitForSingleObject(notify_event,INFINITE);
            printf("Notify thread wait done, answer = %s\n",notify_data?"deny":"permit");

            buf_req.parameters.process_check.permit=notify_data;
            ret=DeviceIoControl(notify_dev,IOCTL_NOTIFY_PROCESS_CHECK,&buf_req,sizeof(buf_req),NULL,0,&bytes,&overlapped);

            if (!ret && (GetLastError()==ERROR_IO_PENDING)) GetOverlappedResult(notify_dev,&overlapped,&bytes,TRUE);
            break;


          case RESPONSE_NOTIFY_FILE_CHECK:
            printf("Notify thread: RESPONSE_NOTIFY_FILE_CHECK received, you should make a decision:\n");
            notify_type=RESPONSE_NOTIFY_FILE_CHECK;
            notify_data=(ULONG)buf_res.parameters.file_check.name;

            printf("permit 0x%.8X access for process pid=%d to file=%S?\n",
                   buf_res.parameters.file_check.access,buf_res.parameters.file_check.caller_pid,notify_data);

            PulseEvent(notify_event);
            WaitForSingleObject(notify_event,INFINITE);
            printf("Notify thread wait done, answer = %s\n",notify_data?"deny":"permit");

            buf_req.parameters.file_check.permit=notify_data;
            ret=DeviceIoControl(notify_dev,IOCTL_NOTIFY_FILE_CHECK,&buf_req,sizeof(buf_req),NULL,0,&bytes,&overlapped);

            if (!ret && (GetLastError()==ERROR_IO_PENDING)) GetOverlappedResult(notify_dev,&overlapped,&bytes,TRUE);
            break;


          default:
            printf("Notify thread error: unknown response buffer type\n");
        }
      }
      CloseHandle(overlapped.hEvent);
    } else printf("Notify thread error: unable to create event\n");

    CloseHandle(notify_dev);
  } else printf("Notify thread error: unable to connect device\n");

  printf("Notify thread terminates\n");
  ExitThread(0);
}


/*
 converts X:\directory\file.ext to \Device\???\directory\file.ext
 where ??? could be HarddiskVolumeX or CdRomX or any other device
 returns 0 if error, nonzero otherwise
*/