procmgr.cpp

远程控制系统,可以实现局域网内进行远程控制计算机,很方便,并且是学习用的非常好的资料.

// procmgr.cpp: implementation of the CProcMgr class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "resource.h"
#include "common.h"

#include "ntshell.h"
#include "sockmgr1.h"
#include "inputdlg.h"
#include "procmgr.h"

LRESULT CProcMgr::OnInitDialog(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM lParam, BOOL& /*bHandled*/)
{
	slave_sock = INVALID_SOCKET;
	processlist = NULL;
	SortColumn = 0;
	SizeSortOrder[0] = 0;
	SizeSortOrder[1] = 1;
	SizeSortOrder[2] = 2;
	SizeSortOrder[3] = 3;
	SizeSortOrder[4] = 4;
	SizeSortOrder[5] = 5;

	m_hList.Attach(GetDlgItem(IDC_LIST));
	m_hList.InsertColumn(0, "进程名", LVCFMT_LEFT, 150, 0);
	m_hList.InsertColumn(1, "PID", LVCFMT_RIGHT, 50, 0);
	m_hList.InsertColumn(2, "CPU", LVCFMT_RIGHT, 50, 0);
	m_hList.InsertColumn(3, "用户名", LVCFMT_LEFT, 130, 0);
	m_hList.InsertColumn(4, "内存使用", LVCFMT_RIGHT, 70, 0);
	m_hList.InsertColumn(5, "映像路径", LVCFMT_LEFT, 300, 0);

	DWORD dwStyle = m_hList.GetExtendedListViewStyle();
	dwStyle |= LVS_EX_FULLROWSELECT | LVS_EX_HEADERDRAGDROP;
	m_hList.SetExtendedListViewStyle(dwStyle);

	return TRUE;
}

LRESULT CProcMgr::OnDestroy(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM lParam, BOOL& /*bHandled*/)
{
	if (slave_sock != INVALID_SOCKET)
	{
		closesocket(slave_sock);
		slave_sock = INVALID_SOCKET;
	}

	while (processlist != NULL)
	{
		PPROCESS_LIST p = processlist;
		processlist = processlist->next;
		delete [] p;
	}

	return TRUE;
}

LRESULT CProcMgr::OnSize(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM lParam, BOOL& /*bHandled*/)
{
	::MoveWindow(GetDlgItem(IDC_LIST), 7, 7, LOWORD(lParam) - 14, HIWORD(lParam) - 14, TRUE);

	return TRUE;
}

LRESULT CProcMgr::OnTimer(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM /*lParam*/, BOOL& /*bHandled*/)
{
	NTSHELL_REQUEST req;
	req.ClientVersion = 1;
	req.Reserved = 0;
	req.RequestClass = PROCMGR_LISTPROCESS;
	req.Reserved2 = 0;

	WriteBuffer((char *)&req, sizeof(NTSHELL_REQUEST));

	FlushBuffer();

	SendPack();
	//CleanupFileList();

	return TRUE;
}

LRESULT CProcMgr::OnRun(WORD /*wNotifyCode*/, WORD /*wID*/, HWND /*hWndCtl*/, BOOL& /*bHandled*/)
{
	NTSHELL_REQUEST req;
	req.ClientVersion = 1;
	req.Reserved = 0;
	req.RequestClass = PROCMGR_RUN;
	req.Reserved2 = 0;

	PROCMGR_PROCESSRUN pr;
	pr.ShowWindow = TRUE;
	pr.Reserved = 0;

	CInputDlg<IDD_CREATEPROCESS> dlg;
	char szImagePath[INPUT_BUFFER_SIZE] = "";

	if (dlg.DoModal(m_hWnd, (LPARAM)szImagePath) == FALSE)
		return 0;

	WriteBuffer((char *)&req, sizeof(NTSHELL_REQUEST));
	WriteBuffer((char *)&pr, sizeof(PROCMGR_PROCESSRUN));
	WriteBuffer(szImagePath, strlen(szImagePath) + 1);

	FlushBuffer();

	if (SendPack() == SOCKMGR_SOCK_FAILED)
		MessageBox("连接不可用", "错误", MB_ICONERROR);

	return 0;
}

LRESULT CProcMgr::OnKillProcess(WORD /*wNotifyCode*/, WORD /*wID*/, HWND /*hWndCtl*/, BOOL& /*bHandled*/)
{
	int nItem = m_hList.GetSelectedIndex();
	PPROCMGR_PROCESSINFO p = &((PPROCESS_LIST)m_hList.GetItemData(nItem))->processinfo;

	CString str;

	str.Format("确定要结束进程 %s ？", (char *)p + p->OffsetOfProcessName);

	if (MessageBox(str, "警告", MB_ICONWARNING | MB_YESNO) == IDNO)
		return 0;

	NTSHELL_REQUEST req;
	req.ClientVersion = 1;
	req.Reserved = 0;
	req.RequestClass = PROCMGR_KILLPROCESS;
	req.Reserved2 = 0;

	PROCMGR_PROCESSKILL pkp;
	pkp.ProcessId = p->ProcessId;

	WriteBuffer((char *)&req, sizeof(NTSHELL_REQUEST));
	WriteBuffer((char *)&pkp, sizeof(PROCMGR_PROCESSKILL));

	FlushBuffer();

	if (SendPack() == SOCKMGR_SOCK_FAILED)
		MessageBox("连接不可用", "错误", MB_ICONERROR);

	return 0;
}

LRESULT CProcMgr::OnForceKillProcess(WORD /*wNotifyCode*/, WORD /*wID*/, HWND /*hWndCtl*/, BOOL& /*bHandled*/)
{
	int nItem = m_hList.GetSelectedIndex();
	PPROCMGR_PROCESSINFO p = &((PPROCESS_LIST)m_hList.GetItemData(nItem))->processinfo;

	CString str;

	str.Format("强制结束功能并不稳定，可能导致远程操作系统崩溃，请慎重使用。\r\n确定要结束进程 %s ？", (char *)p + p->OffsetOfProcessName);

	if (MessageBox(str, "警告", MB_ICONWARNING | MB_YESNO) == IDNO)
		return 0;

	NTSHELL_REQUEST req;
	req.ClientVersion = 1;
	req.Reserved = 0;
	req.RequestClass = PROCMGR_FORCEKILLPROCESS;
	req.Reserved2 = 0;

	PROCMGR_PROCESSKILL pkp;
	pkp.ProcessId = p->ProcessId;

	WriteBuffer((char *)&req, sizeof(NTSHELL_REQUEST));
	WriteBuffer((char *)&pkp, sizeof(PROCMGR_PROCESSKILL));

	FlushBuffer();

	if (SendPack() == SOCKMGR_SOCK_FAILED)
		MessageBox("连接不可用", "错误", MB_ICONERROR);

	return 0;
}

LRESULT CProcMgr::OnListEvent(int wParam, LPNMHDR pNMHDR, BOOL& /*bHandled*/)
{
	NMLISTVIEW *pNMListView = (NMLISTVIEW *)pNMHDR;
	int nItem = pNMListView->iItem;
	HMENU hMenu, hPopupMenu;
	POINT point;

	nItem = m_hList.GetSelectedIndex();

	switch (pNMHDR->code)
	{
	case NM_RCLICK:
		hMenu = LoadMenu(GetModuleHandle(NULL), (LPCTSTR)IDR_SUBMENU);
		hPopupMenu = GetSubMenu(hMenu, 1);

		if (nItem == -1)
		{
			EnableMenuItem(hPopupMenu, IDM_PROCESS_KILL, MF_GRAYED);
			EnableMenuItem(hPopupMenu, IDM_PROCESS_FORCEKILL, MF_GRAYED);
		}
		else
		{
			EnableMenuItem(hPopupMenu, IDM_PROCESS_KILL, MF_ENABLED);
			EnableMenuItem(hPopupMenu, IDM_PROCESS_FORCEKILL, MF_ENABLED);
		}

		GetCursorPos(&point);
		TrackPopupMenu(hPopupMenu, TPM_LEFTALIGN, point.x, point.y, 0, m_hWnd, NULL);

		DestroyMenu(hMenu);
		break;
	}

	return 0;
}

LRESULT CProcMgr::OnSocketRead(UINT /*uMsg*/, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/)
{
	int ret = RecvPack();

	if (ret == SOCKMGR_PENDING)
		return 0;

	if (ret == SOCKMGR_SOCK_FAILED)
	{
		char *p = GetErrorString();
		MessageBox(p, "", MB_ICONERROR);
		LocalFree(p);
	}

	if (ret == SOCKMGR_SUCCESS)
	{
		ULONG nDataSize = GetBufferDataSize();
		PBYTE pData = (PBYTE)HeapAlloc(GetProcessHeap(), 0, nDataSize);

		if (pData != NULL)
		{
			ReadBuffer((char *)pData, nDataSize);
			ParseDataPack(pData, nDataSize);
			HeapFree(GetProcessHeap(), 0, pData);
		}
	}

	return 0;
}

LRESULT CProcMgr::OnSocketWrite(UINT /*uMsg*/, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/)
{
	int ret = SendPack();

	if (ret == SOCKMGR_SOCK_FAILED)
	{
		char *p = GetErrorString();
		MessageBox(p, "", MB_ICONERROR);
		LocalFree(p);
	}

	return 0;
}

int CALLBACK CProcMgr::CompareFunc(LPARAM lParam1, LPARAM lParam2, LPARAM lParamSort)
{
	return 0;
}

LRESULT CProcMgr::OnConnectCreated(UINT /*uMsg*/, WPARAM wParam, LPARAM /*lParam*/, BOOL& bHandled)
{
	SetSocket((SOCKET)wParam);
	NTSHELL_REQUEST req;
	req.ClientVersion = 1;
	req.Reserved = 0;
	req.RequestClass = PROCMGR_LISTPROCESS;
	req.Reserved2 = 0;
	WriteBuffer((char *)&req, sizeof(NTSHELL_REQUEST));
	FlushBuffer();
	CleanupFileList();
	SetTimer(WM_TIMER, 1000);

	return 0;
}

LRESULT CProcMgr::OnConnectClosed(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM /*lParam*/, BOOL& bHandled)
{
	SetSocket(INVALID_SOCKET);
	KillTimer(WM_TIMER);

	return 0;
}

void CProcMgr::CleanupFileList()
{
	m_hList.DeleteAllItems();

	while (processlist != NULL)
	{
		PPROCESS_LIST p = processlist;
		processlist = processlist->next;
		delete []p;
	}
}

void CProcMgr::ParseDataPack(LPBYTE pData, ULONG nDataSize)
{
	PNTSHELL_RESULTSET rs = (PNTSHELL_RESULTSET)pData;
	CString str;
	PBYTE p = pData + sizeof(NTSHELL_RESULTSET);
	ULONG l = nDataSize - sizeof(NTSHELL_RESULTSET);

	if (rs->ErrorCode != ERROR_SUCCESS)
	{
		LPVOID lpMsgBuf;

		FormatMessage(	FORMAT_MESSAGE_ALLOCATE_BUFFER |
						FORMAT_MESSAGE_FROM_SYSTEM |
						FORMAT_MESSAGE_IGNORE_INSERTS,
						NULL,
						rs->ErrorCode,
						0,
						(LPTSTR)&lpMsgBuf,
						0,
						NULL
						);

		MessageBox((LPTSTR)lpMsgBuf, "", MB_ICONERROR);
		LocalFree(lpMsgBuf);
		return;
	}

	if (rs->ResultClass == PROCMGR_LISTPROCESS)
	{
		int i, count;
		CString str;
		PPROCESS_LIST pList, *ppList = &processlist;
		PPROCMGR_PROCESSINFO ppi;

		while (*ppList != NULL)
		{
			for (i = 0, count = 0; ; count++)
			{
				if (i >= (int)l || count >= (int)rs->NumberOfResults)
				{
					LV_FINDINFO lvfi;
					lvfi.flags = LVFI_PARAM;
					lvfi.lParam = (LPARAM)*ppList;
					int nItem = m_hList.FindItem(&lvfi, -1);
					m_hList.DeleteItem(nItem);		//移除不存在的进程
					pList = *ppList;
					*ppList = (*ppList)->next;		//从链表中移除
					delete [] pList;
					goto Loop1;
				}

				ppi = (PPROCMGR_PROCESSINFO)(p + i);

				if ((*ppList)->processinfo.ProcessId == ppi->ProcessId)
				{
					LV_FINDINFO lvfi;
					lvfi.flags = LVFI_PARAM;
					lvfi.lParam = (LPARAM)*ppList;
					int nItem = m_hList.FindItem(&lvfi, -1);

					(*ppList)->processinfo.TotalPrivateBytes = ppi->TotalPrivateBytes;
					(*ppList)->processinfo.TotalVirtualSizeBytes = ppi->TotalVirtualSizeBytes;
					(*ppList)->processinfo.ProcessCpuTime = ppi->ProcessCpuTime;

					str.Format("%dk", ppi->TotalPrivateBytes / 1024);
					m_hList.SetItemText(nItem, 4, str);		//更新进程信息
					break;
				}

				i += ppi->OffsetOfNextEntry;
			}

			ppList = &(*ppList)->next;
Loop1:
			continue;
		}

		for (i = 0, count = 0; i < (int)l && count < (int)rs->NumberOfResults; count++)
		{
			ppi = (PPROCMGR_PROCESSINFO)(p + i);
			pList = processlist;

			while (pList != NULL)
			{
				if (pList->processinfo.ProcessId == ppi->ProcessId)
					break;
				pList = pList->next;
			}

			if (pList == NULL)
			{
				pList = (PPROCESS_LIST)new BYTE[sizeof(PROCESS_LIST) + ppi->OffsetOfNextEntry - sizeof(PROCMGR_PROCESSINFO)];
				memcpy(&(pList->processinfo), ppi, ppi->OffsetOfNextEntry);
				pList->next = processlist;
				processlist = pList;

				char *pName = (char *)ppi + ppi->OffsetOfProcessName;
				m_hList.InsertItem(LVIF_TEXT | LVIF_PARAM, count, pName, 0, 0, 0, (LPARAM)pList);

				str.Format("%d", ppi->ProcessId);
				m_hList.SetItemText(count, 1, str);

				pName = (char *)ppi + ppi->OffsetOfUserName;
				m_hList.SetItemText(count, 3, pName);

				str.Format("%dk", ppi->TotalPrivateBytes / 1024);
				m_hList.SetItemText(count, 4, str);

				pName = (char *)ppi + ppi->OffsetOfImagePath;
				m_hList.SetItemText(count, 5, pName);

				if (ppi->ProcessId == 0)
				{
					m_hList.SetItemText(count, 0, "System Idle Process");
					m_hList.SetItemText(count, 3, "SYSTEM");
				}
			}

			i += ppi->OffsetOfNextEntry;
		}

		m_hList.SortItems(CompareFunc, (LPARAM)SizeSortOrder[SortColumn]);
	}
}