lee_process_manager.cpp

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

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



#include "stdafx.h"
#include "LEE_PROCESS_MANAGER.h"
#include "Process_Functions.h"

#include "ProcessList.h"
#include "ioctls.h"

#include "AssistFunction.h"

#include "LoadDrv/Driver.h"

#define MAX_LOADSTRING 100

// 全局变量:
HINSTANCE hInst;								// 当前实例
TCHAR szTitle[MAX_LOADSTRING];					// 标题栏文本
TCHAR szWindowClass[MAX_LOADSTRING];			// 主窗口类名

HWND hParentWnd;
HWND hList;
HANDLE hDevice;

BYTE *pProcessMemory;

ULONG CurrentProcessCnt;
ULONG CurrentThreadCnt;




DWORD PID;
PROCESSENTRY32 pe32;


// 此代码模块中包含的函数的前向声明:
ATOM				MyRegisterClass(HINSTANCE hInstance);
BOOL				InitInstance(HINSTANCE, int);
LRESULT CALLBACK	WndProc(HWND, UINT, WPARAM, LPARAM);
INT_PTR CALLBACK	About(HWND, UINT, WPARAM, LPARAM);
INT_PTR CALLBACK Thread_Func(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam);


DWORD WINAPI TaskThread( LPVOID lpParam );
typedef DWORD (* TTD)(LPVOID lParam);

int APIENTRY _tWinMain(HINSTANCE hInstance,
                     HINSTANCE hPrevInstance,
                     LPTSTR    lpCmdLine,
                     int       nCmdShow)
{

	__try
	{
		
	}
	__except( 1 )
	{
	}

	TCHAR strDriver[MAX_PATH];
	if ( !GetDriverPath( strDriver, MAX_PATH ) )
	{
		puts( "Can't find driver!" );
		return -1;
	}

	hDevice = LoadDriver( strDriver );
	if ( INVALID_HANDLE_VALUE == hDevice )
	{
		puts( "Load driver failed!" );
		return -1;
	}

	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);

 	// TODO: 在此放置代码。
	MSG msg;
	HACCEL hAccelTable;


	
	///////////////////////建立设备
	//hDevice = 
	//	CreateFile(L"\\\\.\\MyProcessManage",
	//	GENERIC_READ | GENERIC_WRITE,
	//	0,		// share mode none
	//	NULL,	// no security
	//	OPEN_EXISTING,
	//	FILE_ATTRIBUTE_NORMAL,
	//	NULL );		// no template

	//if(hDevice==NULL)
	//{
	//	MessageBox(NULL,L"No Driver!",NULL,MB_OK);
	//}

	////////////////////////////初始化内存

	pProcessMemory=new BYTE[0x10000];
	CurrentProcessCnt=0;
	CurrentThreadCnt=0;


	// 初始化全局字符串
	LoadString(hInstance, IDS_APP_TITLE, szTitle, MAX_LOADSTRING);
	LoadString(hInstance, IDC_LEE_PROCESS_MANAGER, szWindowClass, MAX_LOADSTRING);
	MyRegisterClass(hInstance);


	// 执行应用程序初始化:

		InitCommonControls();
		OleInitialize(NULL);
	if (!InitInstance (hInstance, nCmdShow))
	{
		return FALSE;
	}

	hAccelTable = LoadAccelerators(hInstance, MAKEINTRESOURCE(IDC_LEE_PROCESS_MANAGER));




	if (hDevice == INVALID_HANDLE_VALUE)
	{
		printf("Failed to obtain file handle to device: "
			"%s with Win32 error code: %d\n",
			"MyWDMDevice", GetLastError() );
		
	}


	

	// 主消息循环:
	while (GetMessage(&msg, NULL, 0, 0))
	{
		if (!TranslateAccelerator(msg.hwnd, hAccelTable, &msg))
		{
			

			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}


	delete []pProcessMemory;

	 UnloadDriver( hDevice );

	return (int) msg.wParam;
}



//
//  函数: MyRegisterClass()
//
//  目的: 注册窗口类。
//
//  注释:
//
//    仅当希望
//    此代码与添加到 Windows 95 中的“RegisterClassEx”
//    函数之前的 Win32 系统兼容时，才需要此函数及其用法。调用此函数十分重要，
//    这样应用程序就可以获得关联的
//    “格式正确的”小图标。
//
ATOM MyRegisterClass(HINSTANCE hInstance)
{
	WNDCLASSEX wcex;

	wcex.cbSize = sizeof(WNDCLASSEX);

	wcex.style			= CS_HREDRAW | CS_VREDRAW;
	wcex.lpfnWndProc	= WndProc;
	wcex.cbClsExtra		= 0;
	wcex.cbWndExtra		= 0;
	wcex.hInstance		= hInstance;
	wcex.hIcon			= LoadIcon(hInstance, MAKEINTRESOURCE(IDI_LEE_PROCESS_MANAGER));
	wcex.hCursor		= LoadCursor(NULL, IDC_ARROW);
	wcex.hbrBackground	= (HBRUSH)(COLOR_WINDOW+1);
	wcex.lpszMenuName	= MAKEINTRESOURCE(IDC_LEE_PROCESS_MANAGER);
	wcex.lpszClassName	= szWindowClass;
	wcex.hIconSm		= LoadIcon(wcex.hInstance, MAKEINTRESOURCE(IDI_SMALL));

	return RegisterClassEx(&wcex);
}

//
//   函数: InitInstance(HINSTANCE, int)
//
//   目的: 保存实例句柄并创建主窗口
//
//   注释:
//
//        在此函数中，我们在全局变量中保存实例句柄并
//        创建和显示主程序窗口。
//
BOOL InitInstance(HINSTANCE hInstance, int nCmdShow)
{
   HWND hWnd;

   hInst = hInstance; // 将实例句柄存储在全局变量中

   hWnd = CreateWindow(szWindowClass, szTitle, WS_OVERLAPPEDWINDOW,
      CW_USEDEFAULT, 0, CW_USEDEFAULT, 0, NULL, NULL, hInstance, NULL);

   if (!hWnd)
   {
      return FALSE;
   }


   hParentWnd=hWnd;

   ULONG size=sizeof(SYSTEM_PROCESS_INFORMATION);


  hList=CreateListView(hWnd,nCmdShow);



   InitListViewColumns_Process(hList);

  /* GetProcessList(hList);*/

   GetProcess_From_Drv(hList);


 

   ShowWindow(hWnd, nCmdShow);
   UpdateWindow(hWnd);
   

   return TRUE;
}

//
//  函数: WndProc(HWND, UINT, WPARAM, LPARAM)
//
//  目的: 处理主窗口的消息。
//
//  WM_COMMAND	- 处理应用程序菜单
//  WM_PAINT	- 绘制主窗口
//  WM_DESTROY	- 发送退出消息并返回
//
//
LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
	int wmId, wmEvent;
	PAINTSTRUCT ps;
	HDC hdc;
	



	switch (message)
	{
	case WM_COMMAND:
		wmId    = LOWORD(wParam);
		wmEvent = HIWORD(wParam);
		// 分析菜单选择:
		switch (wmId)
		{
		case IDM_ABOUT:
			DialogBox(hInst, MAKEINTRESOURCE(IDD_ABOUTBOX), hWnd, About);
			break;
		case IDM_EXIT:
			DestroyWindow(hWnd);
			break;

		case ID_RENEW_PROCESS:
			GetProcess_From_Drv(hList);


			break;
		case ID_SET_AFFINITY:
			{
				DWORD BytesReturned;
				BOOL result=DeviceIoControl(hDevice, IOCTL_SETPROCESS, pProcessMemory, 0x10000,NULL, 0,&BytesReturned, NULL);

				break;

			}
		case ID_SET_AFFINITY_ALL:
			{
				DWORD BytesReturned;
				BOOL result=DeviceIoControl(hDevice, IOCTL_SETPROCESS_ALL_AFFINITY, pProcessMemory, 0x10000,NULL, 0,&BytesReturned, NULL);

				break;

			}
		case ID_RAISE_FUNC:
			{

				TTD ttd=TaskThread;

				DWORD BytesReturned;
				BOOL result=DeviceIoControl(hDevice, IOCTL_GIVE_FUNCTION, &ttd, 0x04,NULL, 0,&BytesReturned, NULL);


			}
			break;
			
		default:
			return DefWindowProc(hWnd, message, wParam, lParam);
		}
		break;
	case WM_PAINT:
		hdc = BeginPaint(hWnd, &ps);
		// TODO: 在此添加任意绘图代码...
		EndPaint(hWnd, &ps);
		break;
	case WM_NOTIFY:
		{
			NMHDR *pnmh;
			pnmh=(LPNMHDR )lParam;

			if(pnmh->hwndFrom==hList&&pnmh->code==NM_DBLCLK)
			{
				LRESULT iSelect=SendMessage(hList,LVM_GETNEXTITEM,-1,LVNI_FOCUSED);
				CurrentProcessCnt=(ULONG)iSelect;
				DialogBoxParam(hInst, MAKEINTRESOURCE(IDD_THREAD_DIALOG), hWnd, Thread_Func,(LPARAM)iSelect);

			}


			break;
		}
		
		

	case WM_KEYDOWN:
		//ListProcessThreads(0);
		{
			BYTE *pProcessMemory=new BYTE[0x10000];

			DWORD BytesReturned;
   			BOOL result=DeviceIoControl(hDevice, IOCTL_GETPROCESS,NULL, 0, pProcessMemory, 0x10000,&BytesReturned, NULL);

			DWORD dw=GetLastError();
 			delete []pProcessMemory;


		}

		
		break;
	case WM_DESTROY:
		PostQuitMessage(0);

	
		break;


	default:
		return DefWindowProc(hWnd, message, wParam, lParam);
	}
	return 0;
}

// “关于”框的消息处理程序。
INT_PTR CALLBACK About(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	UNREFERENCED_PARAMETER(lParam);
	switch (message)
	{
	case WM_INITDIALOG:
		return (INT_PTR)TRUE;

	case WM_COMMAND:
		if (LOWORD(wParam) == IDOK || LOWORD(wParam) == IDCANCEL)
		{
			EndDialog(hDlg, LOWORD(wParam));
			return (INT_PTR)TRUE;
		}
		break;
	}
	return (INT_PTR)FALSE;
}


INT_PTR CALLBACK Thread_Func(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	HWND hEditWnd;
	HANDLE hProcessSnap;
	int Slider=0;
	WCHAR tempstr[256];

	// Take a snapshot of all processes in the system.
	hProcessSnap = CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 );
	
	int iSelect;

	UNREFERENCED_PARAMETER(lParam);
	switch (message)
	{
	case WM_INITDIALOG:
		{
			iSelect=(int)lParam;
		
		
						
			HWND hList_Thread=GetDlgItem(hDlg,IDC_THREAD_LIST);
			InitListViewColumns_Thread(hList_Thread);

			HWND hAffinity_Combo=GetDlgItem(hDlg,IDC_AFFINITY_COMBO);
			SendMessage(hAffinity_Combo,CB_ADDSTRING,NULL,(LPARAM)L"0x0001"); 
			SendMessage(hAffinity_Combo,CB_ADDSTRING,NULL,(LPARAM)L"0x0002"); 
			SendMessage(hAffinity_Combo,CB_ADDSTRING,NULL,(LPARAM)L"0x0003"); 

			//HWND hCpu_Check=GetDlgItem(hDlg,IDC_CPU1_CHECK);
			

			
					
			GetThread_To_Thread_List_Index(iSelect,hList_Thread);
				
				
			break;
	

		}
		
	case WM_HSCROLL:

	     Slider =(int) SendDlgItemMessage( hDlg, IDC_CPU_RATE_SLIDER, TBM_GETPOS, 0U, 0L );
		 Slider =GetSliderParam( Slider );
		 hEditWnd=GetDlgItem(hDlg,IDC_CPU_RATE);
		 swprintf_s(tempstr,L"%2d",Slider);
		 SetDlgItemText(hDlg,IDC_CPU_RATE,tempstr);
			break;

	case WM_COMMAND:

		
		if (LOWORD(wParam) == IDOK || LOWORD(wParam) == IDCANCEL)
		{
			EndDialog(hDlg, LOWORD(wParam));
			return (INT_PTR)TRUE;
		}
		else if(LOWORD(wParam)==ID_APPLY)
		{
			LRESULT lResult1 = SendMessage( GetDlgItem(hDlg,IDC_CPU1_CHECK), BM_GETCHECK, 0,0);
			LRESULT lResult2 = SendMessage( GetDlgItem(hDlg,IDC_CPU2_CHECK), BM_GETCHECK, 0,0);
			LRESULT lResult3 = SendMessage( GetDlgItem(hDlg,IDC_CPU3_CHECK), BM_GETCHECK, 0,0);
			LRESULT lResult4 = SendMessage( GetDlgItem(hDlg,IDC_CPU4_CHECK), BM_GETCHECK, 0,0);

			if(lResult1==BST_CHECKED)lResult1=1;
			else lResult1=0;
			if(lResult2==BST_CHECKED)lResult2=1;
			else lResult2=0;
			if(lResult3==BST_CHECKED)lResult3=1;
			else lResult3=0;
			if(lResult4==BST_CHECKED)lResult4=1;
			else lResult4=0;
			
			ULONG Result=ULONG(lResult1|(lResult2<<1)|(lResult3<<2)|(lResult4<<3));

			HWND hList_Thread=GetDlgItem(hDlg,IDC_THREAD_LIST);
			LRESULT iSelect=SendMessage(hList_Thread,LVM_GETNEXTITEM,-1,LVNI_FOCUSED);
			CurrentThreadCnt=(ULONG)iSelect;


			PMY_PROCESS_INFO pCP=(((PMY_PROCESS_INFO)pProcessMemory)+CurrentProcessCnt);
			PMY_THREAD_INFO pCT=((PMY_THREAD_INFO)pCP->ThreadInfos)+CurrentThreadCnt;
			pCT->CpuMask=Result;

		
			LRESULT lResult = SendMessage( hList_Thread, LVM_DELETEITEM, (WPARAM)iSelect, 0 );  

			Install_Thread_ToSub_From_Drv(*pCT,hList_Thread,(int)iSelect);

		




		}
		break;
		

		

	}
	return (INT_PTR)FALSE;

}



DWORD WINAPI TaskThread( LPVOID lpParam )
{
	for (int i=0;i<100000000000;i++)
	{
	}

	HANDLE hProcess=GetCurrentProcess();
	//MessageBox(hParentWnd,L"OK",NULL,MB_OK);
	
	return 1;
}