#include <ntddk.h>          
#include <string.h>

#define NT_DEVICE_NAME      L"\\Device\\pcidisp"
#define DOS_DEVICE_NAME     L"\\DosDevices\\pcidisp"

#define MAX_PCI_BUS 255
#define MAX_PCI_SLOT 32

static KSPIN_LOCK PciLock;
static ULONG BusType = 0;

VOID DrvUnloadDriver(__in PDRIVER_OBJECT pDriverObject);
NTSTATUS DrvCreateClose(__in PDEVICE_OBJECT DeviceObject, __in PIRP Irp);

VOID InitPciBus(VOID);
ULONG GetBusConfigType(VOID);
static ULONG PciConfigUlong(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset);
static ULONG SizeBar(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset);
ULONG ConvertBARsz(ULONG size);

NTSTATUS DriverEntry(__in PDRIVER_OBJECT pDriverObject, __in PUNICODE_STRING RegistryPath)
// Driver EntryPoint
{
	NTSTATUS        Status;
    UNICODE_STRING  DeviceName;
    UNICODE_STRING  SymLink;
    PDEVICE_OBJECT  pDeviceObject = NULL;
	PHYSICAL_ADDRESS bar;
	ULONG out, size;
	PULONG framebuffer = NULL;
	int bus, slot, i;

    RtlInitUnicodeString(&DeviceName,NT_DEVICE_NAME);

    Status = IoCreateDevice(pDriverObject,0,&DeviceName,FILE_DEVICE_UNKNOWN,FILE_DEVICE_SECURE_OPEN,FALSE,&pDeviceObject);

    if (!NT_SUCCESS(Status))
    {
        // Error creating device.
        return Status;
    }

    pDriverObject->MajorFunction[IRP_MJ_CREATE] = DrvCreateClose;
    pDriverObject->MajorFunction[IRP_MJ_CLOSE] = DrvCreateClose;
    pDriverObject->DriverUnload = DrvUnloadDriver;

    RtlInitUnicodeString(&SymLink,DOS_DEVICE_NAME);

    Status = IoCreateSymbolicLink(&SymLink,&DeviceName);

    if (!NT_SUCCESS(Status))
    {
        // Error creating SymLink.
        IoDeleteDevice(pDeviceObject);
        return Status;
    }

	InitPciBus();
	for(bus = 0; bus < MAX_PCI_BUS; bus++)
	{
		for(slot=0; slot < MAX_PCI_SLOT; slot++)
		{
			out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0);
			if((0xFFFF & out) != 0xFFFF && (0xFFFF & out) > 0) {
				out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 8); 
				if((0xFF000000 & out)>>24 == 3) {
					DbgPrint("Bus numero %i:%i (Carte graph)\n", bus, slot);
					out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0);
					DbgPrint("\tDevice ID: %x\n\tVendor ID: %x\n", (0xFFFF0000 & out)>>16, (0xFFFF & out));
					out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 4);
					DbgPrint("\tStatus: %x\n\tCommand: %x\n", (0xFFFF0000 & out)>>16, (0xFFFF & out));
					out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 8);
					DbgPrint("\tClass code: %x\n\tSubclass: %x\n\tRevision ID: %x", (0xFF000000 & out)>>24, (0xFF0000 & out)>>16, (0xFFFF & out));
					out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0x30);
					DbgPrint("\tExpansion ROM base address : %x\n", out);

					for(i = 0; i<6; i++) {
						out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0x10+i*4);
						if(out > 0) {
							DbgPrint("\t\tBase address #%i (BAR%i)\n", i, i);
							size = SizeBar((UCHAR)bus, (UCHAR)slot, 0, 0x10+i*4);
							DbgPrint("\t\t\tPhysical [0x%x - 0x%x]\n", out, out+size);
						}
					}
				}
			} 
		}
	}

    return Status;
}


VOID DrvUnloadDriver(__in PDRIVER_OBJECT pDriverObject)
// Driver Unload
{
    PDEVICE_OBJECT pDeviceObject = pDriverObject->DeviceObject;
    UNICODE_STRING SymLink;

    RtlInitUnicodeString(&SymLink,DOS_DEVICE_NAME);
    IoDeleteSymbolicLink(&SymLink);

    if(pDriverObject != NULL) IoDeleteDevice(pDeviceObject);
}


NTSTATUS DrvCreateClose( __in PDEVICE_OBJECT DeviceObject, __in PIRP Irp)
{
    NTSTATUS Status;

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    Status = STATUS_SUCCESS;

    return Status;
}

VOID InitPciBus(VOID) {
	KeInitializeSpinLock(&PciLock);
	BusType = GetBusConfigType();
	DbgPrint("Bus type: %i\n", BusType);
}

ULONG GetBusConfigType(VOID) {
	ULONG Value;
	KIRQL oldIrql;
	
	KeAcquireSpinLock(&PciLock, &oldIrql);

	WRITE_PORT_UCHAR((PUCHAR)0xcfb, 0x01);
	Value = READ_PORT_ULONG((PULONG)0xCF8);
	WRITE_PORT_ULONG((PULONG)0xcf8, 0x80000000);
	if (READ_PORT_ULONG((PULONG)0xcf8) == 0x80000000) {
		WRITE_PORT_ULONG((PULONG)0xcf8, Value);
		KeReleaseSpinLock(&PciLock, oldIrql);
		return 1;
    }
	WRITE_PORT_ULONG((PULONG)0xcf8, Value);

	WRITE_PORT_UCHAR((PUCHAR)0xcfb, 0x00);
	WRITE_PORT_UCHAR((PUCHAR)0xcf8, 0x00);
	WRITE_PORT_UCHAR((PUCHAR)0xcfa, 0x00);
	if (READ_PORT_UCHAR((PUCHAR)0xcf8) == 0x00 && READ_PORT_UCHAR((PUCHAR)0xcfB) == 0x00) {
		KeReleaseSpinLock(&PciLock, oldIrql);
		return 2;
    }
	KeReleaseSpinLock(&PciLock, oldIrql);
	return 0;
}

static ULONG PciConfigUlong(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset)
{ 
	KIRQL oldIrql;
	ULONG tmp;
	ULONG address;
	ULONG lbus = (ULONG)bus;
	ULONG lslot = (ULONG)slot;
	ULONG lfunc = (ULONG)func;

	address = ((lbus << 16) | (lslot << 11) |(lfunc << 8 ) | (offset & 0xfc) | ((UINT32)0x80000000));

	if(BusType == 1)
	{
		KeAcquireSpinLock(&PciLock, &oldIrql);
		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		tmp = READ_PORT_ULONG((PULONG)0xcfc);
		KeReleaseSpinLock(&PciLock, oldIrql);
		return tmp;
	}
	return STATUS_UNSUCCESSFUL;
}

static ULONG SizeBar(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset)
{ 
	KIRQL oldIrql;
	ULONG tmp, address, bar;
	ULONG lbus = (ULONG)bus;
	ULONG lslot = (ULONG)slot;
	ULONG lfunc = (ULONG)func;

	address = ((lbus << 16) | (lslot << 11) |(lfunc << 8 ) | (offset & 0xfc) | ((UINT32)0x80000000));

	if(BusType == 1)
	{
		KeAcquireSpinLock(&PciLock, &oldIrql);
		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		bar = READ_PORT_ULONG((PULONG)0xcfc);

		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		WRITE_PORT_ULONG((PULONG)0xcfc, 0xffffffff);
		tmp = ConvertBARsz(READ_PORT_ULONG((PULONG)0xcfc));

		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		WRITE_PORT_ULONG((PULONG)0xcfc, bar);
		KeReleaseSpinLock(&PciLock, oldIrql);
		return tmp;
	}
	return 0;
}

ULONG ConvertBARsz(ULONG size)
{
	ULONG realsize;

	realsize = (size&0xfffffff0);
	realsize = (1 + ~realsize);
	if(realsize & 1)
	{
		realsize = (realsize&0xffff);
	} else {
		realsize = (realsize&0xffffffff);
	}
	return realsize;
}