%include "icedump.inc"
%include "vxdn.inc"
%include "wiat.inc"
%include "util.mac"
%include "win32n.inc"


global PbpmSysDynamicDeviceInit
global PbpmSysDynamicDeviceExit
global PbpmThreadInit
global PbpmThreadNotExecuteable
global PbpmVmInit
global PbpmVmNotExecuteable

global Parse_Pbpm
global Service_Pbpm


extern sdata
extern Parser.error
extern Parser.errorMsg
extern ParseAddress
extern ParseExpression
extern Error_V86
extern Error_PM16
extern Error_PMR0
extern SetCB
extern VWIN32.W32_GetThreadContext
extern VWIN32.W32_SetThreadContext
extern VWIN32.W32ServiceTable


bits 32


DRx_size	equ	CONTEXT.cx_Dr7 - CONTEXT.cx_Dr0 + 4

segment _LDATA
init_data_begin ICEDUMP_PBPM
init_data_item	AllocateThreadDataSlot,		FreeThreadDataSlot
init_data_item	HookVxDServices,		UnhookVxDServices
init_data_item	HookW32Services,		UnhookW32Services
init_data_end


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
PbpmSysDynamicDeviceInit:
	init_construct ICEDUMP_PBPM
	jc	.error

	clc
	retn

.error:
	call	PbpmSysDynamicDeviceExit

	stc
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
PbpmSysDynamicDeviceExit:
	init_destruct ICEDUMP_PBPM
	clc
	retn


;-------------------------------------------------------------------------------
; init TDS in new thread
;
; EDI: R0TCB
;
; stc on error
;-------------------------------------------------------------------------------
PbpmThreadInit:
	mov	eax,[TDS]
	and	dword [eax+edi],byte 0
	clc
	retn


;-------------------------------------------------------------------------------
; EDI: R0TCB
;
; stc on error
;-------------------------------------------------------------------------------
PbpmThreadNotExecuteable:
	push	ecx
	push	edx

	mov	ecx,[TDS]
	xor	eax,eax
	xchg	eax,[edi+ecx]
	or	eax,eax
	jz	@F

	VMMCall	_HeapFree, eax, byte 0

@@
	pop	edx
	pop	ecx
	clc
	retn


;-------------------------------------------------------------------------------
; ebx: VMCB
;
; stc on error
;-------------------------------------------------------------------------------
PbpmVmInit:
	VMMCall	Get_Initial_Thread_Handle
	mov	eax,[TDS]
	and	dword [eax+edi],byte 0
	retn


;-------------------------------------------------------------------------------
; ebx: VMCB
;
; stc on error
;-------------------------------------------------------------------------------
PbpmVmNotExecuteable:
	VMMCall	Get_Initial_Thread_Handle
	call	PbpmThreadNotExecuteable
	retn


;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
AllocateThreadDataSlot:
	push	ebx
	push	ecx
	push	edx
	push	edi

	VMMCall	_AllocateThreadDataSlot
	or	eax,eax
	jz	.error

	mov	[TDS],eax
	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi
	mov	ecx,[TDS]

@@
	and	dword [ecx+edi],byte 0
	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	@B

	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	clc
	retn

.error:
	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	stc
	retn


segment _LDATA
	align 4
TDS:	dd 0


segment _LTEXT
;-------------------------------------------------------------------------------
; stc on error
;-------------------------------------------------------------------------------
FreeThreadDataSlot:
	push	ebx
	push	ecx
	push	edx
	push	edi

	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi

.loop:
	call	PbpmThreadNotExecuteable

	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	.loop

	VMMCall	_FreeThreadDataSlot, dword [TDS]

	pop	edi
	pop	edx
	pop	ecx
	pop	ebx
	clc
	retn


;-------------------------------------------------------------------------------
; hook VWIN32 Get/Set Thread Context VxD APIs
;
; stc on error
;-------------------------------------------------------------------------------
HookVxDServices:
	push	esi

	GetDeviceServiceOrdinal eax, _VWIN32_Get_Thread_Context
	mov	esi,HookedGetThreadContext
	VMMCall	Hook_Device_Service
	jc	near .0

	GetDeviceServiceOrdinal eax, _VWIN32_Set_Thread_Context
	mov	esi,HookedSetThreadContext
	VMMCall	Hook_Device_Service
	jc	near .1

	pop	esi
	clc
	retn

.1:
	GetDeviceServiceOrdinal eax, _VWIN32_Get_Thread_Context
	mov	esi,HookedGetThreadContext
	VMMCall	Unhook_Device_Service

.0:
	pop	esi
	stc
	retn


;-------------------------------------------------------------------------------
; unhook VWIN32 Get/Set Thread Context VxD APIs
;
; stc on error
;-------------------------------------------------------------------------------
UnhookVxDServices:
	push	esi

	GetDeviceServiceOrdinal eax, _VWIN32_Set_Thread_Context
	mov	esi,HookedSetThreadContext
	VMMCall	Unhook_Device_Service

	GetDeviceServiceOrdinal eax, _VWIN32_Get_Thread_Context
	mov	esi,HookedGetThreadContext
	VMMCall	Unhook_Device_Service

	pop	esi
	clc
	retn


;-------------------------------------------------------------------------------
; hook VWIN32 Get/Set Thread Context W32 APIs
;
; stc on error
;-------------------------------------------------------------------------------
HookW32Services:
	mov	eax,[VWIN32.W32_GetThreadContext]
	lea	eax,[eax*8+8]
	add	eax,[VWIN32.W32ServiceTable]
	push	dword [eax]
	pop	dword [OrgW32GetThreadContext]
	mov	dword [eax],HookedW32GetThreadContext

	mov	eax,[VWIN32.W32_SetThreadContext]
	lea	eax,[eax*8+8]
	add	eax,[VWIN32.W32ServiceTable]
	push	dword [eax]
	pop	dword [OrgW32SetThreadContext]
	mov	dword [eax],HookedW32SetThreadContext

	clc
	retn


;-------------------------------------------------------------------------------
; unhook VWIN32 Get/Set Thread Context W32 APIs
;
; stc on error
;-------------------------------------------------------------------------------
UnhookW32Services:
	mov	eax,[VWIN32.W32_GetThreadContext]
	lea	eax,[eax*8+8]
	add	eax,[VWIN32.W32ServiceTable]
	push	dword [OrgW32GetThreadContext]
	pop	dword [eax]

	mov	eax,[VWIN32.W32_SetThreadContext]
	lea	eax,[eax*8+8]
	add	eax,[VWIN32.W32ServiceTable]
	push	dword [OrgW32SetThreadContext]
	pop	dword [eax]

	clc
	retn


;-------------------------------------------------------------------------------
; /PBPM [ON|OFF] | coded on summer 2oo2 by ^DAEMON^ | thx to the_owl for help!
;-------------------------------------------------------------------------------
segment _LTEXT
Parse_Pbpm:
	mov	edi,Error_V86
	mov	ebp,[dClient_EFLAGS]
	test	byte [ebp+2],2		; is client in V86 mode?
	jnz	near Parser.errorMsg

	mov	edi,Error_PM16
	mov	ebp,[dClient_CS]
	lar	eax,[ebp]		; is client 32 bit?
	bt	eax,22
	jnc	near Parser.errorMsg

	mov	edi,Error_PMR0
	test	byte [ebp],3		; is client in ring-0?
	jz	near Parser.errorMsg

	call	[pSkipWhiteSpace]	; skip to <ON|OFF>
	jnz	@F

.status:
	mov	esi,status
	call	[pPrintToCommandWindow]
	popad
	retn

@@
	call	[pIsItOnOrOff]
	jb	near Parser.error

	test	al,al
	jz	.off

	cmp	dword [status.onoff],'ON'
	je	.status

	mov	dword [status.onoff],'ON'
	jmp	short .ret

.off:
	cmp	dword [status.onoff],'OFF'
	je	.status

	mov	dword [status.onoff],'OFF'

.ret:
	push	byte SERVICE_PBPM
	mov	ebp,[dClient_EAX]
	pop	dword [ebp]

	call	SetCB

	mov	ebp,[fExecuteMoreCommands]	; set internal Winice flag to 0
	mov	byte [ebp],0

	popad
	retn


;-------------------------------------------------------------------------------------
;
; real service | hot stuff ?
; 
;-------------------------------------------------------------------------------------
Service_Pbpm:
	cmp	dword [status.onoff],'OFF'
	jne	.off

	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi
	mov	ecx,[TDS]

@@
	VMMCall	_HeapAllocate, byte DRx_size, byte HEAPZEROINIT
	mov	[edi+ecx],eax

	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	@B

	jmp	short .ret

.off:
	VMMCall	Get_Sys_VM_Handle
	VMMCall	Get_Initial_Thread_Handle
	mov	ebx,edi

.loop:
	call	PbpmThreadNotExecuteable

	VMMCall	Get_Next_Thread_Handle
	cmp	ebx,edi
	jnz	.loop

.ret:
	popfd
	popad
	retn


;-----------------------------------------------------------------------------
;
;.HookedGetThreadContext:
;
; -> Win32 (_vwin32_get_thread_context)
;
;-----------------------------------------------------------------------------
HookedW32GetThreadContext:
	push	esi
	push	edi
	push	dword [esp+24]		; CONTEXT
	push	dword [esp+24]		; R0TCB, can be NULL
	push	dword [esp+24]
	push	dword [esp+24]
	call	[OrgW32GetThreadContext]

	mov	esi,[esp+24]

	VMMCall	Get_Cur_Thread_Handle

	add	edi,[TDS]
	mov	edi,[edi]
	or	edi,edi
	jz	@F

;---------------------------------------------------------------------------------------
;
; insert fake values ;)
;
;---------------------------------------------------------------------------------------
	push	ecx
	mov	ecx,6
	xchg	esi,edi
	add	edi,04h
	rep	movsd
	pop	ecx

@@:	pop	edi
	pop	esi
	retn	16


segment _LDATA
	align 4
OrgW32GetThreadContext dd 0


segment _LTEXT

;-----------------------------------------------------------------------------
;
;.HookedGetThreadContext:
;
; -> VxD (_vwin32_get_thread_context)
;
;-----------------------------------------------------------------------------

	jmp	short HookedGetThreadContext
	jmp	[OrgGetThreadContext]

HookedGetThreadContext:
	push	esi
	push	edi

	push	dword [esp+16]		; CONTEXT
	push	dword [esp+16]		; R0TCB, can be NULL
	call	[OrgGetThreadContext]
	pop	edi			; oh yeah, these little undoc tricks ;-)
	pop	esi

	add	edi,[TDS]
	mov	edi,[edi]		 ; TraceInfo
	or	edi,edi
	jz	@F


;---------------------------------------------------------------------------------------
;
; insert fake values ;)
;
;---------------------------------------------------------------------------------------
	push	ecx
	mov	ecx,6
	xchg	esi,edi
	add	edi,04h
	repz	movsd
	push	ecx

@@:	pop	edi
	pop	esi
	retn

segment _LDATA
	align 4
OrgGetThreadContext: dd 0


segment _LTEXT
;-----------------------------------------------------------------------------
;
;.HookedSetThreadContext:
;
; -> VxD (_vwin32_get_thread_context)
;
;-----------------------------------------------------------------------------

	jmp	short HookedSetThreadContext
	jmp	[OrgSetThreadContext]

HookedSetThreadContext:
	push	eax
	push	ebx
	push	esi
	push	edi
	push	ebp

	mov	eax,[esp+28]		; CONTEXT
	call	haxx0r

	push	dword [esp+28]		; CONTEXT
	push	dword [esp+28]		; R0TCB, can be NULL
	call	[OrgSetThreadContext]
	pop	edi
	pop	esi

	pop	ebp
	pop	edi
	pop	esi
	pop	ebx
	pop	eax
	retn

segment _LDATA
	align 4
OrgSetThreadContext: dd 0


segment _LTEXT
;-----------------------------------------------------------------------------
;
;.HookedSetThreadContext:
;
; -> Win32 (_vwin32_set_thread_context)
;
;-----------------------------------------------------------------------------
HookedW32SetThreadContext:
	push	eax
	push	ebx
	push	esi
	push	edi
	push	ebp

	mov	eax,[esp+36]		; CONTEXT
	call	haxx0r

	push	dword [esp+36]		; CONTEXT
	push	dword [esp+36]		; R0TCB, can be NULL
	push	dword [esp+36]
	push	dword [esp+36]
	call	[OrgW32SetThreadContext]

	pop	ebp
	pop	edi
	pop	esi
	pop	ebx
	pop	eax
	retn	16


segment _LDATA
	align 4
OrgW32SetThreadContext dd 0


segment _LTEXT

haxx0r:
	lea	esi, [eax+4h]

	VMMCall	Get_Cur_Thread_Handle		; grab R0TCB
	add	edi,[TDS]
	mov	eax,edi
	cmp	dword [edi],byte 0
	jnz	.save

	cmp	dword [status.onoff],'ON'
	je	@F

	retn

@@
	push	ecx
	push	edx

	VMMCall	_HeapAllocate, byte DRx_size, byte HEAPZEROINIT

	pop	edx
	pop	ecx

	test	eax,eax
	jnz	@F

	retn

@@
	mov	[edi],eax

.save:
	mov	edi,[edi]
;	Trace_Out "ICEDUMP: BPM Protection : prevented modification"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR0 = #eax"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR1 = #eax"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR2 = #eax"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR3 = #eax"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR6 = #eax"
	lodsd
	and	dword [esi-4],0
	stosd
;	Trace_Out "			DR7 = #eax"
;	mov	eax,[eax+0b8h]			; eip
;	Trace_Out "			EIP = #eax"
	retn


segment _LDATA
status		db 'PBPM is ',
.onoff		db 'OFF',0
