lee_functions.cpp

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

/************************************************************************
* 文件名称:LEE_Functions.cpp                                           
* 作    者:李骥
*************************************************************************/



#ifdef __cplusplus
extern "C"
{
#endif

#include <NTDDK.h>

#ifdef __cplusplus
}
#endif 

#include "LEE_Functions.h"
#include "NtFunctions.h"

////////////////////定义所用到的全局变量///////////////     

unsigned long OldCr0;   
UNICODE_STRING DeviceNameString;   
UNICODE_STRING LinkDeviceNameString;   

KEVENT  kEvent;

SYSTEM_PROCESS_INFORMATION ControlProcess;

HANDLE hMyThread;

SYSTEM_PROCESS_INFORMATION System_PInformation_List[256];

typedef NTSTATUS (*ST)( IN HANDLE ThreadHandle, OUT PULONG PreviousSuspendCount OPTIONAL ); 
typedef NTSTATUS (*LEE_SetAffinityThread)(IN HANDLE hThread,IN ULONG dwMask);



extern PSSDT   KeServiceDescriptorTable;


ZWQUERYSYSTEMINFORMATION OldZwQuerySystemInformation;  

ST        st;
LEE_SetAffinityThread  LST;

PVOID pf=(PVOID)0x805e145e ;

PLIST_ENTRY pGlobeListEntry[60];


NTSTATUS MyGetProcessInformation(PMY_PROCESS_INFO pMy_Process_Info,PMY_THREAD_INFO pMy_Thread_Info,ULONG TotalSize)
{
	ULONG cbBuffer = 0x80000;                                           //申请内存大小BYTE
	ULONG dwCount  = 0;
	PVOID pBuffer  = NULL;
	PSYSTEM_PROCESS_INFORMATION pInfo;

	PSYSTEM_THREAD_INFORMATION pThread_Info;

	MY_PROCESS_INFO My_Process_Info_Temp; 
	MY_THREAD_INFO My_Thread_Info_Temp; 
	ULONG ProcessCount=0;
	ULONG ThreadCount=0;
	ULONG ThreadTotal=0;

	/*__try
	{
		ProbeForWrite((PVOID)pMy_Process_Info,TotalSize,sizeof(CHAR));

	}
	__except( EXCEPTION_EXECUTE_HANDLER )
	{
		DbgPrint("failed!");
	}*/
		

	pBuffer = ExAllocatePool(PagedPool, cbBuffer);                       //申请分页内存
	// 获取进程信息
	KdPrint(("We Use ZwQuerySystemInformation!\n"));

	SYSTEMINFOCLASS SC=SystemProcessInformation;
	ZwQuerySystemInformation( SC,
		pBuffer,
		cbBuffer,
		NULL);

	pInfo = (PSYSTEM_PROCESS_INFORMATION)pBuffer;
		
	
	

	for( ; ; )
	{
		My_Process_Info_Temp.dThreadCount=pInfo->dThreadCount;
		My_Process_Info_Temp.InheritedFromProcessID=pInfo->dParentProcessID;
		My_Process_Info_Temp.ProcessID=pInfo->dUniqueProcessId;

		RtlZeroMemory(My_Process_Info_Temp.ProcessName,sizeof(WCHAR)*256);

		RtlCopyMemory(My_Process_Info_Temp.ProcessName,pInfo->ProcessName.Buffer,pInfo->ProcessName.Length);

		My_Process_Info_Temp.ThreadInfos=pMy_Thread_Info+ThreadTotal;
		

		for( ThreadCount=0;ThreadCount<pInfo->dThreadCount ;ThreadCount++ )
		{
			pThread_Info=(PSYSTEM_THREAD_INFORMATION(pInfo->ThreadInfos))+ThreadCount;
			My_Thread_Info_Temp.ProcessID=(ULONG)(pThread_Info->ClientID.UniqueProcess);
			My_Thread_Info_Temp.ThreadID=(ULONG)(pThread_Info->ClientID.UniqueThread);
			My_Thread_Info_Temp.ThreadCnt=ThreadCount;

			My_Thread_Info_Temp.CpuMask=0x00000001;
			My_Thread_Info_Temp.CpuRate=100;


			PMY_THREAD_INFO pThread_In=(PMY_THREAD_INFO)(My_Process_Info_Temp.ThreadInfos)+ThreadCount;
			RtlCopyMemory(pThread_In,&My_Thread_Info_Temp,sizeof(MY_THREAD_INFO));	
			ThreadTotal++;

		}

	
		RtlCopyMemory(pMy_Process_Info+ProcessCount,&My_Process_Info_Temp,sizeof(MY_PROCESS_INFO));
		ProcessCount++;

		
		

		

		if (pInfo->NextEntryDelta == 0)
		{
			ULONG *pProcessCount=((ULONG *)pMy_Thread_Info)-1;
			RtlCopyMemory(pProcessCount,&ProcessCount,sizeof(ULONG));
			break;
		}


		pInfo = (PSYSTEM_PROCESS_INFORMATION)(((ULONG)pInfo) + pInfo->NextEntryDelta);
	}
	DbgPrint("\n0x%08X\n0x%08X   over\n",pMy_Process_Info,pMy_Thread_Info);
	ExFreePool(pBuffer);


	return 1;



}

/////////////////////////////////////////////////////////////////////////
NTSTATUS MySetProcessInformation(PMY_PROCESS_INFO pMy_Process_Info,PMY_THREAD_INFO pMy_Thread_Info,ULONG TotalSize)
{
	

	
	MY_THREAD_INFO My_Thread_Info_Temp; 

	ULONG ProcessCount=*(ULONG *)((ULONG *)pMy_Thread_Info-1);
	ULONG ThreadCount=0;

	PVOID p=(PVOID)0x80512fcf ;
	LST=(LEE_SetAffinityThread)p;

	for(ULONG i=0 ;i<ProcessCount ;i++ )
	{
		

			pMy_Thread_Info=(PMY_THREAD_INFO)(pMy_Process_Info->ThreadInfos);

			for(ULONG j=0;j<pMy_Process_Info->dThreadCount;j++)
			{
				My_Thread_Info_Temp.ProcessID=pMy_Thread_Info->ProcessID;
				My_Thread_Info_Temp.ThreadID=pMy_Thread_Info->ThreadID;
				My_Thread_Info_Temp.ThreadCnt=pMy_Thread_Info->ThreadCnt;

				My_Thread_Info_Temp.CpuMask=pMy_Thread_Info->CpuMask;
				My_Thread_Info_Temp.CpuRate=pMy_Thread_Info->CpuRate;

				

				CLIENT_ID cid;
				cid.UniqueProcess=(HANDLE)My_Thread_Info_Temp.ProcessID;
				cid.UniqueThread=(HANDLE)My_Thread_Info_Temp.ThreadID;
				
				
				
				SetThreadAffinity_FromUser(My_Thread_Info_Temp.ThreadID,My_Thread_Info_Temp.CpuMask);
				
				/*NtOpenThread(&hThread,THREAD_ALL_ACCESS,NULL,&cid);
				LST(hThread,(ULONG)My_Thread_Info_Temp.CpuMask);*/

				

				pMy_Thread_Info++;
			}


			
		pMy_Process_Info++;
	}


	return 1;



}
///////////////////////////////////////////////////////////////////////////////////////

#define PIDOFFSET 0X84      //EPROCESS中UniqueProcessId偏移
#define FLINKOFFSET 0x88    //EPROCESS中ActiveProcessLinks偏移
#define ObjectTable 0xc4    //EPROCESS中ObjectTable偏移
#define TableList 0x1c      //ObjectTable中HandleTableList偏移
#define PIDOFFSET2 0x008    //ObjectTable中UniqueProcessId偏移
#define QuotaProcess 0x004  //ObjectTable中QuotaProcess偏移


#define   ThreadListHead_EP  0x190     //ThreadListHead in EPROCESS

#define   ThreadListEntry_KT 0x1b0   //_LIST_ENTRY IN KTHREAD
#define   ThreadListEntry_ET 0x22c// : _LIST_ENTRY in ETHREAD

#define  CidOffset_ET     0x1ec         // _CLIENT_ID In ETHREAD

//通过ActiveProcessLinks遍历进程
VOID 
SetAllThreadAffinity(ULONG Mask)
{
	ULONG eproc=0x0;
	ULONG eth=0x0;
	int current_PID=0;
	int start_PID=0;
	int count=0;

	HANDLE current_TID=0;
	HANDLE start_TID=0;

	//PVOID p=(PVOID)0x805039c6 ;//0x804fe0e0  ;
	//LST=(LEE_SetAffinityThread)p;


	PLIST_ENTRY plist_active_procs;
	DbgPrint("SetAllThreadAffinity:\n");
	eproc=(ULONG)PsGetCurrentProcess();
	start_PID=*((int *)(eproc+PIDOFFSET));
	current_PID=start_PID;



	while(1)
	{
		if((count>=1)&&(start_PID==current_PID))
		{
			return;
		}
		else
		{

			

			plist_active_procs=(PLIST_ENTRY)(eproc+FLINKOFFSET);
			eproc=(ULONG)plist_active_procs->Flink;
			eproc=eproc-FLINKOFFSET;
			current_PID=*((int *)(eproc+PIDOFFSET));
			if(current_PID<=0)
				current_PID=0;

			PLIST_ENTRY TListHead=(PLIST_ENTRY)(eproc+ThreadListHead_EP);
			eth=(ULONG)TListHead->Flink;
			eth=eth-ThreadListEntry_ET;

			PCLIENT_ID pCid=(PCLIENT_ID)(eth+ CidOffset_ET);
			current_TID=pCid->UniqueThread;
			


			ULONG count_t=0;
			while(1)
			{
				if(current_PID==0||current_TID==0)
				{
					break;
				}

				else
				{	DbgPrint("		Num:%02d	thread id %08x\n",count_t,current_TID);

					//KeSetAffinityThread((HANDLE)eth,Mask);

				ULONG *p=(ULONG *)(eth+0x124);
					*p=(ULONG)Mask;


					PLIST_ENTRY TList=(PLIST_ENTRY)(eth+ThreadListEntry_ET);
					eth=(ULONG)(TList->Flink);
					eth=eth-ThreadListEntry_ET;

					pCid=(PCLIENT_ID)(eth+ CidOffset_ET);
					current_TID=pCid->UniqueThread;
					

					

					/*if(pCid->UniqueProcess==(HANDLE)PID&&pCid->UniqueThread==(HANDLE)TID)
					{*/
						

					/*}*/


					
				}

				count_t++;

			}
			




			DbgPrint("process id %08x,address %08x\n",current_PID,eproc);
			count++;
		}
	}
	DbgPrint("Total number is %d\n\n",count);
}

//通过HandleTable遍历进程
VOID
FindByHandleTable()
{
	ULONG eproc;
	PLIST_ENTRY start_list,HandleTableList=NULL;
	ULONG addr;
	int count=0;
	int pid;
	DbgPrint("FindByObjectTable:\n");
	eproc=(ULONG)PsGetCurrentProcess();
	HandleTableList=(PLIST_ENTRY)(*(PULONG)((ULONG)eproc+ObjectTable)+TableList);
	start_list=HandleTableList;
	do {
		pid=*(PULONG)((ULONG)HandleTableList+PIDOFFSET2-TableList);
		if(pid!=4)
			addr=*(PULONG)((ULONG)HandleTableList+QuotaProcess-TableList);
		else 
			addr=0xffffffff;
		DbgPrint("process id %4d,address %8x\n",pid,addr);
		count++;
		if(pid!=4)
			addr=*(PULONG)((ULONG)HandleTableList+PIDOFFSET2-TableList);
		HandleTableList=HandleTableList->Flink;
	} while(start_list!=HandleTableList);
	DbgPrint("Total number is %d\n",count);
}
//////////////////////////////////////////////////////////////////////

VOID SystemThread(IN PVOID pContext)
{
	KIRQL irql;
	KeRaiseIrql(APC_LEVEL,&irql);
	/*for(i=0;i<1000000;i++)
	{

		DbgPrint("SystemThread!\n");

	}*/
	KeLowerIrql(irql);
	PsTerminateSystemThread(STATUS_SUCCESS);
}





VOID MyThread(IN PVOID pContext)
{


	//KIRQL irql;
	//ULONG ProcessID;

	////st =(ST)(pf);

	////st(hMyThread,&x);

	//_SYSTEM_THREADS *pthread,*temp;
	//ULONG ThreadCount=ControlProcess.dThreadCount;

	//pthread=(_SYSTEM_THREADS *)(ControlProcess.ThreadInfos);





	//HANDLE hThread[10];
	//ULONG PreCount[10];


	////	temp=pthread;
	////	for(ULONG i=0;i<ThreadCount;i++)
	////	{



	////		NtOpenThread(&hThread[i],2,0,&(temp->ClientID));
	////		//NtSuspendThread(hThread[i],&PreCount[i]);

	////		temp++;
	////	}	

	////	KeStallExecutionProcessor(30);

	////	for(ULONG i=0;i<ThreadCount;i++)
	////	{

	////		//NtResumeThread(hThread[i],&PreCount[i]);
	////	}	
	////	KeStallExecutionProcessor(70);
	/////*}*/

	//////KeWaitForSingleObject(&kEvent, Executive, KernelMode, FALSE, 0);



	//irql=KeGetCurrentIrql();
	//ProcessID=(ULONG)PsGetCurrentProcessId();

	//DbgPrint("DriverEntry: IRQL %d Process %d",irql,ProcessID);

	//DbgPrint("MyThread!Over!!!\n");





	//KeLowerIrql(oldirql);

	PsTerminateSystemThread(STATUS_SUCCESS);
}











//////////////////////////////////////////////////////////////////////

VOID 
SetThreadAffinity_FromUser(ULONG TID,ULONG Mask)
{
	ULONG eproc=0x0;
	ULONG eth=0x0;
	int current_PID=0;
	int start_PID=0;
	int count=0;

	HANDLE current_TID=0;
	HANDLE start_TID=0;

	//PVOID p=(PVOID)0x805039c6 ;//0x804fe0e0  ;
	//LST=(LEE_SetAffinityThread)p;


	PLIST_ENTRY plist_active_procs;
	DbgPrint("SetAllThreadAffinity:\n");
	eproc=(ULONG)PsGetCurrentProcess();
	start_PID=*((int *)(eproc+PIDOFFSET));
	current_PID=start_PID;



	while(1)
	{
		if((count>=1)&&(start_PID==current_PID))
		{
			return;
		}
		else
		{



			plist_active_procs=(PLIST_ENTRY)(eproc+FLINKOFFSET);
			eproc=(ULONG)plist_active_procs->Flink;
			eproc=eproc-FLINKOFFSET;
			current_PID=*((int *)(eproc+PIDOFFSET));
			if(current_PID<=0)
				current_PID=0;

			PLIST_ENTRY TListHead=(PLIST_ENTRY)(eproc+ThreadListHead_EP);
			eth=(ULONG)TListHead->Flink;
			eth=eth-ThreadListEntry_ET;

			PCLIENT_ID pCid=(PCLIENT_ID)(eth+ CidOffset_ET);
			current_TID=pCid->UniqueThread;



			ULONG count_t=0;
			while(1)
			{
				if(current_PID==0||current_TID==0)
				{
					break;
				}

				else
				{	//DbgPrint("		Num:%02d	thread id %08x\n",count_t,current_TID);

				//KeSetAffinityThread((HANDLE)eth,Mask);

				if((ULONG)current_TID==TID)
				{
					ULONG *p=(ULONG *)(eth+0x124);
					*p=(ULONG)Mask;

				}
				


				PLIST_ENTRY TList=(PLIST_ENTRY)(eth+ThreadListEntry_ET);
				eth=(ULONG)(TList->Flink);
				eth=eth-ThreadListEntry_ET;

				pCid=(PCLIENT_ID)(eth+ CidOffset_ET);
				current_TID=pCid->UniqueThread;




				/*if(pCid->UniqueProcess==(HANDLE)PID&&pCid->UniqueThread==(HANDLE)TID)
				{*/


				/*}*/



				}

				count_t++;

			}





			//DbgPrint("process id %08x,address %08x\n",current_PID,eproc);
			count++;
		}
	}
	//DbgPrint("Total number is %d\n\n",count);
}