#include <ntddk.h>
#include <common.h>

WCHAR DeviceName[] = L"\\Device\\antimida";
WCHAR SymLinkName[] = L"\\DosDevices\\antimida";

UNICODE_STRING usDeviceName;
UNICODE_STRING usSymbolicLinkName;

typedef struct _DEVICE_CONTEXT
{
	PDRIVER_OBJECT  pDriverObject;       
    PDEVICE_OBJECT  pDeviceObject;
}
DEVICE_CONTEXT, *PDEVICE_CONTEXT, **PPDEVICE_CONTEXT;

PDEVICE_OBJECT  g_pDeviceObject  = NULL;
PDEVICE_CONTEXT g_pDeviceContext = NULL;

#define FILE_DEVICE_ANTIMIDA 0x8000

#define CODE_RESTORE_INFO	CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x800, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_READ_MEM		CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x801, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_QUERY_MEM		CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x802, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_WRITE_MEM		CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x803, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_ALLOC_MEM		CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x804, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_CREATE_THREAD	CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x805, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)
#define CODE_DBG_CONTINUE	CTL_CODE(FILE_DEVICE_ANTIMIDA, 0x806, \
	METHOD_BUFFERED, FILE_ANY_ACCESS)


NTSTATUS DriverInitialize(PDRIVER_OBJECT  pDriverObject,
						  PUNICODE_STRING pusRegistryPath);

NTSTATUS DriverEntry(PDRIVER_OBJECT  pDriverObject,
					 PUNICODE_STRING pusRegistryPath);

#ifdef ALLOC_PRAGMA

#pragma alloc_text (INIT, DriverInitialize)
#pragma alloc_text (INIT, DriverEntry)

#endif

VOID NTAPI KeAttachProcess(IN PEPROCESS);

//
// ZwReadVirtualMemory stuff
//

typedef struct _Input_ZwReadVirtualMemory
{
	HANDLE  ProcessHandle;
	PVOID  BaseAddress;
	PVOID  Buffer;
	ULONG  BufferLength;
	PULONG  ReturnLength;
} Input_ZwReadVirtualMemory;

NTSTATUS (*pZwReadVirtualMemory)(IN HANDLE  ProcessHandle,
								 IN PVOID  BaseAddress,
								 OUT PVOID  Buffer,
								 IN ULONG  BufferLength,
								 OUT PULONG  ReturnLength  OPTIONAL);


KMUTEX DumpMutex;

//
// ZwQueryVirtualMemory stuff
//

typedef enum _MEMORY_INFORMATION_CLASS {
	MemoryBasicInformation,
	MemoryWorkingSetList,
	MemorySectionName,
	MemoryBasicVlmInformation
} MEMORY_INFORMATION_CLASS;


NTSTATUS (*pZwQueryVirtualMemory)(IN HANDLE  ProcessHandle,
								  IN PVOID  BaseAddress,
								  IN MEMORY_INFORMATION_CLASS  MemoryInformationClass,
								  OUT PVOID  MemoryInformation,
								  IN ULONG  MemoryInformationLength,
								  OUT PULONG  ReturnLength  OPTIONAL);

typedef struct _Input_ZwQueryVirtualMemory
{
	HANDLE ProcessHandle;
	PVOID  BaseAddress;
	MEMORY_INFORMATION_CLASS  MemoryInformationClass;
	ULONG  MemoryInformationLength;
	PVOID  MemoryInformation;
	PULONG  ReturnLength;
} Input_ZwQueryVirtualMemory;

//
// ZwWriteVirtualMemory stuff
//

NTSTATUS (*pZwWriteVirtualMemory)(IN HANDLE  ProcessHandle,
								 IN PVOID  BaseAddress,
								 IN PVOID  Buffer,
								 IN ULONG  BufferLength,
								 OUT PULONG  ReturnLength  OPTIONAL);

typedef struct _Input_ZwWriteVirtualMemory
{
	HANDLE  ProcessHandle;
	PVOID  BaseAddress;
	PVOID  Buffer;
	ULONG  BufferLength;
	PULONG  ReturnLength;
} Input_ZwWriteVirtualMemory;

//
// ZwAllocateVirtualMemory stuff
//

NTSTATUS (*pZwAllocateVirtualMemory)(IN HANDLE ProcessHandle,
									 IN OUT PVOID *BaseAddress,
									 IN ULONG_PTR ZeroBits,
									 IN OUT PSIZE_T RegionSize,
									 IN ULONG AllocationType,
									 IN ULONG Protect);

typedef struct _Input_ZwAllocateVirtualMemory
{
	HANDLE ProcessHandle;
	PVOID *BaseAddress;
	ULONG_PTR ZeroBits;
	PSIZE_T RegionSize;
	ULONG AllocationType;
	ULONG Protect;
} Input_ZwAllocateVirtualMemory;

//
// ZwCreateThread stuff
//

typedef struct _USER_STACK {
	PVOID  FixedStackBase;
	PVOID  FixedStackLimit;
	PVOID  ExpandableStackBase;
	PVOID  ExpandableStackLimit;
	PVOID  ExpandableStackBottom;
} USER_STACK, *PUSER_STACK;


NTSTATUS (*pZwCreateThread)(OUT PHANDLE  ThreadHandle,
							IN ACCESS_MASK  DesiredAccess,
							IN POBJECT_ATTRIBUTES  ObjectAttributes,
							IN HANDLE  ProcessHandle,
							OUT PCLIENT_ID  ClientId,
							IN PCONTEXT  ThreadContext,
							IN PUSER_STACK  UserStack,
							IN BOOLEAN  CreateSuspended);

typedef struct _Input_ZwCreateThread
{
	PHANDLE  ThreadHandle;
	ACCESS_MASK  DesiredAccess;
	POBJECT_ATTRIBUTES  ObjectAttributes;
	HANDLE  ProcessHandle;
	PCLIENT_ID  ClientId;
	PCONTEXT  ThreadContext;
	PUSER_STACK  UserStack;
	BOOLEAN  CreateSuspended;
} Input_ZwCreateThread;

//
// ZwDebugContinue stuff
// i was too lazy to find the correct declaration
//

NTSTATUS (*pZwDebugContinue)(PVOID *A, PVOID *B, PVOID *C);

typedef struct _Input_ZwDebugContinue
{
	PVOID *A;
	PVOID *B;
	PVOID *C;
} Input_ZwDebugContinue;


//
// restore info struct
//

typedef struct _RESTORE_INFO
{
	DWORD ZwAllocateVirtualMemory;
	DWORD ZwCreateThread;
	DWORD ZwDebugContinue;
	DWORD ZwQueryVirtualMemory;
	DWORD ZwReadVirtualMemory;
	DWORD ZwTerminateProcess;
	DWORD ZwWriteVirtualMemory;

	BYTE KeAttachProcessPatch[5];

} RESTORE_INFO, *PRESTORE_INFO;

RESTORE_INFO RestoreInfo;

//
// SDT hook stuff
// 

extern PSERVICE_DESCRIPTOR_TABLE KeServiceDescriptorTable;

PVOID *KeServiceTablePointers;
PMDL KeServiceTableMdl;
BOOLEAN *ServiceIsHooked;

PPVOID MapServiceTable(BOOLEAN **);
VOID UnmapServiceTable(PVOID);

#ifdef ALPHA
#define FUNCTION_PTR(_Function) (*(PULONG) _Function) & 0x0000FFFF
#else
#define FUNCTION_PTR(_Function) *(PULONG)((PUCHAR) _Function + 1)
#endif

#define HOOK_SYSCALL(_Function, _Hook, _Orig)									\
	if (!ServiceIsHooked[FUNCTION_PTR(_Function)]) {							\
	_Orig = (PVOID) InterlockedExchange((PLONG)							\
	&KeServiceTablePointers[FUNCTION_PTR(_Function)], (LONG) _Hook );	\
	ServiceIsHooked[ FUNCTION_PTR(_Function) ] = TRUE; }


#define UNHOOK_SYSCALL(_Function, _Hook, _Orig)								\
	if (ServiceIsHooked[FUNCTION_PTR(_Function)] &&							\
	KeServiceTablePointers[FUNCTION_PTR(_Function) ] == (PVOID) _Hook ) {	\
	InterlockedExchange((PLONG) &KeServiceTablePointers[			\
	FUNCTION_PTR(_Function)], (LONG) _Orig );						\
	ServiceIsHooked[FUNCTION_PTR(_Function)] = FALSE; }

NTSTATUS (*RealZwCreateFile)(OUT PHANDLE  FileHandle,
							 IN ACCESS_MASK  DesiredAccess,
							 IN POBJECT_ATTRIBUTES  ObjectAttributes,
							 OUT PIO_STATUS_BLOCK  IoStatusBlock,
							 IN PLARGE_INTEGER  AllocationSize  OPTIONAL,
							 IN ULONG  FileAttributes,
							 IN ULONG  ShareAccess,
							 IN ULONG  CreateDisposition,
							 IN ULONG  CreateOptions,
							 IN PVOID  EaBuffer  OPTIONAL,
							 IN ULONG  EaLength
							 );

NTSTATUS (*RealZwOpenKey)(OUT PHANDLE, IN ACCESS_MASK, IN POBJECT_ATTRIBUTES);

NTSTATUS HookedZwCreateFile(OUT PHANDLE  FileHandle,
							IN ACCESS_MASK  DesiredAccess,
							IN POBJECT_ATTRIBUTES  ObjectAttributes,
							OUT PIO_STATUS_BLOCK  IoStatusBlock,
							IN PLARGE_INTEGER  AllocationSize  OPTIONAL,
							IN ULONG  FileAttributes,
							IN ULONG  ShareAccess,
							IN ULONG  CreateDisposition,
							IN ULONG  CreateOptions,
							IN PVOID  EaBuffer  OPTIONAL,
							IN ULONG  EaLength)
{
	DbgPrint("%ws\n", ObjectAttributes->ObjectName->Buffer);

	return RealZwCreateFile(FileHandle, DesiredAccess, ObjectAttributes,
		IoStatusBlock, AllocationSize, FileAttributes, ShareAccess,
		CreateDisposition, CreateOptions, EaBuffer, EaLength);
}

NTSTATUS HookedZwOpenKey(OUT PHANDLE KeyHandle, IN ACCESS_MASK DesiredAccess,
						 IN POBJECT_ATTRIBUTES ObjectAttributes)
{
	return RealZwOpenKey(KeyHandle, DesiredAccess, ObjectAttributes);
}

PPVOID MapServiceTable(BOOLEAN **ServiceIsHooked)
{
	PVOID Mem;

	Mem = ExAllocatePoolWithTag(0, 
		KeServiceDescriptorTable->ntoskrnl.ServiceLimit,
		0x206B6444);

	if (Mem == NULL)
		return NULL;

	*ServiceIsHooked = (BOOLEAN *) Mem;

	memset(Mem, 0, KeServiceDescriptorTable->ntoskrnl.
		ServiceLimit);

	KeServiceTableMdl = MmCreateMdl(NULL, 
		KeServiceDescriptorTable->ntoskrnl.ServiceTable,
		(KeServiceDescriptorTable->ntoskrnl.ServiceLimit * 
		sizeof (POINTER)));

	if (KeServiceTableMdl == NULL)
		return NULL;

	MmBuildMdlForNonPagedPool(KeServiceTableMdl);

	return (PPVOID) MmMapLockedPages(KeServiceTableMdl, 0);
}

VOID UnmapServiceTable(PVOID KeServiceTablePointers)
{
	if (KeServiceTableMdl == NULL)
		return;

	MmUnmapLockedPages(KeServiceTablePointers,
		KeServiceTableMdl);

	ExFreePool(KeServiceTableMdl);
}

//
// patch KeAttachProcess
//

VOID RebuildNtoskrnl()
{
	PMDL Mdl;
	DWORD CR0Backup;
	DWORD *ptr;
	BYTE *Patch;

	ptr = (DWORD *)(2 + (DWORD) &KeAttachProcess);

	ptr = (DWORD*) *ptr;

	Patch = (BYTE *) *ptr;

	Mdl = MmCreateMdl(0, (PVOID) Patch, 5);

	if (Mdl == NULL)
		return;

	MmProbeAndLockPages(Mdl, 0, 0);

	if (*Patch == 0xE9)
	{
		__asm
		{
			mov eax, cr0
			mov CR0Backup, eax
			and eax, 0xFFFEFFFF
			mov cr0, eax
		}

		//
		// patch
		//

		memcpy(Patch, RestoreInfo.KeAttachProcessPatch, 5);

		__asm
		{
			mov eax, CR0Backup
			mov cr0, eax
		}
	}

	MmUnlockPages(Mdl);
	IoFreeMdl(Mdl);
}


NTSTATUS ControlDispatcher(PDEVICE_CONTEXT pDeviceContext, DWORD dwCode,
						   BYTE *pInput, DWORD dwInputSize,
						   BYTE *pOutput, DWORD dwOutputSize, DWORD *pdwInfo)
{
	switch (dwCode)
	{

	case CODE_RESTORE_INFO:
		{
			RtlCopyMemory(&RestoreInfo, pInput, sizeof (RESTORE_INFO));

			pZwReadVirtualMemory = (NTSTATUS (*)(HANDLE,
				PVOID, PVOID, ULONG, PULONG))
				RestoreInfo.ZwReadVirtualMemory;

			pZwQueryVirtualMemory = (NTSTATUS (*)(HANDLE,
				PVOID, MEMORY_INFORMATION_CLASS, PVOID, ULONG,
				PULONG)) RestoreInfo.ZwQueryVirtualMemory;

			pZwWriteVirtualMemory = (NTSTATUS (*)(HANDLE, PVOID,
				PVOID, ULONG, PULONG))
				RestoreInfo.ZwWriteVirtualMemory;

			pZwAllocateVirtualMemory = (NTSTATUS (*)(HANDLE,
				PVOID *, ULONG_PTR, PSIZE_T, ULONG, ULONG))
				RestoreInfo.ZwAllocateVirtualMemory;

			pZwCreateThread = (NTSTATUS (*)(PHANDLE, ACCESS_MASK,
				POBJECT_ATTRIBUTES, HANDLE, PCLIENT_ID,
				PCONTEXT, PUSER_STACK, BOOLEAN))
				RestoreInfo.ZwCreateThread;

			pZwDebugContinue = (NTSTATUS (*)(PVOID *A, PVOID *B, PVOID *C))
				RestoreInfo.ZwDebugContinue;

			RebuildNtoskrnl();

			break;
		}

	case CODE_READ_MEM:
		{
			Input_ZwReadVirtualMemory Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwReadVirtualMemory));

			return pZwReadVirtualMemory(Input.ProcessHandle, Input.BaseAddress,
				Input.Buffer, Input.BufferLength, Input.ReturnLength);
		}

	case CODE_QUERY_MEM:
		{
			Input_ZwQueryVirtualMemory Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwQueryVirtualMemory));

			return pZwQueryVirtualMemory(Input.ProcessHandle, Input.BaseAddress,
				Input.MemoryInformationClass, Input.MemoryInformation,
				Input.MemoryInformationLength, Input.ReturnLength);
		}

	case CODE_WRITE_MEM:
		{
			Input_ZwWriteVirtualMemory Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwWriteVirtualMemory));

			return pZwWriteVirtualMemory(Input.ProcessHandle, Input.BaseAddress,
				Input.Buffer, Input.BufferLength, Input.ReturnLength);
		}

	case CODE_ALLOC_MEM:
		{
			Input_ZwAllocateVirtualMemory Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwAllocateVirtualMemory));

			return pZwAllocateVirtualMemory(Input.ProcessHandle, Input.BaseAddress,
				Input.ZeroBits, Input.RegionSize, Input.AllocationType, Input.Protect);
		}

	case CODE_CREATE_THREAD:
		{
			Input_ZwCreateThread Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwCreateThread));

			return pZwCreateThread(Input.ThreadHandle, Input.DesiredAccess,
				Input.ObjectAttributes, Input.ProcessHandle, Input.ClientId, 
				Input.ThreadContext, Input.UserStack, Input.CreateSuspended);
		}

	case CODE_DBG_CONTINUE:
		{
			Input_ZwDebugContinue Input;

			RtlCopyMemory(&Input, pInput, sizeof (Input_ZwDebugContinue));

			return pZwDebugContinue(Input.A, Input.B, Input.C);
		}

		
	default:
	 	return STATUS_INVALID_PARAMETER;
	}
	 
	return STATUS_SUCCESS;
}

NTSTATUS DeviceDispatcher(PDEVICE_CONTEXT pDeviceContext, PIRP pIrp)
{
    PIO_STACK_LOCATION pisl;
    DWORD dwInfo = 0;
    NTSTATUS ns = STATUS_NOT_IMPLEMENTED;

    pisl = IoGetCurrentIrpStackLocation(pIrp);

    switch (pisl->MajorFunction)
	{
	
	case IRP_MJ_CREATE:
	case IRP_MJ_CLEANUP:
	case IRP_MJ_CLOSE:
		{
			ns = STATUS_SUCCESS;
            break;
		}

	case IRP_MJ_DEVICE_CONTROL:
		{
			MUTEX_ACQUIRE(DumpMutex);

			ns = ControlDispatcher(pDeviceContext,
				pisl->Parameters.DeviceIoControl.IoControlCode,
				(BYTE *) pIrp->AssociatedIrp.SystemBuffer,
				pisl->Parameters.DeviceIoControl.InputBufferLength,
				(BYTE *) pIrp->AssociatedIrp.SystemBuffer,
				pisl->Parameters.DeviceIoControl.OutputBufferLength,
				&dwInfo);

			MUTEX_RELEASE(DumpMutex);
			
			break;
		}
	}

    pIrp->IoStatus.Status = ns;
    pIrp->IoStatus.Information = dwInfo;

    IoCompleteRequest (pIrp, IO_NO_INCREMENT);

    return ns;
}

NTSTATUS DriverDispatcher(PDEVICE_OBJECT pDeviceObject, PIRP pIrp)
{
	return (pDeviceObject == g_pDeviceObject ? 
		DeviceDispatcher(g_pDeviceContext, pIrp)
		: STATUS_INVALID_PARAMETER_1);
}


VOID DriverUnload(PDRIVER_OBJECT pDriverObject)
{
	UNHOOK_SYSCALL(ZwCreateFile, HookedZwCreateFile, RealZwCreateFile);
	UNHOOK_SYSCALL(ZwOpenKey, HookedZwOpenKey, RealZwOpenKey);
	UnmapServiceTable(KeServiceTablePointers);

    IoDeleteSymbolicLink(&usSymbolicLinkName);
    IoDeleteDevice(pDriverObject->DeviceObject);
}

NTSTATUS DriverInitialize(PDRIVER_OBJECT pDriverObject,
						  PUNICODE_STRING pusRegistryPath)
{
    PDEVICE_OBJECT pDeviceObject = NULL;
    NTSTATUS ns = STATUS_DEVICE_CONFIGURATION_ERROR;
    
    RtlInitUnicodeString(&usDeviceName, DeviceName);
    RtlInitUnicodeString(&usSymbolicLinkName, SymLinkName);

    if ((ns = IoCreateDevice(pDriverObject, sizeof (DEVICE_CONTEXT),
		&usDeviceName, FILE_DEVICE_ANTIMIDA, 0, FALSE,
		&pDeviceObject)) == STATUS_SUCCESS)
	{
        if ((ns = IoCreateSymbolicLink(&usSymbolicLinkName, 
			&usDeviceName)) == STATUS_SUCCESS)
		{
            g_pDeviceObject  = pDeviceObject;
            g_pDeviceContext = pDeviceObject->DeviceExtension;

            g_pDeviceContext->pDriverObject = pDriverObject;
            g_pDeviceContext->pDeviceObject = pDeviceObject;

		}
		else
		{
			IoDeleteDevice(pDeviceObject);
		}
	}
	
	return ns;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObject,
					 PUNICODE_STRING pusRegistryPath)
{
	PDRIVER_DISPATCH *ppdd;
    NTSTATUS ns = STATUS_DEVICE_CONFIGURATION_ERROR;

    if ((ns = DriverInitialize(pDriverObject, pusRegistryPath)) == STATUS_SUCCESS)
	{
        ppdd = pDriverObject->MajorFunction;

        ppdd[IRP_MJ_CREATE                  ] =
        ppdd[IRP_MJ_CREATE_NAMED_PIPE       ] =
        ppdd[IRP_MJ_CLOSE                   ] =
        ppdd[IRP_MJ_READ                    ] =
        ppdd[IRP_MJ_WRITE                   ] =
        ppdd[IRP_MJ_QUERY_INFORMATION       ] =
        ppdd[IRP_MJ_SET_INFORMATION         ] =
        ppdd[IRP_MJ_QUERY_EA                ] =
        ppdd[IRP_MJ_SET_EA                  ] =
        ppdd[IRP_MJ_FLUSH_BUFFERS           ] =
        ppdd[IRP_MJ_QUERY_VOLUME_INFORMATION] =
        ppdd[IRP_MJ_SET_VOLUME_INFORMATION  ] =
        ppdd[IRP_MJ_DIRECTORY_CONTROL       ] =
        ppdd[IRP_MJ_FILE_SYSTEM_CONTROL     ] =
        ppdd[IRP_MJ_DEVICE_CONTROL          ] =
        ppdd[IRP_MJ_INTERNAL_DEVICE_CONTROL ] =
        ppdd[IRP_MJ_SHUTDOWN                ] =
        ppdd[IRP_MJ_LOCK_CONTROL            ] =
        ppdd[IRP_MJ_CLEANUP                 ] =
        ppdd[IRP_MJ_CREATE_MAILSLOT         ] =
        ppdd[IRP_MJ_QUERY_SECURITY          ] =
        ppdd[IRP_MJ_SET_SECURITY            ] =
        ppdd[IRP_MJ_POWER                   ] =
        ppdd[IRP_MJ_SYSTEM_CONTROL          ] =
        ppdd[IRP_MJ_DEVICE_CHANGE           ] =
        ppdd[IRP_MJ_QUERY_QUOTA             ] =
        ppdd[IRP_MJ_SET_QUOTA               ] =
        ppdd[IRP_MJ_PNP                     ] = DriverDispatcher;
        pDriverObject->DriverUnload           = DriverUnload;

		MUTEX_INIT(DumpMutex);

		KeServiceTablePointers = MapServiceTable(&ServiceIsHooked);

		HOOK_SYSCALL(ZwCreateFile, HookedZwCreateFile, RealZwCreateFile);
		HOOK_SYSCALL(ZwOpenKey, HookedZwOpenKey, RealZwOpenKey);

	}

	return ns;
}
