BeneficialCode-WinArk/WinArk/ProcessModuleTable.cpp

#include "stdafx.h"
#include "ProcessModuleTable.h"
#include <algorithm>
#include "SortHelper.h"
#include "resource.h"

LRESULT CProcessModuleTable::OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return 0;
}
LRESULT CProcessModuleTable::OnDestroy(UINT uMsg, WPARAM wParam, LPARAM lparam, BOOL& /*bHandled*/) {
	return 0;
}
LRESULT CProcessModuleTable::OnPaint(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	PaintTable(m_hWnd);
	return 0;
}

LRESULT CProcessModuleTable::OnHScroll(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	Tablefunction(m_hWnd, uMsg, wParam, lParam);
	return 0;
}

LRESULT CProcessModuleTable::OnVScroll(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	Tablefunction(m_hWnd, uMsg, wParam, lParam);
	return 0;
}

LRESULT CProcessModuleTable::OnUserVabs(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	Tablefunction(m_hWnd, uMsg, wParam, lParam);
	return 0;
}

LRESULT CProcessModuleTable::OnUserVrel(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	Tablefunction(m_hWnd, uMsg, wParam, lParam);
	return 0;
}

LRESULT CProcessModuleTable::OnUserChgs(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	Tablefunction(m_hWnd, uMsg, wParam, lParam);
	return 0;
}

LRESULT CProcessModuleTable::OnMouseWheel(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	int zDelta = GET_WHEEL_DELTA_WPARAM(wParam);
	return Tablefunction(m_hWnd, WM_VSCROLL, zDelta >= 0 ? 0 : 1, wParam);
}
LRESULT CProcessModuleTable::OnMouseMove(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnLBtnDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnLBtnUp(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnRBtnDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnUserSts(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnWindowPosChanged(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnKeyDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}
LRESULT CProcessModuleTable::OnSysKeyDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& /*bHandled*/) {
	return Tablefunction(m_hWnd, uMsg, wParam, lParam);
}

LRESULT CProcessModuleTable::OnGetDlgCode(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM /*lParam*/, BOOL& /*bHandled*/) {
	return DLGC_WANTARROWS;
}


CProcessModuleTable::CProcessModuleTable(BarInfo& bars, TableInfo& table,DWORD pid,HWND hwnd)
	: CTable(bars, table),m_Tracker(pid),_hDlg(hwnd) {
	SetTableWindowInfo(bars.nbar);
	Refresh();
}

CProcessModuleTable::CProcessModuleTable(BarInfo& bars, TableInfo& table, HANDLE hProcess)
	: CTable(bars, table), m_Tracker(hProcess) {
	SetTableWindowInfo(bars.nbar);
	//DoRefresh();
}

CString CProcessModuleTable::CharacteristicsToString(WinSys::DllCharacteristics ch) {
	using namespace WinSys;

	CString result;
	if ((ch & DllCharacteristics::HighEntropyVA) != DllCharacteristics::None)
		result += L"High Entropy VA, ";
	if ((ch & DllCharacteristics::DynamicBase) != DllCharacteristics::None)
		result += L"Dynamic Base, ";
	if ((ch & DllCharacteristics::ForceIntegrity) != DllCharacteristics::None)
		result += L"Force Integrity, ";
	if ((ch & DllCharacteristics::NxCompat) != DllCharacteristics::None)
		result += L"NX Compat, ";
	if ((ch & DllCharacteristics::NoIsolation) != DllCharacteristics::None)
		result += L"No Isolation, ";
	if ((ch & DllCharacteristics::NoSEH) != DllCharacteristics::None)
		result += L"No SEH, ";
	if ((ch & DllCharacteristics::NoBind) != DllCharacteristics::None)
		result += L"No Bind, ";
	if ((ch & DllCharacteristics::AppContainer) != DllCharacteristics::None)
		result += L"App Container, ";
	if ((ch & DllCharacteristics::WDMDriver) != DllCharacteristics::None)
		result += L"WDM Driver, ";
	if ((ch & DllCharacteristics::ControlFlowGuard) != DllCharacteristics::None)
		result += L"CFG, ";
	if ((ch & DllCharacteristics::TerminalServerAware) != DllCharacteristics::None)
		result += L"TS Aware, ";

	if (!result.IsEmpty())
		result = result.Left(result.GetLength() - 2);
	return result;
}

int CProcessModuleTable::ParseTableEntry(CString& s, char& mask, int& select, std::shared_ptr<WinSys::ModuleInfo>& info, int column) {
	switch (column) {
	case 0:
		s = info->Name.empty() ? L"<Pagefile Backed>" : info->Name.c_str();
		break;
	case 1:
		s = info->Type == WinSys::MapType::Image ? L"Image" : L"Data";
		break;
	case 2:
		s.Format(L"0x%08X", info->ModuleSize);
		break;
	case 3:
		s.Format(L"0x%p", info->Base);
		break;
	case 4:
		if (info->Type == WinSys::MapType::Image)
			s.Format(L"0x%p", info->ImageBase);
		else
			s = L"N/A";
		break;
	case 5:
		s = info->Type == WinSys::MapType::Image ? CharacteristicsToString(info->Characteristics) : CString();
		break;
	case 6:
		s = info->Path.c_str();
		break;
	}
	return s.GetLength();
}

CProcessModuleTable::ModuleInfoEx& CProcessModuleTable::GetModuleEx(WinSys::ModuleInfo* mi) {
	auto it = m_ModulesEx.find(mi);
	if (it != m_ModulesEx.end())
		return it->second;

	ModuleInfoEx mx;
	m_ModulesEx.insert({ mi,mx });
	return GetModuleEx(mi);
}


void CProcessModuleTable::Refresh() {
	if (!m_Tracker.IsRunning()) {
		AtlMessageBox(*this, L"Process has terminated", IDS_TITLE, MB_OK|MB_ICONWARNING);
		EndDialog(_hDlg, 0);
		return;
	}
	auto first = m_Table.data.info.empty();
	m_Tracker.EnumModules();
	if (first) {
		m_Table.data.info = m_Tracker.GetModules();
	}
	else {
		auto count = static_cast<int>(m_Table.data.info.size());
		for (int i = 0; i < count; i++) {
			auto& mi = m_Table.data.info[i];
			auto& mx = GetModuleEx(mi.get());
			if (mx.IsUnloaded && ::GetTickCount64() > mx.TargetTime) {
				ATLTRACE(L"Module unload end: %u %s (0x%p)\n", ::GetTickCount(), mi->Name.c_str(), mi->Base);
				m_ModulesEx.erase(mi.get());
				m_Table.data.info.erase(m_Table.data.info.begin() + i);
				i--;
				count--;
			}
			else if (mx.IsNew && !mx.IsUnloaded && ::GetTickCount64() > mx.TargetTime) {
				mx.IsNew = false;
			}
		}

		for (auto& mi : m_Tracker.GetNewModules()) {
			m_Table.data.info.push_back(mi);
			auto& mx = GetModuleEx(mi.get());
			mx.IsNew = true;
			mx.TargetTime = ::GetTickCount64() + 2000;
		}
		for (auto& mi : m_Tracker.GetUnloadedModules()) {
			ATLTRACE(L"Module unload start: %u %s (0x%p)\n", ::GetTickCount(), mi->Name.c_str(), mi->Base);
			auto& mx = GetModuleEx(mi.get());
			ATLASSERT(!mx.IsUnloaded);
			mx.IsNew = false;
			mx.IsUnloaded = true;
			mx.TargetTime = ::GetTickCount64() + 2000;
		}
	}

	auto count = static_cast<int>(m_Table.data.info.size());
	m_Table.data.n = count;

	return;
}