#include "pch.h"

#include <windows.h>

#include <psapi.h>

#include <tlhelp32.h>

#include <locale>

#include <codecvt>

#include <vector>

#include <map>

#include "..\driver\public.h"

#include "findPFNDatabase.h"

class statistics {

public:

statistics()

{

scannedPages = ignoredPagesNX = scannedProcesses = modifiedPages = 0;

}

unsigned int scannedPages;

unsigned int ignoredPagesNX;

unsigned int scannedProcesses;

unsigned int modifiedPages;

};

class modifiedPage {

public:

modifiedPage(DWORD newProcessID, wchar_t* newModuleName, void* newPageBase, BYTE newSectionName[8], unsigned long long newSectionOffset);

unsigned long processID;

std::wstring moduleName;

unsigned long long pageBase;

std::wstring sectionName;

unsigned long long sectionOffset;

};

modifiedPage::modifiedPage(DWORD newProcessID, wchar_t* newModuleName, void* newPageBase, BYTE newSectionName[8], unsigned long long newSectionOffset)

: processID(newProcessID), moduleName(newModuleName), pageBase((unsigned long long)newPageBase), sectionName(L""), sectionOffset(newSectionOffset)

{

std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;

std::wstring secName = std::wstring(converter.from_bytes((char*)newSectionName, (char*)&newSectionName[7]));

size_t nullPos = secName.find(L'\0');

if (nullPos != secName.npos)

secName = secName.substr(0, nullPos);

sectionName.append(secName);

}

BOOL EnableDebugPrivilege(BOOL bEnable)

{

HANDLE hToken = nullptr;

LUID luid;

if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken)) return FALSE;

if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) return FALSE;

TOKEN_PRIVILEGES tokenPriv;

tokenPriv.PrivilegeCount = 1;

tokenPriv.Privileges[0].Luid = luid;

tokenPriv.Privileges[0].Attributes = bEnable ? SE_PRIVILEGE_ENABLED : 0;

if (!AdjustTokenPrivileges(hToken, FALSE, &tokenPriv, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) return FALSE;

return TRUE;

}

int scanProcess(HANDLE driverHnd, DWORD targetPID, HANDLE toScanHandle, std::vector<modifiedPage> *resultsOut, statistics* stats)

{

DWORD cbNeeded;

int s = EnumProcessModules(toScanHandle, NULL, 0, &cbNeeded);

if (s == 0)

{

printf("Couldn't call EnumProcessModules to get buffer size, gle %d\n", GetLastError());

return -1;

}

HMODULE* moduleList = (HMODULE*)malloc(cbNeeded);

memset(moduleList, 0, cbNeeded);

s = EnumProcessModules(toScanHandle, moduleList, cbNeeded, &cbNeeded);

if (s == 0)

{

// This'll happen sometimes if there's a module loaded between our calls.

// TODO: we can retry in this case.

printf("Couldn't call EnumProcessModules to get modules, gle %d.\n", GetLastError());

return -1;

}

for (HMODULE* thisModPtr = &moduleList[0]; thisModPtr < &moduleList[cbNeeded / sizeof(HMODULE)]; thisModPtr++)

{

HMODULE thisModule = *thisModPtr;

TCHAR szModName[MAX_PATH];

memset(szModName, 0, MAX_PATH * sizeof(TCHAR));

if (GetModuleFileNameEx(toScanHandle, thisModule, szModName, sizeof(szModName) / sizeof(TCHAR)) == 0)

{

printf("GetModuleFileNameEx failed, GLE %d\n", GetLastError());

continue;

}

IMAGE_DOS_HEADER mz;

SIZE_T bytesRead;

if (!ReadProcessMemory(toScanHandle, thisModule, &mz, sizeof(IMAGE_DOS_HEADER), &bytesRead))

{

printf("Can't read module MZ header\n");

return -1;

}

if (mz.e_magic != IMAGE_DOS_SIGNATURE)

{

printf("MZ header not found\n");

continue;

}

IMAGE_NT_HEADERS pe;

unsigned long long peAddress = (((unsigned long long)thisModule) + mz.e_lfanew);

if (!ReadProcessMemory(toScanHandle, (void*)peAddress, &pe, sizeof(IMAGE_NT_HEADERS), &bytesRead))

{

printf("Can't read module PE header\n");

return -1;

}

if (pe.Signature != IMAGE_NT_SIGNATURE)

{

printf("PE header not found\n");

continue;

}

IMAGE_SECTION_HEADER* sect;

unsigned long long firstSectionAddress = peAddress + FIELD_OFFSET(IMAGE_NT_HEADERS, OptionalHeader) + sizeof(IMAGE_OPTIONAL_HEADER);

sect = (IMAGE_SECTION_HEADER*)malloc(sizeof(IMAGE_SECTION_HEADER) * pe.FileHeader.NumberOfSections);

if (!ReadProcessMemory(toScanHandle, (LPCVOID)(firstSectionAddress), sect, sizeof(IMAGE_SECTION_HEADER) * pe.FileHeader.NumberOfSections, &bytesRead))

{

printf("Can't read first section of module\n");

return -1;

}

for (unsigned long sectionIndex = 0; sectionIndex < pe.FileHeader.NumberOfSections; sectionIndex++)

{

IMAGE_SECTION_HEADER* thisSection = &sect[sectionIndex];

unsigned long long relocatedSectionBase = thisSection->VirtualAddress + (unsigned long long)thisModule;

// We are interested only in executable sections.

// TODO: check that discardable pages are zero'ed out?

// TODO: check that non-executable pages haven't been made executable?

if ((thisSection->Characteristics & IMAGE_SCN_MEM_EXECUTE) == 0)

{

stats->ignoredPagesNX += (thisSection->SizeOfRawData / 0x1000);

// printf("%ls!%s (at %p) is not executable, skipping\n", szModName, thisSection->Name, relocatedSectionBase);

continue;

}

//printf("scanning %ls!%s (at %p), size 0x%08lx\n", szModName, thisSection->Name, relocatedSectionBase, thisSection->SizeOfRawData);

int dirtyPages = 0;

int errorPages = 0;

getPageInfoRequest req;

req.pageToCheck = relocatedSectionBase;

// Work out how many pages we will check

req.numberOfPagesToCheck = thisSection->Misc.VirtualSize / 0x1000;

if (thisSection->Misc.VirtualSize % 0x1000 != 0)

req.numberOfPagesToCheck++;

req.targetPID = targetPID;

getPageInfoResponse* resp = (getPageInfoResponse*)malloc(sizeof(getPageInfoRequest) * req.numberOfPagesToCheck);

memset(resp, 0x00, sizeof(getPageInfoResponse) * req.numberOfPagesToCheck);

DWORD bytesRet;

s = DeviceIoControl(driverHnd, IOCTL_DRIVER_QUERY_VA, &req, sizeof(req), resp, sizeof(getPageInfoResponse) * req.numberOfPagesToCheck, &bytesRet, NULL);

if (s == 0)

{

errorPages++;

printf("DeviceIoControl failed, GLE %d\n", GetLastError());

return -1;

}

stats->scannedPages += req.numberOfPagesToCheck;

for (unsigned int n = 0; n < req.numberOfPagesToCheck; n++)

{

unsigned long long pageAddress = relocatedSectionBase + (n * 0x1000);

if (!resp[n].isValid)

{

printf("Page at 0x%016llx (%ls!%s) not valid (maybe it's paged out?) 0x%08lx\n", pageAddress, szModName, thisSection->Name, thisSection->Characteristics);

errorPages++;

continue;

}

if (resp[n].isDirty)

{

dirtyPages++;

resultsOut->push_back(modifiedPage(targetPID, (wchar_t*)szModName, (void*)pageAddress, thisSection->Name, (pageAddress - relocatedSectionBase)));

stats->modifiedPages++;

}

}

// if (dirtyPages == 0)

// printf("Module %ls: OK\n", szModName);

// else

// printf("Module %ls: detected %d dirty pages!\n", szModName, dirtyPages);

}

}

return 0;

}

int setPFNDatabase(HANDLE driverHnd, unsigned long long PFNDatabaseStart)

{

setPFNDatabaseRequest req;

req.offsetToMmPfnDatabaseInNtDllFromExAllocatePoolWithTag = PFNDatabaseStart;

DWORD bytesRet;

int s = DeviceIoControl(driverHnd, IOCTL_DRIVER_SET_PFN_DATABASE, &req, sizeof(req), NULL, 0, &bytesRet, NULL);

if (s == 0)

{

printf("Failed to set PFN database to 0x%016llx: GLE %d\n", PFNDatabaseStart, GetLastError());

return -1;

}

return 0;

}

int main()

{

EnableDebugPrivilege(TRUE);

HANDLE driverHnd = CreateFile(L"\\\\.\\cowspot", GENERIC_ALL, 0, NULL, OPEN_EXISTING, 0, NULL);

if (driverHnd == INVALID_HANDLE_VALUE)

{

printf("Couldn't open driver device '%ls', gle %d\n", DOS_DEVICE_NAME, GetLastError());

return -1;

}

if (setPFNDatabase(driverHnd, findPFNDatabase()) != 0)

return -1;

HANDLE snapshotHnd = CreateToolhelp32Snapshot(TH32CS_SNAPALL, 0);

if (snapshotHnd == INVALID_HANDLE_VALUE)

{

printf("CreateToolhelp32Snapshot failed, GLE %d\n", GetLastError());

return -1;

}

PROCESSENTRY32 proc;

memset(&proc, 0, sizeof(PROCESSENTRY32));

proc.dwSize = sizeof(PROCESSENTRY32);

if (!Process32First(snapshotHnd, &proc))

{

printf("Process32First failed, GLE %d\n", GetLastError());

return -1;

}

statistics stat;

std::vector<modifiedPage> results;

unsigned long start = GetTickCount();

while (Process32Next(snapshotHnd, &proc))

{

HANDLE toScanHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, proc.th32ProcessID);

if (toScanHandle == NULL)

{

printf("Couldn't open target process with PID %d ('%ls'), gle %d\n", proc.th32ProcessID, proc.szExeFile, GetLastError());

continue;

}

stat.scannedProcesses++;

if (scanProcess(driverHnd, proc.th32ProcessID, toScanHandle, &results, &stat) != 0)

printf("Failed to scan process '%ls'\n", proc.szExeFile);

// else

// printf("Scanned process '%ls'\n", proc.szExeFile);

CloseHandle(toScanHandle);

}

CloseHandle(snapshotHnd);

unsigned long end = GetTickCount();

printf("Scan took %dms\n", (end - start));

// Print some stats and the results. Each page is 4KB, so 256 pages is 1MB.

printf("Scanned %d pages (%d MB), ignored %d NX pages (%d MB), totalling %d pages (%d MB). Found %d modified pages (%.02f MB).\n",

stat.scannedPages, stat.scannedPages / 256,

stat.ignoredPagesNX, stat.scannedPages / 256,

stat.ignoredPagesNX + stat.scannedPages, (stat.ignoredPagesNX + stat.scannedPages) / 256,

stat.modifiedPages, ((float)stat.modifiedPages) / 256);

for (unsigned int n = 0; n < results.size(); n++)

{

modifiedPage thisModifiedPage = results[n];

printf("PID %04d module '%ls', section %S, offset 0x%08llx\n", thisModifiedPage.processID, thisModifiedPage.moduleName.c_str(), thisModifiedPage.sectionName.c_str(), thisModifiedPage.sectionOffset);

/*

HANDLE toScanHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, thisModifiedPage.processID);

if (toScanHandle == NULL)

{

printf("Couldn't open target process with PID %d ('%ls'), gle %d\n", proc.th32ProcessID, proc.szExeFile, GetLastError());

continue;

}

SIZE_T bytesRead;

unsigned char* pageContents[0x2000];

memset(pageContents, 0, 0x2000);

if (!ReadProcessMemory(toScanHandle, (LPCVOID)thisModifiedPage.pageBase, pageContents, 0x2000, &bytesRead))

{

printf("ReadProcessMemory failed\n");

continue;

}

for (unsigned int n = 0; n < 0x2001; n++)

{

printf("0x%02hhx ", (unsigned)pageContents[n]);

if (n % 0x10 == 0)

printf("\n0x%08lx: ", n);

}

CloseHandle(toScanHandle);*/

}

return 0;

}