www.pudn.com > ntshell.rar > 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 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]); 
	} 
}