/*
Enumerates all loaded drivers on the system
Author: illustrissimus

To build and run with cmd.exe:
> cd "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build"
> vcvars64.bat
> cd "C:\source"
> cl.exe /nologo /Tp enum_drivers.cpp /link /out:enum_drivers.exe
> enum_drivers.exe
*/

#include <windows.h>
#include <tchar.h>
#include <psapi.h>
#include <stdio.h>

typedef struct {
	LPVOID *drivers;
	DWORD driver_count;
	DWORD current;
} driver_iterator;

typedef struct {
	TCHAR* driver_name;
	LPVOID driver_address;
} driver;

bool create_driver_iterator(driver_iterator* di) {
	LPVOID dummy[1];
	DWORD dummy_bytes_needed, drivers_bytes_needed;
	/* first time we fail on purpose to learn the required buffer size */
	if (!EnumDeviceDrivers(dummy, sizeof(dummy), &dummy_bytes_needed))
		return false;
	if (dummy_bytes_needed < 1)
		return false;
	di->drivers = (LPVOID*)VirtualAlloc(NULL,
		dummy_bytes_needed,
		MEM_COMMIT | MEM_RESERVE,
		PAGE_READWRITE);
	if (di->drivers == NULL)
		return false;
	if (!EnumDeviceDrivers(di->drivers, dummy_bytes_needed, &drivers_bytes_needed)) {
		VirtualFree(di->drivers, 0, MEM_RELEASE);
		return false;
	}
	if (dummy_bytes_needed < drivers_bytes_needed) {
		/* TOCTOU: some driver has been added between calls so let's bail */
		VirtualFree(di->drivers, 0, MEM_RELEASE);
		return false;
	}
	di->driver_count = drivers_bytes_needed / sizeof(di->drivers[0]);
	di->current = 0;
	return true;
}

TCHAR* get_driver_name(LPVOID addr) {
	/* the API does not help us determine the required buffer size, so we guess */
	DWORD buffer_size = 16, copied_characters;
	TCHAR* buffer;
	do {
		buffer = (TCHAR*)VirtualAlloc(NULL,
			buffer_size * sizeof(TCHAR),
			MEM_COMMIT | MEM_RESERVE,
			PAGE_READWRITE);
		if (buffer == NULL)
			return NULL;
		copied_characters = GetDeviceDriverBaseName(addr, buffer, buffer_size - 1);
		if (copied_characters == 0) {
			VirtualFree(buffer, 0, MEM_RELEASE);
			return NULL;
		}
		if (copied_characters != buffer_size - 1) {
			break;
		}
		buffer_size *= 2;
		VirtualFree(buffer, 0, MEM_RELEASE);
	} while (true);
	return buffer;
}

driver* driver_iterator_next(driver_iterator* di) {
	if (di->current >= di->driver_count)
		return NULL;
	driver* d = (driver*)VirtualAlloc(NULL,
		sizeof(driver),
		MEM_COMMIT | MEM_RESERVE,
		PAGE_READWRITE);
	if (d == NULL)
		return NULL;
	
	d->driver_name = get_driver_name(di->drivers[di->current]);
	d->driver_address = di->drivers[di->current];
	di->current += 1;
	return d;
}

bool driver_iterator_has_next(driver_iterator* di) {
	return di->current < di->driver_count;
}

void release_driver_iterator(driver_iterator* di) {
	VirtualFree(di->drivers, 0, MEM_RELEASE);
	di->driver_count = 0;
	di->current = 0;
}

void enumerate_drivers() {
	driver_iterator di;
	if (!create_driver_iterator(&di))
		return;
	while (driver_iterator_has_next(&di)) {
		driver *d = driver_iterator_next(&di);
		if (d == NULL) {
			release_driver_iterator(&di);
			return;
		}
		if (d->driver_name != NULL) {
			_tprintf(TEXT("0x%p %s\n"), d->driver_address, d->driver_name);
			VirtualFree(d->driver_name, 0, MEM_RELEASE);
		}
		else {
			_tprintf(TEXT("0x%p %s\n"), d->driver_address, TEXT("?"));
		}
		VirtualFree(d, 0, MEM_RELEASE);
	}
	release_driver_iterator(&di);
}

int _tmain(int argc, TCHAR** argv) {
	enumerate_drivers();
	return 0;
}