;@echo off
;goto make

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;
;  KbdSpy - This is an example of a simple legacy keyboard filter driver.
;
;  WARNING: You will fail to attach to USB-keyboard stack
;
;  We create two device objects. The first one is a control device. It provides
;  the interface to our user-mode client. The second device is a filter device.
;  It intercepts keyboard data being passed from the keyboard class driver.
;  So, it is an upper filter.  The intercepted data is collected in the list.
;  Upon timer triggering we signal shared event object to notify our user-mode
;  client about there is something interesting for it.  The user-mode client
;  issues control request and we copy all info into the buffer.
;
;  Written by Four-F (four-f@mail.ru)
;
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

.486
.model flat, stdcall
option casemap:none

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                  I N C L U D E   F I L E S                                        
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

include \masm32\include\w2k\ntstatus.inc
include \masm32\include\w2k\ntddk.inc
include \masm32\include\w2k\ntoskrnl.inc
include \masm32\include\w2k\ntddkbd.inc
include \masm32\include\w2k\hal.inc

includelib \masm32\lib\w2k\ntoskrnl.lib
includelib \masm32\lib\w2k\hal.lib

include \masm32\Macros\Strings.mac

include ..\common.inc

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                       M A C R O S                                                 
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

; spin lock macros

LOCK_ACQUIRE MACRO lck:REQ
	; Returns old IRQL in al
	IF (OPATTR (lck)) AND 00010000y
		;; Is a register value
		IFDIFI <lck>, <ecx>	;; don't move ecx onto itself
			mov ecx, lck
		ENDIF
	ELSEIF (OPATTR (lck)) AND 01000000y
		;; relative to SS
		lea ecx, lck
	ELSE
		mov ecx, offset lck
	ENDIF
	fastcall KfAcquireSpinLock, ecx
ENDM

LOCK_RELEASE MACRO lck:REQ, NewIrql:REQ
	IF (OPATTR (lck)) AND 00010000y
		;; Is a register value
		IFDIFI <lck>, <ecx>	;; don't move ecx onto itself
			mov ecx, lck
		ENDIF
	ELSEIF (OPATTR (lck)) AND 01000000y
		;; relative to SS
		lea ecx, lck
	ELSE
		mov ecx, offset lck
	ENDIF

	IFDIFI <NewIrql>, <dl>	;; don't move dl onto itself
		mov dl, NewIrql
	ENDIF

	.if dl == DISPATCH_LEVEL
		fastcall KefReleaseSpinLockFromDpcLevel, ecx
	.else
		and edx, 0FFh		;; for shure (KIRQL is BYTE)
		fastcall KfReleaseSpinLock, ecx, edx
	.endif
ENDM

; mutex macros

MUTEX_INIT MACRO mtx:REQ
	IF (OPATTR (mtx)) AND 00010000y
		;; Is a register value
		invoke KeInitializeMutex, mtx, 0
	ELSEIF (OPATTR (mtx)) AND 01000000y
		;; relative to SS
		invoke KeInitializeMutex, addr mtx, 0
	ELSE
		invoke KeInitializeMutex, offset mtx, 0
	ENDIF
ENDM

MUTEX_ACQUIRE MACRO mtx:REQ
	IF (OPATTR (mtx)) AND 00010000y
		;; Is a register value
		invoke KeWaitForMutexObject, mtx, Executive, KernelMode, FALSE, NULL
	ELSEIF (OPATTR (mtx)) AND 01000000y
		;; relative to SS
		invoke KeWaitForMutexObject, addr mtx, Executive, KernelMode, FALSE, NULL
	ELSE
		invoke KeWaitForMutexObject, offset mtx, Executive, KernelMode, FALSE, NULL
	ENDIF
ENDM

MUTEX_RELEASE MACRO mtx:REQ
	IF (OPATTR (mtx)) AND 00010000y
		;; Is a register value
		invoke KeReleaseMutex, mtx, FALSE
	ELSEIF (OPATTR (mtx)) AND 01000000y
		;; relative to SS
		invoke KeReleaseMutex, addr mtx, FALSE
	ELSE
		invoke KeReleaseMutex, offset mtx, FALSE
	ENDIF
ENDM

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                     S T R U C T U R E S                                           
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

KEY_DATA_ENTRY STRUCT
	ListEntry	LIST_ENTRY	<>	; For memory blocks tracking.
	KeyData		KEY_DATA	<>
KEY_DATA_ENTRY ENDS

FiDO_DEVICE_EXTENSION STRUCT

	; The top of the stack before this filter was added

	pNextLowerDeviceObject	PDEVICE_OBJECT	?

	; The referenced pointer to file object that represents
	; the corresponding device object.  This pointer we get
	; from IoGetDeviceObjectPointer and must dereference
	; while detaching.

	pTargetFileObject	PFILE_OBJECT	?

FiDO_DEVICE_EXTENSION ENDS
PFiDO_DEVICE_EXTENSION typedef ptr FiDO_DEVICE_EXTENSION

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                 R E A D O N L Y    D A T A                                        
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

.const

CCOUNTED_UNICODE_STRING	"\\Device\\KbdSpy", g_usControlDeviceName, 4
CCOUNTED_UNICODE_STRING	"\\DosDevices\\KbdSpy", g_usSymbolicLinkName, 4

CCOUNTED_UNICODE_STRING	"\\Device\\KeyboardClass0", g_usTargetDeviceName, 4

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                              U N I N I T I A L I Z E D  D A T A                                   
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

.data?

g_pDriverObject					PDRIVER_OBJECT	?

g_pControlDeviceObject			PDEVICE_OBJECT	?	; Control Device Object pointer
g_pFilterDeviceObject			PDEVICE_OBJECT	?	; Filter Device Object pointer

g_pEventObject					PKEVENT			?
; This spin-lock let us be sure that no one will dereference event object pointer
; while we compare it agaist NULL and then call KeSetEvent in our completion routine
g_EventSpinLock					KSPIN_LOCK		?	; locks key data list

g_fCDO_Opened					BOOL			?
g_fFiDO_Attached				BOOL			?
g_fSpy							BOOL			?

g_dwPendingRequests				DWORD			?

align 4
g_pKeyDataLookaside				PNPAGED_LOOKASIDE_LIST	?

align 4
g_KeyDataListHead				LIST_ENTRY		<>	; accessed under lock

; Holds number of KEY_DATA_ENTRYs in list. Should not exceed MAX_KEY_DATA_ENTRIES.
g_cKeyDataEntries				SDWORD			?	; accessed under lock

; This spin-lock let us be sure that only one thread is working with key data at a time
g_KeyDataSpinLock				KSPIN_LOCK		?

; This mutex let us be sure no one will try to do some unpredictable things.
; For example: no one can try to attach while we in the middle of the detaching.
align 4
g_mtxCDO_State					KMUTEX		<>

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                         C O D E                                                   
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

.code

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                              _ExAllocateFromNPagedLookasideList                                   
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

_ExAllocateFromNPagedLookasideList proc uses esi Lookaside:PNPAGED_LOOKASIDE_LIST

;; Function Description:
;;    This function removes (pops) the first entry from the specified
;;    nonpaged lookaside list.
;;
;; Arguments:
;;    Lookaside - Supplies a pointer to a nonpaged lookaside list structure.
;;
;; Return Value:
;;    If an entry is removed from the specified lookaside list, then the
;;    address of the entry is returned as the function value. Otherwise,
;;    NULL is returned.

	mov esi, Lookaside
	assume esi:ptr NPAGED_LOOKASIDE_LIST
	mov ecx, esi
	inc [esi].L.TotalAllocates
	lea edx, [esi]._Lock
	fastcall ExInterlockedPopEntrySList, ecx, edx
	.if eax == NULL
		push [esi].L.Tag
		inc [esi].L.AllocateMisses
		push [esi].L._Size
		push [esi].L._Type
		call [esi].L.Allocate
	.endif
	assume esi:nothing

	ret

_ExAllocateFromNPagedLookasideList endp

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                  _ExFreeToNPagedLookasideList                                     
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

_ExFreeToNPagedLookasideList proc Lookaside:PNPAGED_LOOKASIDE_LIST, Entry:PVOID

;; Function Description:
;;    This function inserts (pushes) the specified entry into the specified
;;    nonpaged lookaside list.
;;
;; Arguments:
;;    Lookaside - Supplies a pointer to a nonpaged lookaside list structure.
;;    Entry - Supples a pointer to the entry that is inserted in the lookaside list.
;;
;; Return Value:
;;    None.

	mov ecx, Lookaside
	assume ecx:ptr NPAGED_LOOKASIDE_LIST
	inc [ecx].L.TotalFrees
	mov ax, [ecx].L.ListHead.Depth
	.if ax >= [ecx].L.Depth
		push Entry
		inc [ecx].L.FreeMisses
		call [ecx].L.Free
	.else
		mov edx, Entry
		lea eax, [ecx]._Lock
		fastcall ExInterlockedPushEntrySList, ecx, edx, eax
	.endif
	assume ecx:nothing

	ret

_ExFreeToNPagedLookasideList endp
		
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                         AddEntry                                                  
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

AddEntry proc uses ebx pKeyData:PKEY_DATA

; We have to access g_cKeyDataEntries and g_KeyDataListHead
; under lock protection. Since keyboard hits occur relatively rare
; we simply protect whole code.  So, this proc may be optimized a little.

	LOCK_ACQUIRE g_KeyDataSpinLock
	mov bl, al			; old IRQL

	.if g_cKeyDataEntries < MAX_KEY_DATA_ENTRIES

		; Allocate new entry from lookaside list
	
		invoke _ExAllocateFromNPagedLookasideList, g_pKeyDataLookaside
		.if eax != NULL

			mov edx, eax
			assume edx:ptr KEY_DATA_ENTRY

			mov ecx, pKeyData
			assume ecx:ptr KEY_DATA

			mov eax, [ecx].dwScanCode
			mov [edx].KeyData.dwScanCode, eax

			mov eax, [ecx].Flags
			mov [edx].KeyData.Flags, eax

			assume ecx:nothing

			; Add to head

			lea ecx, [edx].ListEntry
			InsertHeadList addr g_KeyDataListHead, ecx

			assume edx:nothing

			inc g_cKeyDataEntries
    
		.endif
	.endif

	LOCK_RELEASE g_KeyDataSpinLock, bl

	ret

AddEntry endp

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                       RemoveEntry                                                 
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

RemoveEntry proc uses ebx pBuffer:PVOID

local cbCopied:DWORD

	and cbCopied, 0
	
	; We have to access g_cKeyDataEntries and g_KeyDataListHead
	; under lock protection. Since keyboard hits occur relatively rare
	; we simply protect whole code.  So, this proc may be optimized a little.
	
	LOCK_ACQUIRE g_KeyDataSpinLock
	mov bl, al							; old IRQL

	IsListEmpty addr g_KeyDataListHead
	.if eax != TRUE							; Is there something to remove?
			
		; Remove from tail

		RemoveTailList addr g_KeyDataListHead

		mov edx, eax						; edx -> KEY_DATA_ENTRY.ListEntry
		sub edx, KEY_DATA_ENTRY.ListEntry	; edx -> KEY_DATA_ENTRY

		assume edx:ptr KEY_DATA_ENTRY

		mov ecx, pBuffer
		assume ecx:ptr KEY_DATA

		mov eax, [edx].KeyData.dwScanCode
		mov [ecx].dwScanCode, eax

		mov eax, [edx].KeyData.Flags
		mov [ecx].Flags, eax

		mov cbCopied, sizeof KEY_DATA

		assume ecx:nothing
		assume edx:nothing

		; Put a block back onto lookaside list

		invoke _ExFreeToNPagedLookasideList, g_pKeyDataLookaside, edx

		dec g_cKeyDataEntries

	.endif

	LOCK_RELEASE g_KeyDataSpinLock, bl

	mov eax, cbCopied
	ret

RemoveEntry endp

;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
;                                      KeyboardAttach                                               
;:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

KeyboardAttach proc

local status:NTSTATUS
local pTargetDeviceObject:PDEVICE_OBJECT
local pTargetFileObject:PFILE_OBJECT

	invoke DbgPrint, $CTA0("KbdSpy: Entering KeyboardAttach\n")

	mov status, STATUS_UNSUCCESSFUL

	.if ( g_pFilterDeviceObject != NULL )

		; Filter device object exist and should be attached

		mov status, STATUS_SUCCESS

	.else

		; Let's attach to keyboard device stack
		;
		; Create unnamed device because filter device objects should never be named.
		; We are going to attach it to existing keyboard device stack. So no one may
		; directly open filter device by name.
		
		mov eax, g_pControlDeviceObject
		mov ecx, (DEVICE_OBJECT PTR [eax]).DriverObject

		invoke IoCreateDevice, ecx, sizeof FiDO_DEVICE_EXTENSION, NULL, \
					FILE_DEVICE_UNKNOWN, 0, FALSE, addr g_pFilterDeviceObject
		.if eax == STATUS_SUCCESS
	
			; Supply a name for any device object in the stack we are about to attach to.
			; IoGetDeviceObjectPointer returns the pointer to upper most device object in the stack.

			invoke IoGetDeviceObjectPointer, addr g_usTargetDeviceName, FILE_READ_DATA, \
										addr pTargetFileObject, addr pTargetDeviceObject
			.if eax == STATUS_SUCCESS
    
				; Here we have two pointers: pointer to the topmost device in the keyboard stack
				; and pointer to the corresponding file object.  IoGetDeviceObjectPointer
				; references file object but not the device object.
				;
				; We are just one line from attaching to our target.  We must prevent
				; our driver from unloading while it intercepts keyboard IRPs.
				; We could use RemoveLock, but the easiest solution is to remove pointer