/*
 *	PatchFinder for Windows 2000
 *	Joanna Rutkowska, joanna at mailsnare dot net
 *	(c) 2003
 *
 */
#include <ntddk.h>
#include <stdio.h>
#include "driver.h"
#include "sst.h"
#include "interrupt.h"
#include "dbprotect.h"
#include "..\inc\kernel_iface.h"
#include "..\inc\config.h"

PNTPROC serviceTbl;
PCHAR	argumentTbl;
NTPROC origService;
HANDLE traceThreadId;	// id of the thread which will be traced
int origKiSystemService;
int origDebugHandler;
unsigned PF_traps;		// the most important dword...
unsigned PF_traps_frozen;

#ifdef TRACEPATH
volatile int* PF_trace = NULL;
unsigned PF_tracemaxlen = 4*50000;
#endif

int pfsrv_started = 0;	// pfService cannot be restarted!

__declspec(naked) NewKiSystemService()
{
	__asm {
		
		pushad
		push	fs
		mov		bx, 0x30			// kernel fs
		mov		fs, bx
		push	ds
		push	es


		mov		eax, fs:[0x124];	// pETHREAD
		mov		eax, [eax+0x1e4];	// Cid
		cmp		eax, traceThreadId;
		jne		goon;
		
		pushfd;
		pop		eax;
		or		eax, 0x100;
		push	eax;
		popfd;
		nop;
		nop;
		nop;

goon:	
		pop es
		pop ds
		pop fs
		popad
		jmp		origKiSystemService;
	}
}




__declspec(naked) NewDebugHandler2()
{
	__asm {
		cli;
		pushfd;
		push eax;
		mov eax, DR6;
		test ax, 0x100f;	// BD |B3|B2|B1|B0
		jz goon;
							// FIXME: do something more inteligent
		int 3;				// then just crashing the system!
goon:
		push ecx;
		// FIXME: push _KERNEL_DS
#ifdef TRACEPATH
		mov ecx, PF_trace;
		jecxz pf_debug_end;
#endif

		// FIXME: move before 'goon', eax is already set with DR6
		mov eax, DR6;		// check BS (signgle step) bit in DR6
		test ah, 0x40;
		jz pf_debug_end;

#ifdef TRACEPATH
		mov eax, PF_traps;
		shl eax, 2;
		cmp eax, PF_tracemaxlen;
		jg pf_debug_end;

		add ecx, eax;
		// FIXME: esp + X, change X when put some pushes before
		// FIXME: maybe better it'll be to use ebp...
		mov eax, dword ptr [esp + 12];	// get the debug handler retaddr
		
		mov dword ptr [ecx], eax;
#endif		
		inc	PF_traps;
pf_debug_end:	
		pop ecx;
		pop eax;
		popfd;	
		sti;
		iretd;
	}
}


__declspec(naked) NewDebugHandler1()
{
	__asm {
		jmp NewDebugHandler2;
	}
}





void OnUnload( IN PDRIVER_OBJECT DriverObject ) {

#ifdef DR_PROTECTION
	KdPrint(("pf: can not be unloaded!\n"));
	return;
#else
	KdPrint(("pf: unloading...\n"));

	KdPrint(("pf: removing pfServie...\n"));
	serviceTbl [pfServiceNo] = (NTPROC) origService;

	KdPrint(("pf: unhooking int 0x2e...\n"));
	setIntHandler (NT_SYSTEM_SERVICE_INT, origKiSystemService);

	KdPrint(("pf: unhooking debug int...\n"));
	setIntHandler (NT_DEBUG_INT, origDebugHandler);

#ifdef TRACEPATH
	KdPrint(("pf: freeing PF_trace...\n"));
	if (PF_trace) 
		MmFreeNonCachedMemory ((PVOID)PF_trace, PF_tracemaxlen);

	KdPrint(("pf: unloaded\n"));
#endif
#endif
}


NTSTATUS pfService(int magic, int cmd, int *tracebuf) {		
#ifdef TRACEPATH
	int nb;
#endif
	// assuming that origService takes at least 2 arguments! (FIXME: autocheck)
	if (magic == serviceMagic) {
		switch (cmd) {
			
			case PF_START:
				traceThreadId = PsGetCurrentThreadId();
				PF_traps = 0;
				PF_traps_frozen = 0;
				KdPrint (("PF_START: threadID = %#x, argSz = %#x\n",
					traceThreadId, argumentTbl[pfServiceNo]));
				return 0;
				// we don't care about 'ret xx', becouse
				// KiSystemService is doing 'mov esp, ebp'
				// after returning from our function. 
			
			case PF_GET_COUNT:
				if (PF_traps_frozen > 0) 
					return PF_traps_frozen;	// to guarantee that
							// (pfStop();...;pfStop()) == pfStop();

				traceThreadId = (HANDLE)0xffffffff;
				PF_traps_frozen = PF_traps;			
				
				KdPrint (("PF_GET_COUNT: PF_traps = %#x, frozen = %#x\n",
					PF_traps, PF_traps_frozen));
				
				return PF_traps_frozen;

			case PF_GET_TRACE:
				
				KdPrint (("PF_GET_TRACE: PF_traps_frozen = %#x\n", PF_traps_frozen));
#ifdef TRACEPATH
				nb = ((unsigned)PF_traps_frozen < (PF_tracemaxlen/sizeof(int))) ? 
				(PF_traps_frozen*sizeof(int)) : PF_tracemaxlen;
				
				KdPrint (("PF_GET_TRACE: tracebuf = %#x, PF_trace = %#x, nb = %#x\n",
					tracebuf, PF_trace, nb));
				memcpy ((char*)tracebuf, (char*)PF_trace, nb);
#endif
				return PF_traps_frozen;

			case PF_QUERY:
				KdPrint (("PF_QUERY\n"));
				return PF_ANSWER;

			case PF_SERVICE_STARTED:
				return pfsrv_started++;

			default:
				KdPrint (("pfService called with unknown cmd\n"));
			
		}
		
		return 0;

	}
	_asm {
		mov		esp, ebp;
		pop		ebp;
		jmp		origService;
	}
	
}






NTSTATUS DriverEntry( IN PDRIVER_OBJECT driverObject,
		      IN PUNICODE_STRING registryPath ) {
    
	KdPrint (("pf: loading...\n"));
	
	serviceTbl = KeServiceDescriptorTable->ntoskrnl.ServiceTable;
	argumentTbl = KeServiceDescriptorTable->ntoskrnl.ArgumentTable;
	driverObject->DriverUnload  = OnUnload; 
	traceThreadId = (HANDLE)0xffffffff;


	KdPrint (("pf: installing at service no. %#x...\n", pfServiceNo));
    origService = serviceTbl [pfServiceNo];
	serviceTbl [pfServiceNo] = (NTPROC) pfService;
	KdPrint (("pf: servTbl[%#x] at %#x\n",
		pfServiceNo, &serviceTbl[pfServiceNo] ));

	
	KdPrint (("pf: hooking int 0x2e...\n"));
	origKiSystemService = getIntHandler (NT_SYSTEM_SERVICE_INT);
	setIntHandler (NT_SYSTEM_SERVICE_INT, (int)NewKiSystemService);
	

	KdPrint (("pf: hooking debug int...\n"));
	origDebugHandler = getIntHandler (NT_DEBUG_INT);
	setIntHandler (NT_DEBUG_INT, (int)NewDebugHandler1);
	KdPrint (("pf: idt[%d] at %#x\n",
		NT_DEBUG_INT, getIntGateAddr(NT_DEBUG_INT) ));


#ifdef TRACEPATH
	KdPrint (("pf: allocating memory for PF_trace buffer (%d bytes)\n",
		PF_tracemaxlen));
	PF_trace = MmAllocateNonCachedMemory (PF_tracemaxlen);	// FIXME
	
	KdPrint (("pf: PF_trace at %#x\n", PF_trace));
#endif
	
#ifdef DR_PROTECTION	
	KdPrint (("pf: setting hardware protection...\n"));

	dbProtect (DB_DR0, (int)getIntGateAddr(NT_DEBUG_INT), DB_LEN_4B, DB_PROT_WRITE);
	dbProtect (DB_DR1, (int)getIntGateAddr(NT_DEBUG_INT)+4, DB_LEN_4B, DB_PROT_WRITE);
	dbProtect (DB_DR2, (int)NewDebugHandler1, DB_LEN_4B, DB_PROT_RW);
	dbSetGeneralProtection ();
#endif

	
	
	
	KdPrint (("pf: loaded succesfully.\n"));
	return 0;	
}