keprocessm_drv.cpp

驱动枚举进程,控制线程AFFINITY,通过修改EPROCESS,ETHREAD ,KTHREAD 等结构,修改AFFINITY

/************************************************************************
* 文件名称:KeProcessM_Drv.cpp                                              
* 作    者:李骥
*************************************************************************/
#ifdef __cplusplus
extern "C"
{
#endif

#include <NTDDK.h>

#include "InterruptHook/hook.h"
#include "InterruptHook/handler.h"

#ifdef __cplusplus
}
#endif 

#include "KeProcessM_Drv.h"


#include "LEE_Functions.h"


#include "ioctls.h"


/************************************************************************
* 函数名称:DriverEntry
* 功能描述:初始化驱动程序，定位和申请硬件资源，创建内核对象
* 参数列表:
pDriverObject:从I/O管理器中传进来的驱动对象
pRegistryPath:驱动程序在注册表的中的路径
* 返回 值:返回初始化驱动状态
*************************************************************************/
	extern HANDLE hMyThread;
	typedef NTSTATUS (*TTD)(LPARAM lParam);
#pragma INITCODE
extern "C" NTSTATUS DriverEntry (
								 IN PDRIVER_OBJECT pDriverObject,
								 IN PUNICODE_STRING pRegistryPath	) 
{
	
	NTSTATUS status;
	KdPrint(("Enter DriverEntry\n"));

	//设置卸载函数
	pDriverObject->DriverUnload = Unload;

	//设置派遣函数
	for (int i = 0; i < arraysize(pDriverObject->MajorFunction); ++i)
		pDriverObject->MajorFunction[i] = DispatchRoutin;

	pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceIOControl;

	//创建驱动设备对象
	status = CreateDevice(pDriverObject);


	// 初始化事件
//	KeInitializeEvent(&kEvent, SynchronizationEvent, FALSE);



	PsCreateSystemThread(&hMyThread,0,NULL,NtCurrentProcess(),NULL,SystemThread,NULL);
	

	BackupNtVectors();
	SetExternalHandlers();
	

	KdPrint(("Leave DriverEntry\n"));
	return status;
}

/************************************************************************
* 函数名称:CreateDevice
* 功能描述:初始化设备对象
* 参数列表:
pDriverObject:从I/O管理器中传进来的驱动对象
* 返回 值:返回初始化状态
*************************************************************************/
#pragma INITCODE
NTSTATUS CreateDevice (
					   IN PDRIVER_OBJECT	pDriverObject) 
{
	NTSTATUS status;
	PDEVICE_OBJECT pDevObj;
	PDEVICE_EXTENSION pDevExt;

	//创建设备名称
	UNICODE_STRING devName;
	RtlInitUnicodeString(&devName,L"\\Device\\MyProcessManage");

	//创建设备
	status = IoCreateDevice( pDriverObject,
		sizeof(DEVICE_EXTENSION),
		&(UNICODE_STRING)devName,
		FILE_DEVICE_UNKNOWN,
		0, TRUE,
		&pDevObj );
	if (!NT_SUCCESS(status))
		return status;

	pDevObj->Flags |= DO_BUFFERED_IO;
	pDevExt = (PDEVICE_EXTENSION)pDevObj->DeviceExtension;
	pDevExt->pDevice = pDevObj;
	pDevExt->ustrDeviceName = devName;

	//申请模拟文件的缓冲区
	pDevExt->buffer = (PUCHAR)ExAllocatePool(PagedPool,MAX_FILE_LENGTH);
	//设置模拟文件大小
	pDevExt->file_length = 0;

	//创建符号链接
	UNICODE_STRING symLinkName;
	RtlInitUnicodeString(&symLinkName,L"\\??\\MyProcessManage");
	pDevExt->ustrSymLinkName = symLinkName;
	status = IoCreateSymbolicLink( &symLinkName,&devName );
	if (!NT_SUCCESS(status)) 
	{
		IoDeleteDevice( pDevObj );
		return status;
	}
	return STATUS_SUCCESS;
}

/************************************************************************
* 函数名称:HelloDDKUnload
* 功能描述:负责驱动程序的卸载操作
* 参数列表:
pDriverObject:驱动对象
* 返回 值:返回状态
*************************************************************************/
#pragma PAGEDCODE
VOID Unload (IN PDRIVER_OBJECT pDriverObject) 
{
	PDEVICE_OBJECT	pNextObj;
	KdPrint(("Enter DriverUnload\n"));
	pNextObj = pDriverObject->DeviceObject;
	while (pNextObj != NULL) 
	{
		PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)
			pNextObj->DeviceExtension;
		if (pDevExt->buffer)
		{
			ExFreePool(pDevExt->buffer);
			pDevExt->buffer = NULL;
		}

		//删除符号链接
		UNICODE_STRING pLinkName = pDevExt->ustrSymLinkName;
		IoDeleteSymbolicLink(&pLinkName);
		pNextObj = pNextObj->NextDevice;
		IoDeleteDevice( pDevExt->pDevice );
	}
	RestoreNtVectors();

}

/************************************************************************
* 函数名称:DispatchRoutin
* 功能描述:对读IRP进行处理
* 参数列表:
pDevObj:功能设备对象
pIrp:从IO请求包
* 返回 值:返回状态
*************************************************************************/
#pragma PAGEDCODE
NTSTATUS DispatchRoutin(IN PDEVICE_OBJECT pDevObj,
								IN PIRP pIrp) 
{
	KdPrint(("Enter DispatchRoutin\n"));

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	//建立一个字符串数组与IRP类型对应起来
	static char* irpname[] = 
	{
		"IRP_MJ_CREATE",
		"IRP_MJ_CREATE_NAMED_PIPE",
		"IRP_MJ_CLOSE",
		"IRP_MJ_READ",
		"IRP_MJ_WRITE",
		"IRP_MJ_QUERY_INFORMATION",
		"IRP_MJ_SET_INFORMATION",
		"IRP_MJ_QUERY_EA",
		"IRP_MJ_SET_EA",
		"IRP_MJ_FLUSH_BUFFERS",
		"IRP_MJ_QUERY_VOLUME_INFORMATION",
		"IRP_MJ_SET_VOLUME_INFORMATION",
		"IRP_MJ_DIRECTORY_CONTROL",
		"IRP_MJ_FILE_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CONTROL",
		"IRP_MJ_INTERNAL_DEVICE_CONTROL",
		"IRP_MJ_SHUTDOWN",
		"IRP_MJ_LOCK_CONTROL",
		"IRP_MJ_CLEANUP",
		"IRP_MJ_CREATE_MAILSLOT",
		"IRP_MJ_QUERY_SECURITY",
		"IRP_MJ_SET_SECURITY",
		"IRP_MJ_POWER",
		"IRP_MJ_SYSTEM_CONTROL",
		"IRP_MJ_DEVICE_CHANGE",
		"IRP_MJ_QUERY_QUOTA",
		"IRP_MJ_SET_QUOTA",
		"IRP_MJ_PNP",
	};

	UCHAR type = stack->MajorFunction;
	if (type >= arraysize(irpname))
		KdPrint((" - Unknown IRP, major type %X\n", type));
	else
		KdPrint(("\t%s\n", irpname[type]));

	NTSTATUS status = STATUS_SUCCESS;
	// 完成IRP
	pIrp->IoStatus.Status = status;
	pIrp->IoStatus.Information = 0;	// bytes xfered
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );

	KdPrint(("Leave DispatchRoutin\n"));

	return status;
}

VOID
GetDeviceObjectInfo( PDEVICE_OBJECT DevObj )
{
	POBJECT_HEADER ObjectHeader;
	POBJECT_HEADER_NAME_INFO ObjectNameInfo; 

	if ( DevObj == NULL )
	{
		DbgPrint( "DevObj is NULL!\n" );
		return;
	}
	// 得到对象头
	ObjectHeader = OBJECT_TO_OBJECT_HEADER( DevObj );

	if ( ObjectHeader )
	{
		// 查询设备名称并打印
		ObjectNameInfo = OBJECT_HEADER_TO_NAME_INFO( ObjectHeader );

		if ( ObjectNameInfo && ObjectNameInfo->Name.Buffer )
		{
			DbgPrint( "Driver Name:%wZ - Device Name:%wZ - Driver Address:0x%x - Device Address:0x%x\n",
				&DevObj->DriverObject->DriverName,
				&ObjectNameInfo->Name,
				DevObj->DriverObject,
				DevObj );
		}

		// 对于没有名称的设备，则打印 NULL
		else if ( DevObj->DriverObject )
		{
			DbgPrint( "Driver Name:%wZ - Device Name:%S - Driver Address:0x%x - Device Address:0x%x\n",
				&DevObj->DriverObject->DriverName,
				L"NULL",
				DevObj->DriverObject,
				DevObj );
		}
	}
}

VOID
GetAttachedDeviceInfo( PDEVICE_OBJECT DevObj )
{
	PDEVICE_OBJECT DeviceObject;

	if ( DevObj == NULL )
	{
		DbgPrint( "DevObj is NULL!\n" );
		return;
	}

	DeviceObject = DevObj->AttachedDevice;

	while ( DeviceObject )
	{
		DbgPrint( "Attached Driver Name:%wZ,Attached Driver Address:0x%x,Attached DeviceAddress:0x%x\n",
			&DeviceObject->DriverObject->DriverName,
			DeviceObject->DriverObject,
			DeviceObject );

		DeviceObject = DeviceObject->AttachedDevice;
	}
}

PDRIVER_OBJECT
EnumDeviceStack( PWSTR pwszDeviceName )
{
	UNICODE_STRING DriverName;
	PDRIVER_OBJECT DriverObject = NULL;
	PDEVICE_OBJECT DeviceObject = NULL;


	RtlInitUnicodeString( &DriverName, pwszDeviceName );


	ObReferenceObjectByName( &DriverName,
		OBJ_CASE_INSENSITIVE,
		NULL,
		0,
		( POBJECT_TYPE ) IoDriverObjectType,
		KernelMode,
		NULL,
		(PVOID*)&DriverObject );

	if ( DriverObject == NULL )
	{
		return NULL;
	}

	DeviceObject = DriverObject->DeviceObject;

	while ( DeviceObject )
	{
		GetDeviceObjectInfo( DeviceObject );

		// 判断当前设备上是否有过滤驱动（Filter Driver）
		if ( DeviceObject->AttachedDevice )
		{
			GetAttachedDeviceInfo( DeviceObject );
		}

		// 进一步判断当前设备上 VPB 中的设备
		if ( DeviceObject->Vpb && DeviceObject->Vpb->DeviceObject )
		{
			GetDeviceObjectInfo( DeviceObject->Vpb->DeviceObject );

			if ( DeviceObject->Vpb->DeviceObject->AttachedDevice )
			{
				GetAttachedDeviceInfo( DeviceObject->Vpb->DeviceObject );
			}
		}

		// 得到建立在此驱动上的下一个设备 DEVICE_OBJECT 
		DeviceObject = DeviceObject->NextDevice;
	}

	return DriverObject;
}
NTSTATUS DeviceIOControl( PDEVICE_OBJECT DeviceObject, PIRP irp )
{

	PIO_STACK_LOCATION irpStack;
	PVOID              inputBuffer;
	PVOID              outputBuffer;
	ULONG              inputBufferLength;
	ULONG              outputBufferLength;
	ULONG              ioControlCode;
	ULONG              ioFunctionCode;
	NTSTATUS           ntstatus;

	KdPrint(("Enter DeviceIOControl\n"));

	ntstatus = irp->IoStatus.Status = STATUS_SUCCESS;
	irp->IoStatus.Information = 0;

	irpStack = IoGetCurrentIrpStackLocation( irp );

	inputBuffer             = irpStack->Parameters.DeviceIoControl.Type3InputBuffer;
	inputBufferLength       = irpStack->Parameters.DeviceIoControl.InputBufferLength;
	outputBuffer            = irp->UserBuffer;
	outputBufferLength      = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
	ioControlCode           = irpStack->Parameters.DeviceIoControl.IoControlCode;


	ioFunctionCode=GetIO_FUNCTIONCODE(ioControlCode);
	if(ioFunctionCode>0x800&&ioFunctionCode<=0x809)
	{
	ntstatus = IoDeviceControlProcess( irpStack->FileObject, TRUE,
		inputBuffer, inputBufferLength, 
		outputBuffer, outputBufferLength,
		ioControlCode, &irp->IoStatus, DeviceObject );
	}
	else if(ioFunctionCode>0x810&&ioFunctionCode<=0x820)
	{
	ntstatus=IoDeviceControlInterrupt( DeviceObject, irp);

	}
	IoCompleteRequest( irp, IO_NO_INCREMENT );

	return ntstatus;   
}

NTSTATUS IoDeviceControlProcess( IN PFILE_OBJECT FileObject,IN BOOLEAN Wait, IN PVOID InputBuffer,IN ULONG InputBufferLength, OUT PVOID OutputBuffer,IN ULONG OutputBufferLength,IN ULONG IoControlCode,OUT PIO_STATUS_BLOCK IoStatus,IN PDEVICE_OBJECT DeviceObject )
{
	NTSTATUS status = STATUS_SUCCESS;
	KdPrint(("Enter IoCtlDispath\n"));

	//得到IOCTL码
	ULONG code = IoControlCode;

	ULONG info = 0;

	switch (code)
	{						// process request
	case IOCTL_DUMP_DEVICE_STACK:
		{
			KdPrint(("IOCTL_DUMP_DEVICE_STACK\n"));
			//缓冲区方式IOCTL
			//显示输入缓冲区数据
			WCHAR* InputBuffer = (WCHAR*)InputBufferLength;
			KdPrint(("%ws\n",InputBuffer));

			EnumDeviceStack( InputBuffer );

			break;
		}
	case IOCTL_GETPROCESS:


		MyGetProcessInformation((PMY_PROCESS_INFO)OutputBuffer,(PMY_THREAD_INFO)((CHAR *)OutputBuffer+0x8000),OutputBufferLength);
		DbgPrint("Ok Succece!");

		break;

	case IOCTL_SETPROCESS:
		if(InputBuffer==NULL)
		{
			DbgPrint("NULL InputBuffer!");
			break;
		}
		MySetProcessInformation((PMY_PROCESS_INFO)InputBuffer,(PMY_THREAD_INFO)((CHAR *)InputBuffer+0x8000),InputBufferLength);

//		SetAllThreadAffinity(0x01);
		break;
	case IOCTL_SETPROCESS_ALL_AFFINITY:
		if(InputBuffer==NULL)
		{
			DbgPrint("NULL InputBuffer!");
			break;
		}
	
		SetAllThreadAffinity(0x03);
		break;
	case IOCTL_GIVE_FUNCTION:
		{

			TTD ttd=*(TTD *)InputBuffer;
			KIRQL irql;
			KeRaiseIrql(HIGH_LEVEL,&irql);

			ttd(NULL);

			KeLowerIrql(irql);



		}

		break;
	default:
		status = STATUS_INVALID_VARIANT;
	}

	

	return status;
}


NTSTATUS IoDeviceControlInterrupt(IN	PDEVICE_OBJECT	pDeviceObject,  IN	PIRP	pIrp )
{
	NTSTATUS				iStatus = STATUS_SUCCESS;
	PIO_STACK_LOCATION		pStack;
	ULONG					iTransfered = 0;
	IDT						Idt;
	INT_VECTOR				Vec;

	pStack = IoGetCurrentIrpStackLocation( pIrp );

	switch( pStack->Parameters.DeviceIoControl.IoControlCode )
	{
	case IOCTL_HOOK_INT:
		if( pStack->Parameters.DeviceIoControl.InputBufferLength != sizeof(UCHAR) )
		{
			iStatus = STATUS_INVALID_PARAMETER;
			iTransfered = sizeof(UCHAR);
			break;
		}

		DPRINT( "IoDeviceControl: Hooking Int 0x%.2X\n", *(UCHAR *)pIrp->AssociatedIrp.SystemBuffer );

		LoadIDT( &Idt );
		LoadINTVector( 
			&Idt, 
			*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
			&Vec );
		DWORD_TO_VEC_OFFSET( Vec, InternalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer] );
		SaveINTVector( 
			&Idt, 
			*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
			&Vec );

		DPRINT( "IoDeviceControl: Hooking successful\n" );
		break;

	case IOCTL_UNHOOK_INT:
		if( pStack->Parameters.DeviceIoControl.InputBufferLength != sizeof(UCHAR) )
		{
			iStatus = STATUS_INVALID_PARAMETER;
			iTransfered = sizeof(UCHAR);
			break;
		}

		DPRINT( "IoDeviceControl: Unhooking Int 0x%.2X\n", *(UCHAR *)pIrp->AssociatedIrp.SystemBuffer );

		LoadIDT( &Idt );
		LoadINTVector( 
			&Idt, 
			*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
			&Vec );
		Vec.wHighOffset = OriginalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer].wHighOffset;
		Vec.wLowOffset = OriginalHandlers[*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer].wLowOffset;
		SaveINTVector( 
			&Idt, 
			*(UCHAR *)pIrp->AssociatedIrp.SystemBuffer, 
			&Vec );

		DPRINT( "IoDeviceControl: Unhooking successful\n" );
		break;

	case IOCTL_HOOK_ALL_INT:											//hook所有中断函数
		if( StartTime.QuadPart != 0 )
			break;

		DPRINT( "IoDeviceControl: Hooking all interrupts\n" );
		HookAllInterrupts();
		DPRINT( "IoDeviceControl: Hooking successful\n" );
		break;

	case IOCTL_DUMP_IDT:
		if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(INT_VECTOR) * 256 )
		{
			iStatus = STATUS_BUFFER_TOO_SMALL;
			iTransfered = sizeof(INT_VECTOR) * 256;
			break;
		}

		memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&OriginalHandlers, sizeof(INT_VECTOR) * 256 );
		iTransfered = sizeof(INT_VECTOR) * 256;
		break;

	case IOCTL_GET_INT_COUNTS:
		if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(__int64) * 256 )
		{
			iStatus = STATUS_BUFFER_TOO_SMALL;
			iTransfered = sizeof(__int64) * 256;
			break;
		}

		DPRINT( "IoDeviceControl: Retrieving interrupt counts\n" );

		__asm{ PUSHFD };
		__asm{ CLI };
		memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&InterruptCounts, sizeof(__int64) * 256 );
		__asm{ POPFD };
		iTransfered = sizeof(__int64) * 256;
		break;

	case IOCTL_GET_START_TIME:
		if( pStack->Parameters.DeviceIoControl.OutputBufferLength < sizeof(LARGE_INTEGER) )
		{
			iStatus = STATUS_BUFFER_TOO_SMALL;
			iTransfered = sizeof(LARGE_INTEGER);
			break;
		}

		memcpy( (void *)pIrp->AssociatedIrp.SystemBuffer, (const void *)&StartTime, sizeof(LARGE_INTEGER) );
		iTransfered = sizeof(LARGE_INTEGER);
		break;

	default:
		iStatus = STATUS_INVALID_PARAMETER;
		break;
	}

	pIrp->IoStatus.Status		= iStatus;
	pIrp->IoStatus.Information	= iTransfered;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return iStatus;
}