#pragma once
#include "loaddrv.h"


#include <windows.h>
#include <WinSvc.h>
#include <stdlib.h>
#include <string.h>
#include <WinIoCtl.h>
#define FILE_DEVICE_FILEMON	0x00008300
//
// Version #
//
#define FILEMONVERSION    430

#define LOGBUFSIZE ((ULONG)(64*0x400-(3*sizeof(ULONG)+1)))

#define IOCTL_FILEMON_SETDRIVES   (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x00, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_ZEROSTATS   (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x01, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_GETSTATS    (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x02, METHOD_NEITHER, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_UNLOADQUERY (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x03, METHOD_NEITHER, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_STOPFILTER  (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x04, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_STARTFILTER (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x05, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_SETFILTER   (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x06, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_VERSION     (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x07, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_HOOKSPECIAL (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x08, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_UNHOOKSPECIAL (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x09, METHOD_BUFFERED, FILE_ANY_ACCESS )

#define IOCTL_FILEMON_ADDPROTECTEDFILE (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x0A, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_DELPROTECTEDFILE (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x0B, METHOD_BUFFERED, FILE_ANY_ACCESS )
#define IOCTL_FILEMON_GETPROTECTEDFILELIST (ULONG) CTL_CODE( FILE_DEVICE_FILEMON, 0x0C, METHOD_BUFFERED, FILE_ANY_ACCESS )

typedef struct _FilemonNameParam{
	char  *SourceFileName;
	char  *CacheFileName;
}FilemonNameParam,*PFilemonNameParam;




class CProtectDrv :
	public CLoadDrv
{
public:
	CProtectDrv(LPCTSTR driveName,LPCTSTR sysPath);
public:
	~CProtectDrv(void);

	BOOL StartFilter();

	BOOL HookDrivers()
	{
		try
		{
			DWORD nb;
			ULONG HookedDrive = 0;
			DWORD drive = GetLogicalDrives();
			if(! DeviceIoControl(m_sysHandle, IOCTL_FILEMON_ZEROSTATS,
				NULL, 0, NULL, 0, &nb, NULL ))
			{
				printf("zero stats unsuccess!\n");

				return FALSE;
			}
			if(! DeviceIoControl(m_sysHandle,IOCTL_FILEMON_SETDRIVES,
				&drive,sizeof(drive),&HookedDrive,sizeof(HookedDrive),&nb,NULL))
				return FALSE;
			else
				return TRUE;
		}
		catch (...)
		{
			return FALSE;
		}
		return TRUE;
	}

	BOOL SetFile(char*path,bool bAdd=true)
	{
		DWORD nb;
		FilemonNameParam Dir;
		Dir.SourceFileName = (char *)malloc(256);
		strcpy(Dir.SourceFileName,path);
		Dir.CacheFileName = NULL;
		if( !DeviceIoControl(m_sysHandle,bAdd?IOCTL_FILEMON_ADDPROTECTEDFILE:IOCTL_FILEMON_DELPROTECTEDFILE,
			&Dir,sizeof(FilemonNameParam),NULL,
			0,&nb,NULL) )
		{
			free(Dir.SourceFileName);
			return FALSE;
		}
		else{
			free(Dir.SourceFileName);
			return TRUE;
		}
	}

	BOOL GetProtectFiles(LPVOID LPOutBuffer,LPDWORD lpLen)
	{
		DWORD length = 0;
		if (  !DeviceIoControl(m_sysHandle, IOCTL_FILEMON_GETPROTECTEDFILELIST,
			NULL, 0, LPOutBuffer, LOGBUFSIZE,
			&length, NULL ) )
		{
			return FALSE;
		}
		else
		{
			*lpLen = length;
			return TRUE;
		}
	}
};