Driver.cpp

#include "Driver.h"

#define SystemModuleInformation 11

typedef struct _RTL_PROCESS_MODULE_INFORMATION
{
    HANDLE Section;
    PVOID MappedBase;
    DWORD ImageBase;
    ULONG ImageSize;
    ULONG Flags;
    USHORT LoadOrderIndex;
    USHORT InitOrderIndex;
    USHORT LoadCount;
    USHORT OffsetToFileName;
    UCHAR FullPathName[256];
} RTL_PROCESS_MODULE_INFORMATION, *PRTL_PROCESS_MODULE_INFORMATION;
 
typedef struct _RTL_PROCESS_MODULES
{
    ULONG NumberOfModules;
    RTL_PROCESS_MODULE_INFORMATION Modules[1];
} RTL_PROCESS_MODULES, *PRTL_PROCESS_MODULES;


Driver::Driver(
	const wstring& drvName, 
	const wstring& alternativeDrvName,
	const IoctlNamesMap* ioctlNames
) 
	: driverName(drvName), lastError(ERROR_SUCCESS), ioctlNamesMap(ioctlNames)
{
	shared_ptr<HANDLE> drvHandle(new HANDLE(
		CreateFileW(
			wstring(wstring(L"\\\\.\\") + driverName).c_str(), 
			GENERIC_READ | GENERIC_WRITE,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED,
			NULL
	)),
	[](const HANDLE *ptr) {
		if( *ptr != (HANDLE)INVALID_HANDLE_VALUE && ptr != nullptr) {
			//wcout << L"[.] Closing driver's handle." << endl;
			CloseHandle(*ptr);
		}
		delete ptr;
	});
	
	if( *drvHandle == static_cast<HANDLE>(INVALID_HANDLE_VALUE)) {
		lastError = GetLastError();
		wcerr << L"[!] Could not open \\Device\\" << driverName << L" driver's handle. Error: 0x" 
			<< lastError << endl;
		return;
	}
	
	driverHandle = move(drvHandle);
	wstring moduleName = (alternativeDrvName.size() > 0)? alternativeDrvName : drvName;
	
	driverImageBase = Driver::GetModuleImageBase(moduleName);
	if(!driverImageBase) {
		lastError = GetLastError();
	}
}


Driver::~Driver() {
	
}


DWORD 
Driver::GetModuleImageBase(
	const wstring &imageNameParam
) {
	static const wstring ending(L".sys");	

	wstring imageName(imageNameParam);
	transform(imageName.begin(), imageName.end(), imageName.begin(), ::tolower);
	
	if(!equal(ending.rbegin(), ending.rend(), imageName.rbegin())) {
		imageName += ending;
	}
	
	shared_ptr<RTL_PROCESS_MODULES> modulesInfo(
		new RTL_PROCESS_MODULES[Max_Number_Of_Modules]
	);
	
	if(!modulesInfo.get()) {
		wcerr << L"[!] Could not allocate memory for modulesInfo!" << endl;
		return 0;
	}
	
	NTSTATUS status = NtQuerySystemInformation(
		(SYSTEM_INFORMATION_CLASS)SystemModuleInformation, 
		modulesInfo.get(), 
		sizeof(RTL_PROCESS_MODULES) * Max_Number_Of_Modules, 
		NULL
	);
	
	if(!NT_SUCCESS(status)) {
		wcerr << L"[!] Could not obtain list of loaded driver modules! Status: 0x" 
			<< hex << setfill(L'0') << setw(8) << status << endl;
		return 0;
	}
	
	for(size_t i = 0; i < modulesInfo->NumberOfModules; i++ ) {
		
		const DWORD imageBase = modulesInfo->Modules[i].ImageBase;
		const char *moduleImageNamePointer = reinterpret_cast<const char*>(
			modulesInfo->Modules[i].FullPathName 
			+ modulesInfo->Modules[i].OffsetToFileName);
			
		const string moduleImageName(moduleImageNamePointer);

		wstring moduleImageNameWide(moduleImageName.begin(), moduleImageName.end());
		transform(
			moduleImageNameWide.begin(), 
			moduleImageNameWide.end(), 
			moduleImageNameWide.begin(), 
			::tolower);
			
		if(imageName == moduleImageNameWide) {
			wcout << L"[+] Found " << imageName << L" driver's base: 0x" 
				<< hex << setfill(L'0') << setw(8) << imageBase << endl;
			return imageBase;
		}
	}
	
	wcerr << L"[!] Could not find the " << imageName 
		<< L" module among the loaded modules." << endl;
	
	return 0;
}


tuple<ULONG, DWORD, wstring> 
Driver::GetKernelModuleInfos() {
	
	shared_ptr<RTL_PROCESS_MODULES> modulesInfo(
		new RTL_PROCESS_MODULES[Max_Number_Of_Modules]
	);
	
	if(!modulesInfo.get()) {
		wcerr << L"[!] Could not allocate memory for modulesInfo!" << endl;
		return make_tuple(0, 0, L"");
	}
	
	NTSTATUS status = NtQuerySystemInformation(
		(SYSTEM_INFORMATION_CLASS)SystemModuleInformation, 
		modulesInfo.get(), 
		sizeof(RTL_PROCESS_MODULES) * Max_Number_Of_Modules, 
		NULL
	);
	
	if(!NT_SUCCESS(status)) {
		wcerr << L"[!] Could not obtain list of loaded driver modules! Status: 0x" 
			<< hex << setfill(L'0') << setw(8) << status << endl;
		return make_tuple(0, 0, L"");
	}
	
	const char *moduleImageNamePointer = reinterpret_cast<const char*>(
		modulesInfo->Modules[0].FullPathName 
		+ modulesInfo->Modules[0].OffsetToFileName);
			
	const string moduleImageName(moduleImageNamePointer);

	wstring moduleImageNameWide(moduleImageName.begin(), moduleImageName.end());
	transform(
		moduleImageNameWide.begin(), 
		moduleImageNameWide.end(), 
		moduleImageNameWide.begin(), 
		::tolower);
	
	return make_tuple(
		modulesInfo->Modules[0].ImageBase,
		modulesInfo->Modules[0].ImageSize,
		moduleImageNameWide
	);
}


bool 
Driver::SendIOCTL(
	DWORD ioctlCode, 
	LPVOID inputBuffer, 
	DWORD inputSize, 
	LPVOID outputBuffer, 
	DWORD outSize, 
	LPDWORD writtenBytes,
	BOOL quiet
) {
	if(!quiet) {
		auto ioctlCodeString = [=]() -> const wchar_t* {
			wstringstream wss;
			wss << L"0x" << hex << setw(8) << setfill(L'0') << ioctlCode;
			return wss.str().c_str();
		};
		
		if (ioctlNamesMap) {
			auto ioctl = ioctlNamesMap->find(ioctlCode);
			if ( ioctl != ioctlNamesMap->end()) {
				wcout << endl << L"[+] Issuing IOCTL: " << ioctl->second << endl;
			} else {
				wcout << endl << L"[+] Issuing IOCTL: " << ioctlCodeString() << endl;
			}
		} else {
			wcout << endl << L"[+] Issuing IOCTL: " << ioctlCodeString() << endl;
		}	
	}
	
	DWORD written = 0;
	bool ownAllocation = false;
	LPDWORD writtenPtr = (writtenBytes != nullptr)? writtenBytes : &written;
	
	if (inputBuffer == nullptr) {
		ownAllocation = true;
		inputBuffer = VirtualAlloc ((LPVOID)0,
									inputSize,
									MEM_COMMIT | MEM_RESERVE,
									PAGE_EXECUTE_READWRITE	
		);
		
		if (inputBuffer == nullptr) {
			wcout << L"[!] Could not allocate memory for " << inputSize << L" bytes! Error: " 
					<< GetLastError() << endl;
			throw runtime_error("allocation error");
		}
	}
	
	if(!quiet) {
		wcout << hex << setw(8) << setfill(L'0')
				<< L"[.] Input buffer: 0x" << inputBuffer
				<< L", size: " << dec << inputSize << L" bytes. Output: 0x" 
				<< hex << outputBuffer << endl;
	}
			
	SetLastError(0);
	BOOL ret = DeviceIoControl(*driverHandle, 
								ioctlCode,
								inputBuffer,
								inputSize,
								outputBuffer,
								outSize,
								writtenPtr,
								NULL
	);
	
	const DWORD lastError = GetLastError();
	
	if(!quiet) {
		wcout << hex << setw(8) << setfill(L'0')
				<< L"[>] After sending IOCTL. ret: " << ((ret)? L"TRUE" : L"FALSE")
				<< L", last error: " << dec << lastError << L", written bytes: " 
				<< dec << *writtenPtr << endl;
			
		if (lastError != ERROR_SUCCESS) {
			wcout << L"\tError message: " << getErrorString(lastError) << endl;
		}
		
		wcout << endl;
	}
	
	if (ownAllocation) {
		VirtualFree(inputBuffer, inputSize, MEM_DECOMMIT);
	}
	
	return ret;
}