procfilterdlg.cpp

Strace： Webcrazy原来写的很好的代码

#include "stdafx.h"
#include "StraceGui.h"
#include "ListCtrlEx.h"
#include "ProcFilterDlg.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

DWORD PsGetProcessIdFromModuleName(LPCSTR szName);
typedef DWORD (__stdcall *EnumProcessesProc)(DWORD* ProcessesId,
						DWORD SizeofProcessesIds,
						DWORD* done);
typedef DWORD (__stdcall *GetModuleBaseNameAProc)(HANDLE hProcess,HMODULE hMod,
							char* szProcessName,
							DWORD SizeofszProcessName/* sizeof szProcessName*/ );

typedef DWORD (__stdcall 
			*EnumProcessModulesProc)(HANDLE hProcess,
							 HMODULE* hModule /*array*/,
							 DWORD SizeofhModule/* sizeof(hModule) */,
							 DWORD* done);

EnumProcessesProc pEnumProcesses;
GetModuleBaseNameAProc pGetModuleBaseNameA;
EnumProcessModulesProc pEnumProcessModules;


#define NUM_COLUMNS	2

_TCHAR *gProFilterColumnLabel[NUM_COLUMNS] =
{
	_T("Process Name"), _T("PID")
};

int gProFilterColumnFormat[NUM_COLUMNS] = 
{
	LVCFMT_LEFT, LVCFMT_RIGHT
};

int gProFilterColumnWidth[NUM_COLUMNS] = 
{
	80, 20 // %
};




/////////////////////////////////////////////////////////////////////////////
// CProcFilterDlg dialog


CProcFilterDlg::CProcFilterDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CProcFilterDlg::IDD, pParent)
{
	//{{AFX_DATA_INIT(CProcFilterDlg)
		// NOTE: the ClassWizard will add member initialization here
	//}}AFX_DATA_INIT


}
CProcFilterDlg::~CProcFilterDlg()
{

}


void CProcFilterDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	//{{AFX_DATA_MAP(CProcFilterDlg)
	DDX_Control(pDX, IDC_PROCESSLIST, m_ListCtrl);
	//}}AFX_DATA_MAP
}


BEGIN_MESSAGE_MAP(CProcFilterDlg, CDialog)
	//{{AFX_MSG_MAP(CProcFilterDlg)
	ON_WM_DRAWITEM()
	ON_BN_CLICKED(IDC_REFRESH, OnRefresh)
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// CProcFilterDlg message handlers

void CProcFilterDlg::OnDrawItem(int nIDCtl, LPDRAWITEMSTRUCT lpDrawItemStruct) 
{
	// TODO: Add your message handler code here and/or call default
	
	CDialog::OnDrawItem(nIDCtl, lpDrawItemStruct);
}

BOOL CProcFilterDlg::OnInitDialog() 
{
	CDialog::OnInitDialog();

	m_ListCtrl.SetFullRowSel(FALSE);
	
	// TODO: Add extra initialization here
	m_StateImageList.Create(IDB_STATEICONS, 16, 1, RGB(255, 0, 0));


	m_ListCtrl.SetImageList(&m_StateImageList, LVSIL_STATE);

// insert columns

	int i;
	LV_COLUMN lvc;

	lvc.mask = LVCF_FMT | LVCF_WIDTH | LVCF_TEXT | LVCF_SUBITEM;

	RECT rc;
	m_ListCtrl.GetClientRect(&rc);
	
	int ListCtrl_cx = rc.right - rc.left - GetSystemMetrics(SM_CXVSCROLL);

	for(i = 0; i<NUM_COLUMNS; i++)
	{
		lvc.iSubItem = i;
		lvc.pszText = gProFilterColumnLabel[i];
		lvc.cx = (ListCtrl_cx * gProFilterColumnWidth[i]) / 100;
		lvc.fmt = gProFilterColumnFormat[i];
		m_ListCtrl.InsertColumn(i,&lvc);
	}

// insert items
	BuildProcessList();


	return TRUE;  // return TRUE unless you set the focus to a control
	              // EXCEPTION: OCX Property Pages should return FALSE
}

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

BOOL AdjustPrivileges(LPCTSTR privilege)
{
	HANDLE hToken;              // handle to process token 
	TOKEN_PRIVILEGES tkp;        // ptr. to token structure 
 
	BOOL fResult;               
 
 
	OpenProcessToken(GetCurrentProcess(), 
        TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken) ;
 
	// Get the LUID for debug privilege. 
 
	LookupPrivilegeValue(NULL, privilege, 
			&tkp.Privileges[0].Luid); 
 
	tkp.PrivilegeCount = 1;  // one privilege to set    
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; 
 
 
	fResult = AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, 
		(PTOKEN_PRIVILEGES) NULL, 0); 

    CloseHandle(hToken);
    return fResult;
}



int CProcFilterDlg::BuildProcessList()
{
	DWORD ProcessesId[1024];
	DWORD done;
	HINSTANCE hPSAPI=LoadLibrary("PSAPI.DLL");

	pEnumProcesses=(EnumProcessesProc)GetProcAddress(hPSAPI,"EnumProcesses");
	pEnumProcessModules=(EnumProcessModulesProc)GetProcAddress(hPSAPI,"EnumProcessModules");
	pGetModuleBaseNameA=(GetModuleBaseNameAProc)GetProcAddress(hPSAPI,"GetModuleBaseNameA");

	if(!pEnumProcesses(ProcessesId,sizeof(ProcessesId),&done))
	{
		return (DWORD)-1;
	}

    AdjustPrivileges(SE_DEBUG_NAME);

	int cProcesses = done / sizeof(DWORD);
	int i;
	char szProcessName[255];
	DWORD pid;
	HANDLE hProcess;
	HMODULE hMod;

	for(i=0;i<cProcesses;i++)
	{
		strcpy(szProcessName,"unknown");
		pid = ProcessesId[i];
	
		hProcess = OpenProcess(PROCESS_QUERY_INFORMATION|
								PROCESS_VM_READ,
								FALSE,pid);
		if(!hProcess){
			continue;
		}
		
		
		if(pEnumProcessModules(hProcess,&hMod,sizeof(hMod),&done) )
		{
			pGetModuleBaseNameA(hProcess,hMod,szProcessName,sizeof(szProcessName));
		}

		CloseHandle(hProcess);

		Pid2ImageName[pid] = szProcessName;
	}

	Pid2ImageName[8] = "System";

	LV_ITEM lvi;
	CString ProcessImageName;
	POSITION iterator = Pid2ImageName.GetStartPosition();
	CString str_pid;

	m_ListCtrl.ModifyStyle(LVS_SORTASCENDING,0);
	
	for(i = 0; i < Pid2ImageName.GetCount(); i++)
	{
		Pid2ImageName.GetNextAssoc( iterator,pid,ProcessImageName );

		lvi.mask = LVIF_TEXT | LVIF_IMAGE | LVIF_STATE | LVIF_PARAM;
		lvi.iItem = i;
		lvi.iSubItem = 0;
		lvi.pszText = (LPSTR)(LPCTSTR)ProcessImageName;
		lvi.iImage = i;
		lvi.stateMask = LVIS_STATEIMAGEMASK;
		lvi.state = INDEXTOSTATEIMAGEMASK(1);
		lvi.lParam = pid;

		m_ListCtrl.InsertItem(&lvi);
		str_pid.Format("%X",pid);
		m_ListCtrl.SetItemText(i,1,str_pid);

		for(int j=0;j < PidSelected.GetSize();j++)
		{
			if( pid == PidSelected[j] )
				m_ListCtrl.CheckItem(i);
		}
	}
	m_ListCtrl.ModifyStyle(0,LVS_SORTASCENDING);

	m_ListCtrl.SortItems( CompareFunc, 0);

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

void CProcFilterDlg::OnRefresh() 
{
	// TODO: Add your control notification handler code here
	m_ListCtrl.DeleteAllItems();
	Pid2ImageName.RemoveAll();
		
	BuildProcessList();
	
}

void CProcFilterDlg::OnOK() 
{
	// TODO: Add extra validation here
	PidSelected.RemoveAll();

	int i = -1;
	while( (i = m_ListCtrl.GetNextItem(i,LVNI_ALL)) != -1)
	{
		if(m_ListCtrl.IsChecked(i)){
	
			PidSelected.Add((WORD)m_ListCtrl.GetItemData(i));
		}
	}
	
	CDialog::OnOK();
}

int CProcFilterDlg::GetResults(SERVICEPTR2NAME& Pid2ImageName, CWordArray & PidSelected)
{
	Pid2ImageName.RemoveAll();
	POSITION iterator = this->Pid2ImageName.GetStartPosition();
	DWORD pid;
	CString ProcessImageName;
	for(int i = 0; i < this->Pid2ImageName.GetCount(); i++)
	{
		this->Pid2ImageName.GetNextAssoc( iterator,pid,ProcessImageName );
		Pid2ImageName[pid] = ProcessImageName;
	}
	PidSelected.Copy(this->PidSelected);
	return 0;
}

int CProcFilterDlg::SetSelected(CWordArray & PidSelected)
{
	this->PidSelected.Copy(PidSelected);
	return 0;

}