#include <stdio.h>
#include <stdlib.h>
#include <windows.h>
#include <mmsystem.h>
#include <d3d9.h>
#include "Zdisasm.h"

#pragma comment (lib, "winmm.lib")

#define FUNC_NUM	0x86 // SetStreamSource fonction 86
#define OFFSET_TABLE	0x15C38
#define VERTEX_NUM	40 // nombre de sommet dans le programme

LPDIRECT3DDEVICE9		gDevice = NULL;
LPDIRECT3DVERTEXBUFFER9 	gVertexBuff = NULL;
PVOID				pVertices = NULL; // buffer contenant les couleurs
DWORD				gVertexSize; // taille d'un sommet
DWORD				saveSize = 0; // taille des opcodes sauvegardés de SetStreamSource

void init();
PBYTE hook(PBYTE SetStreamSource, DWORD SetStreamSourceHook, DWORD offset);
void unhook(PBYTE SetStreamSource, PBYTE save);
PIMAGE_NT_HEADERS PEHeaderFromHModule(HMODULE hModule);
void __declspec() SetStreamSourceHook();


BOOL APIENTRY DllMain(HINSTANCE hModule, DWORD dwReason, PVOID lpReserved)
{
	switch (dwReason) 
	{
	case DLL_PROCESS_ATTACH:
		init();
		break;
	}
	return TRUE;
}

void init() {
	HMODULE modD3d9;
	PIMAGE_NT_HEADERS headD3d9;
	PIMAGE_OPTIONAL_HEADER headD3d9Op;
	PBYTE SetStreamSource, save = NULL;
	PDWORD baseTableD3DDEVICE;
	DWORD baseCodeSection, hookFuncSize = 0, InstSize, OldProtect;

	modD3d9 = GetModuleHandle("d3d9.dll");
	headD3d9 = PEHeaderFromHModule(modD3d9);
	headD3d9Op = &headD3d9->OptionalHeader;
	baseCodeSection = headD3d9Op->ImageBase + headD3d9Op->BaseOfCode;
	baseTableD3DDEVICE = (DWORD *)(baseCodeSection + OFFSET_TABLE);
	SetStreamSource = (PBYTE)*(baseTableD3DDEVICE + FUNC_NUM);

	while(*((PBYTE)(&SetStreamSourceHook)+hookFuncSize) != 0x90) {
		GetInstLenght((PDWORD)((PBYTE)&SetStreamSourceHook+hookFuncSize), &InstSize);
		hookFuncSize+=InstSize;
	}

	save = hook(SetStreamSource, (DWORD)&SetStreamSourceHook, hookFuncSize);
	Sleep(100);

	if(save != NULL) {
		unhook(SetStreamSource, save);
	} else return;
	if(gVertexBuff == NULL && gVertexSize == NULL) return;

	
	gVertexBuff->Lock(0, 10, (void **)&pVertices, 0);
	gVertexBuff->Unlock();

	// Je fais mumuse avec les couleurs, on peu faire une multitude d'autre chose...
	if(VirtualProtect(pVertices, VERTEX_NUM*gVertexSize, PAGE_EXECUTE_READWRITE, &OldProtect) != 0) {
		DWORD color;
		
		for(int i=0; i<VERTEX_NUM; i++) { // 12 | 21
			if(i == 0 || i == 12 || i == 24) {
				color = 0xFFFFFFFF & timeGetTime();
				Sleep(100);
			}
			CopyMemory((PDWORD)((DWORD)pVertices + gVertexSize -4 + i*gVertexSize), &color, 4);
		}
		VirtualProtect(pVertices, VERTEX_NUM*gVertexSize, OldProtect, &OldProtect);
	} 
}

PIMAGE_NT_HEADERS PEHeaderFromHModule(HMODULE hModule)
{
	PIMAGE_NT_HEADERS pNTHeader = 0;

	__try  {
		if(PIMAGE_DOS_HEADER(hModule)->e_magic != IMAGE_DOS_SIGNATURE) return 0;
		pNTHeader = PIMAGE_NT_HEADERS(PBYTE(hModule) + PIMAGE_DOS_HEADER(hModule)->e_lfanew);
		if(pNTHeader->Signature != IMAGE_NT_SIGNATURE) pNTHeader = 0;
	}
	__except(EXCEPTION_EXECUTE_HANDLER) { return 0; }

	return pNTHeader;
}

void __declspec() SetStreamSourceHook() {
	__asm { 
		push	ebp;
		mov	ebp, esp;
		push	eax;
		mov	eax, [ebp+4*2];
		mov	gDevice, eax;
		mov	eax, [ebp+4*4];
		mov	gVertexBuff, eax;
		mov	eax, [ebp+4*6];
		mov	gVertexSize, eax;
		pop	eax;
		pop	ebp;
		nop;	// laisser les 5 nop
		nop;	// qui seront remplacé par un jump
		nop;
		nop;
		nop;
	}
}

PBYTE hook(PBYTE SetStreamSource, DWORD SetStreamSourceHook, DWORD offset) {
	PBYTE save = NULL;
	DWORD OldProtect, InstSize;
	BYTE jump[5];
	
	jump[0] = 0xE9;

	// hook de la fonction SetStreamSource
	*((PDWORD)(jump+1)) = SetStreamSourceHook - (DWORD)SetStreamSource - 5;
	while(saveSize<5) {
		GetInstLenght((PDWORD)((PBYTE)SetStreamSource+saveSize), &InstSize);
		saveSize+=InstSize;
	} 
	// On sauvegarde le debut de la fonction
	save = (PBYTE)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, saveSize+5);
	if(save == NULL) return NULL;
	CopyMemory((PVOID)(save), SetStreamSource, saveSize);
	// On place le jump
	if(VirtualProtect(SetStreamSource, saveSize, PAGE_EXECUTE_READWRITE, &OldProtect) == 0) return NULL;
	FillMemory(SetStreamSource, saveSize, 0x90);
	CopyMemory(SetStreamSource, jump, 5);
	VirtualProtect(SetStreamSource, saveSize, OldProtect, &OldProtect);

	// On place un jump vers les opcode sauvegardé à la fin de la fonction de hook
	if(VirtualProtect((LPVOID)((PBYTE)SetStreamSourceHook+offset), 5, PAGE_EXECUTE_READWRITE, &OldProtect) == 0) return NULL;
	*((PDWORD)(jump+1)) = (DWORD)save - (DWORD)((PBYTE)SetStreamSourceHook+offset) - 5;
	CopyMemory((PVOID)((PBYTE)SetStreamSourceHook+offset), jump, 5);
	VirtualProtect((LPVOID)((PBYTE)SetStreamSourceHook+offset), 5, OldProtect, &OldProtect);

	// jump à la fin des opcode sauvegardé vers le reste de la fonction SetStreamSource
	*((PDWORD)(jump+1)) = (DWORD)SetStreamSource - (DWORD)save - 5;
	CopyMemory((PVOID)(save+saveSize), jump, 5);

	return save;
}

void unhook(PBYTE SetStreamSource, PBYTE save) {
	DWORD OldProtect;
	
	VirtualProtect((PVOID)SetStreamSource, saveSize, PAGE_EXECUTE_READWRITE, &OldProtect);		
	CopyMemory((PVOID)SetStreamSource, save, saveSize);		
	VirtualProtect((PVOID)SetStreamSource, saveSize, OldProtect, &OldProtect);

	HeapFree(GetProcessHeap(), HEAP_NO_SERIALIZE, save);
}