// testHook.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
//#include "resources.h"

#include <windows.h>
#include <stdio.h>
#include <process.h>
#include <ImageHlp.h>
#include "ntundoc.h"

#pragma comment(lib,"Advapi32.lib")
#pragma comment(lib,"ImageHlp.lib")
#pragma comment(lib,"User32.lib")


/* Function de substitution aux API */

int __stdcall HookFnOpen(HKEY, LPCTSTR, DWORD, REGSAM, PHKEY);
int __stdcall HookFnQuery(HKEY, LPCTSTR, LPDWORD, LPDWORD, LPBYTE, LPDWORD );
int __stdcall HookFnCreate(HKEY, LPCTSTR, ULONG, LPTSTR, ULONG, ULONG, LPSECURITY_ATTRIBUTES, PHKEY, PULONG);


/* Constances ACCESS_TYPE */

#define ACCESS_CREATE 0
#define ACCESS_OPEN   1
#define ACCESS_READ   2


#define REG_WRITE "Tentative d'écriture dans la base de registre :"
#define REG_PROG  "Programme concerné :"  
#define REG_QUST  "Autoriser l'écriture ?"


typedef struct {

    char szProgName[64], szProgPath[256], szUserName[32], szKeyName[128], 
			szKeyValue[128], szTypeKey[30], szTime[30];

	DWORD dwAccessType, dwAccessGranted, 
			dwReply, dwCount, dwPID;
} DATA, *PDATA;

bool MakeHook(bool, PROC, PROC);

inline void ExtractProgName(char *, char *);
inline void GetMainKey(HKEY, char *);
inline void GetAlertTime(char *);

inline HWND GetRGHwnd();

HINSTANCE hThisMod;
HHOOK hHook; HWND hWnd;

PROC *AddressOfFunction;

const char *szDllName = "Advapi32.dll";
char szProgPath[256]; 

DWORD dwCount = 0;
DATA Data;

PROCESS_INFORMATION ProcessInfo; //Structure PROCESS_INFORMATION
STARTUPINFO StartupInfo; //Structure STARTUPINFO

bool MakeHook(bool bToActive, PROC Address, PROC FunctionHook)
{	
	if(bToActive)
	{
	   ULONG uSize;
	   HANDLE hproc= ProcessInfo.hProcess;	

       PIMAGE_IMPORT_DESCRIPTOR pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)
	   ImageDirectoryEntryToData(hproc, true, IMAGE_DIRECTORY_ENTRY_IMPORT, &uSize);
    
	   if(!pImportDesc) return false;

	   // Recherche de l'import descriptor correspondant à advapi32.dll
	   for(; pImportDesc->Name ; pImportDesc++) {

		  PSTR pszModName = (PSTR)((BYTE *) hproc + pImportDesc->Name);
          if(!_strcmpi((char *)pszModName, szDllName)) break;
	   }

	   if(!pImportDesc->Name) return false;
	
	   // Recupère l'address table de l'import desriptor 
       PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)((BYTE *) hproc + pImportDesc->FirstThunk);

	   // Liste toutes les fonctions référencées dans l'address table 
       for(; pThunk->u1.Function ; pThunk++) {

          PROC *pFn = (PROC *)&pThunk->u1.Function;
		  if(*pFn == Address) {
			  
			 AddressOfFunction = pFn;   // Overwrite à l'emplacement de l'adresse
			 break;
		  }
	   }

	   if(pThunk->u1.Function) {

	      DWORD dwOld;

	      // Remplacement de l'adresse en question...
	      VirtualProtect(AddressOfFunction, 4, PAGE_EXECUTE_READWRITE, &dwOld);
       
	      *AddressOfFunction = FunctionHook;
	      VirtualProtect(AddressOfFunction, 4, dwOld, 0);
	   }
	}

    else {

       DWORD dwOld;

	   // Set l'adresse d'origine
	   VirtualProtect(AddressOfFunction, 4, PAGE_EXECUTE_READWRITE, &dwOld);

	   *AddressOfFunction = Address;
	   VirtualProtect(AddressOfFunction, 4, dwOld, 0);
	}

	return true;
}

int __stdcall HookFnCreate(HKEY hKey, LPCTSTR lpSubKey, DWORD Reserved, LPTSTR lpClass, DWORD dwOptions,
						   REGSAM samDesired, LPSECURITY_ATTRIBUTES lpSecurityAttributes,
						   PHKEY phkResult, LPDWORD lpdwDisposition)
{
    char szDisplayDetect[1024], szProgName[128], szMainKey[6], szTime[30];
	int nReply = 0, nLen; 

	/* Récupère les infos de l'alerte */
    ExtractProgName(szProgPath, szProgName);
	ZeroMemory(&Data, sizeof(DATA));
	
	Data.dwPID = _getpid(); /* TEST PID TRANFERT */
	
	nLen = sizeof(Data.szUserName);
	GetUserName((LPWSTR)Data.szUserName, (unsigned long *)&nLen);
    
	GetMainKey(hKey, szMainKey);
    GetAlertTime(szTime);
    dwCount++;

	/* Remplit la structure DATA */
    _snprintf_s(Data.szKeyName, sizeof(Data.szKeyName), "%s%s", szMainKey, lpSubKey);
    strncpy_s(Data.szProgName, szProgName, sizeof(Data.szProgName));

	Data.dwAccessType = ACCESS_CREATE;   // Type d'accés CREATE (création dans la registry)
	Data.dwCount = dwCount;

	// Avertissement d'une tentative d'écriture
    _snprintf_s(szDisplayDetect, sizeof(szDisplayDetect), "%s %s\n%s %s\n\n%s", 
		      REG_WRITE, lpSubKey, REG_PROG, szProgName, REG_QUST);
   	
	nReply = MessageBox(GetRGHwnd(), (LPCWSTR)szDisplayDetect, (LPCWSTR)"Registry Guard - Détection", MB_ICONWARNING | MB_YESNO);

	if(nReply == IDYES)	Data.dwReply = (RegCreateKeyEx(hKey, lpSubKey, Reserved, lpClass, dwOptions, samDesired, 
	   lpSecurityAttributes, phkResult, lpdwDisposition) == ERROR_SUCCESS ? ERROR_SUCCESS : GetLastError());

	/* Envoi de l'alerte à RG */
    COPYDATASTRUCT cds;
	cds.cbData = sizeof(DATA);
	cds.lpData = &Data;
    
    //SendMessage(GetRGHwnd(), WM_COPYDATA, NULL, (LONG)&cds);

    return 0;
}

int __stdcall HookFnOpen(HKEY g, LPCTSTR lpSubKey, DWORD e, REGSAM z, PHKEY a)
{
	char szProgName[128], szMainKey[6], szTime[30];
	int nLen;

	/* Infos de l'alerte */
	ExtractProgName(szProgPath, szProgName);
    ZeroMemory(&Data, sizeof(DATA));

	Data.dwPID = _getpid(); 

	nLen = sizeof(Data.szUserName);
	GetUserName((LPWSTR)Data.szUserName, (unsigned long *)&nLen);

    GetMainKey(g, szMainKey);
	GetAlertTime(szTime);
    dwCount++;
	 
	/* Remplit DATA et envoi de l'alerte */
	_snprintf_s(Data.szKeyName, sizeof(Data.szKeyName), "%s%s", szMainKey, lpSubKey);
    strncpy_s(Data.szProgName, szProgName, sizeof(Data.szProgName));

    Data.dwReply = (RegOpenKeyEx(g, lpSubKey, e, z, a) == ERROR_SUCCESS ? ERROR_SUCCESS : GetLastError());
	Data.dwAccessType = ACCESS_OPEN;   // Type d'accés OPEN (ouverture dans la registry)
	Data.dwCount = dwCount;

    COPYDATASTRUCT cds;
	cds.cbData = sizeof(DATA);
	cds.lpData = &Data;

	SendMessage(GetRGHwnd(), WM_COPYDATA, 0, (LONG)&cds);

	return 0;
}

int _stdcall HookFnQuery(HKEY hKey, LPCTSTR lpValueName, LPDWORD lpReserved, LPDWORD lpType, 
						 LPBYTE lpData, LPDWORD lpcbData)
{
	char szProgName[128], szTime[30];
	int nLen;

	/* infos de l'alerte */
	ZeroMemory(&Data, sizeof(DATA));
	GetAlertTime(szTime);
    dwCount++;

	Data.dwPID = _getpid(); /* TEST PID TRANFERT */ 
	nLen = sizeof(Data.szUserName);
	GetUserName((LPWSTR)Data.szUserName, (unsigned long *)&nLen);
	
	ExtractProgName(szProgPath, szProgName);
    strncpy_s(Data.szProgName, szProgName, sizeof(Data.szProgName));

	strncpy_s(Data.szProgPath, szProgPath, sizeof(Data.szProgPath));
    
	/* Remplit DATA et envoi de l'alerte*/
	Data.dwAccessType = ACCESS_READ;
	Data.dwCount = dwCount;

	Data.dwReply = (RegQueryValueEx(hKey, lpValueName, lpReserved, lpType, lpData, lpcbData) == ERROR_SUCCESS
		            ? ERROR_SUCCESS : GetLastError());

    _snprintf_s(Data.szKeyName, sizeof(Data.szKeyName), "...\\%s", lpValueName);
	strncpy_s(Data.szKeyValue, (const char *)lpData, sizeof(Data.szKeyValue));

	COPYDATASTRUCT cds;
	cds.cbData = sizeof(DATA);
	cds.lpData = &Data;

	SendMessage(GetRGHwnd(), WM_COPYDATA, 0, (LONG)&cds);

	return 0;
}

inline void ExtractProgName(char *ProgPath, char *ProgName)
{
    char *cPtr;
	int j; 

	for(unsigned int i = 0 ; i < strlen(ProgPath) ; i++) {
		
		char cCurrent = ProgPath[i];
 		if(cCurrent == '\\') cPtr = &ProgPath[i+1];
	}

	for(j = 0 ; *cPtr != 0 ; j++) {
		
		ProgName[j] = *cPtr;
		cPtr++;
	}

	ProgName[j] = 0;
}

inline void GetMainKey(HKEY hKey, char *szMainKey)
{
	if(hKey == HKEY_CLASSES_ROOT) strcpy_s(szMainKey, sizeof(szMainKey), "HKCR\\");
	else if(hKey == HKEY_CURRENT_CONFIG) strcpy_s(szMainKey, sizeof(szMainKey), "HKCC\\");
	else if(hKey == HKEY_CURRENT_USER) strcpy_s(szMainKey, sizeof(szMainKey), "HKCU\\");
	else if(hKey == HKEY_LOCAL_MACHINE) strcpy_s(szMainKey, sizeof(szMainKey), "HKLM\\");
	else if(hKey == HKEY_USERS) strcpy_s(szMainKey, sizeof(szMainKey), "HKU\\");
}

inline void GetAlertTime(char *szTime)
{
    SYSTEMTIME sysTime;
	GetSystemTime(&sysTime);

   _snprintf_s(szTime, 30, 30, "%d:%d %d sec", sysTime.wHour+2, sysTime.wMinute, sysTime.wSecond);
    strncpy_s(Data.szTime, szTime, sizeof(Data.szTime)); 
}

inline HWND GetRGHwnd()
{
	HWND hWnd;
	hWnd = FindWindow((LPCWSTR)"RGWndClass", NULL);

	if(!hWnd) {

		MessageBox(NULL, (LPCWSTR)"Failed to find Registry Guard's top-most window", (LPCWSTR)"Fatal Error", MB_OK);
		return NULL;
	}
    
	return hWnd;
}

LRESULT CALLBACK HookProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	return CallNextHookEx(hHook, nCode, wParam, lParam);
}

extern "C" void InstallHook()
{   
	// Lance le hook
	if(!hHook) hHook = SetWindowsHookEx(WH_CBT, HookProc, hThisMod, 0);
}

extern "C" void UninstallHook()
{   
	// Stop le hook
	UnhookWindowsHookEx(hHook);
	hHook = NULL;
}

int main(int argc, char* argv[])
{
	BOOL bFuncRetn = FALSE;

	ZeroMemory( &ProcessInfo, sizeof(PROCESS_INFORMATION) );; //On complète StartupInfo de 0
	ZeroMemory( &StartupInfo, sizeof(STARTUPINFO) );
	
	bFuncRetn = CreateProcess(TEXT("C:\\Program Files\\Internet Explorer\\IEXPLORE.EXE"), NULL, NULL, NULL, TRUE, 0, NULL, NULL, &StartupInfo, &ProcessInfo); //On démarre le processus, notre programme
	if (bFuncRetn == 0){
		printf_s("CreateProcess failed\n");
	}
	 // Path du prog pour le logging
	GetModuleFileName(0, (LPWCH)szProgPath, sizeof(szProgPath));  

	hThisMod = (HINSTANCE)GetCurrentProcess();

	/* On lance les hooks */
	MakeHook(true, (PROC)RegQueryValueEx, (PROC)HookFnQuery);
	MakeHook(true, (PROC)RegCreateKeyEx, (PROC)HookFnCreate);  
	MakeHook(true, (PROC)RegOpenKeyEx, (PROC)HookFnOpen);


  return 0;
}