#include <Ntifs.h>
#include <Ntdddisk.h>

#include "fat16.h"

#define BUFSIZE                10
#define FILE_DEVICE_UNKNOWN    0x00000022

typedef unsigned int DWORD;
typedef char*        PCHAR;
typedef int          BOOL;

//
DRIVER_UNLOAD Unload;
NTSTATUS      DriverEntry(PDRIVER_OBJECT pDriverObj, PUNICODE_STRING pRegistryPath);
NTSTATUS      MakeReadWriteRequest(PDEVICE_OBJECT pDev, DWORD request, PCHAR pBuf, DWORD bufSize, PLARGE_INTEGER pOffset);
BOOL          FindDiskGeometry(PDEVICE_OBJECT pDev, PDISK_GEOMETRY pDiskGeo);
void          hexdump(unsigned char *data, unsigned int amount);
//

extern
NTSTATUS
ObReferenceObjectByName (
IN PUNICODE_STRING ObjectName,
IN ULONG Attributes,
IN PACCESS_STATE AccessState OPTIONAL,
IN ACCESS_MASK DesiredAccess OPTIONAL,
IN POBJECT_TYPE ObjectType,
IN KPROCESSOR_MODE AccessMode,
IN OUT PVOID ParseContext OPTIONAL,
OUT PVOID *Object
);

extern
POBJECT_TYPE*
IoDriverObjectType;

typedef struct
{
    PCHAR pBuf;
    DISK_GEOMETRY diskGeometry;
} DEVICE_EXTENSION,
  *PDEVICE_EXTENSION;


NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObj, PUNICODE_STRING pRegistryPath)
{
    UNICODE_STRING uDriverName = {0};
    NTSTATUS status            = 0;
    PDRIVER_OBJECT pDiskDriv   = NULL;
    PDEVICE_OBJECT pHdDevice   = NULL;
    PCHAR pBuf                 = NULL;
    LARGE_INTEGER offset       = {0};
    PDEVICE_OBJECT pDevice     = NULL;
    PDEVICE_EXTENSION pDevExt  = NULL;
	DWORD bufsize              = 0;


    pDriverObj->DriverUnload = Unload;
    DbgPrint("[ Loading .. ]\n", pDriverObj);

    /** <INIT> **/

    DbgPrint("Initialization..");

    RtlInitUnicodeString(&uDriverName, L"\\Driver\\Disk");
    status = IoCreateDevice(pDriverObj,
                            sizeof(DEVICE_EXTENSION),
                            NULL,
                            FILE_DEVICE_UNKNOWN,
                            FILE_AUTOGENERATED_DEVICE_NAME,
                            FALSE,
                            &pDevice);
    if(!NT_SUCCESS(status))
    {
        ExFreePoolWithTag(pBuf, 't4pz');

        DbgPrint("\n[!] Error at IoCreateDevice() : 0x%x.\n", status);
        return STATUS_UNSUCCESSFUL;
    }

    pDevExt = (PDEVICE_EXTENSION)pDevice->DeviceExtension;
    RtlSecureZeroMemory(pDevExt, sizeof(DEVICE_EXTENSION));

    DbgPrint("[OK] -> 0x%x.\n", pDevExt);

    /** </INIT> **/
	/** <RAW DISK DEVICE POINTER> **/

    DbgPrint("Retrieving \\Driver\\Disk.sys' driver object pointer..");
    status = ObReferenceObjectByName(&uDriverName,
                                     OBJ_CASE_INSENSITIVE,
                                     NULL,
                                     0,
                                     *IoDriverObjectType,
                                     KernelMode,
                                     0,
                                     &pDiskDriv
                                     );

    if(!NT_SUCCESS(status))
    {
        DbgPrint("\n[!] Error at ObReferenceObjectsByName() : 0x%x.\n", status);
        return STATUS_UNSUCCESSFUL;
    }

    DbgPrint("[OK] -> 0x%x.\n", pDiskDriv);
    pHdDevice = pDiskDriv->DeviceObject;
    DbgPrint("HardDisk 0 device pointer : 0x%x\n", pHdDevice);

	/** </RAW DISK DEVICE POINTER> **/
	/** <DISK GEOMETRY> **/

	DbgPrint("Trying to obtain disk geometry..\n");
    if(FindDiskGeometry(pHdDevice, &pDevExt->diskGeometry) == FALSE)
    {
        DbgPrint("\n[!] Error at FindDiskGeometry().\n");
        return STATUS_UNSUCCESSFUL;
    }

    DbgPrint("[OK]\nYou have %d%d cylinders\nA cylinders is composed of %d tracks\nA Track is composed of %d sectors\nA Sector is composed of %d bytes.\n",
             pDevExt->diskGeometry.Cylinders.HighPart, pDevExt->diskGeometry.Cylinders.LowPart,
             pDevExt->diskGeometry.TracksPerCylinder,
             pDevExt->diskGeometry.SectorsPerTrack,
             pDevExt->diskGeometry.BytesPerSector);
	
	bufsize = pDevExt->diskGeometry.BytesPerSector;
	/** </DISK GEOMETRY> **/
	/** <ALLOCATION> **/

	DbgPrint("Allocating memory..");
	pBuf = ExAllocatePoolWithTag(NonPagedPool, bufsize, 't4pz');

    if(pBuf == NULL)
    {
        DbgPrint("\n[!] Error at ExAllocatePoolWithTag().\n");
        return STATUS_UNSUCCESSFUL;
    }

	RtlSecureZeroMemory(pBuf, BUFSIZE);
	pDevExt->pBuf = pBuf;
	DbgPrint("[OK] -> 0x%x\n", pBuf);

	/** <ALLOCATION> **/
	/** <READ REQUEST> **/

    DbgPrint("Making a read request..\n");

    status = MakeReadWriteRequest(pHdDevice,
                                  IRP_MJ_READ,
                                  pBuf,
                                  bufsize,
                                  &offset);
    if(!NT_SUCCESS(status))
    {
        ExFreePoolWithTag(pBuf, 't4pz');

        DbgPrint("[!] Error at MakeReadWriteRequest().\n");
        return STATUS_UNSUCCESSFUL;
    }

    DbgPrint("[OK] -> 0x%x.\n", pBuf);

    //RtlSecureZeroMemory(pBuf, bufsize);

	/** </ READ REQUEST> **/

    return STATUS_SUCCESS;
}

BOOL FindDiskGeometry(PDEVICE_OBJECT pDevHdDevice, PDISK_GEOMETRY pDiskGeo)
{
    KEVENT event              = {0};
    PIRP pIrp                 = NULL;
    IO_STATUS_BLOCK ioStatus  = {0};
    NTSTATUS status           = 0;
    PDEVICE_EXTENSION pDevExt = NULL;

    pDevExt = (PDEVICE_EXTENSION)pDevHdDevice->DeviceExtension;
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    DbgPrint("Building IRP which will be send to the HardDisk 0 device..");
    pIrp = IoBuildDeviceIoControlRequest(IOCTL_DISK_GET_DRIVE_GEOMETRY,
                                         pDevHdDevice,
                                         NULL,
                                         0,
                                         pDiskGeo,
                                         sizeof(DISK_GEOMETRY),
                                         FALSE,
                                         &event,
                                         &ioStatus
                                        );
    if(pIrp == NULL)
    {
        DbgPrint("\n[!] Error at IoBuildSynchronousFsdRequest().\n");
        return STATUS_UNSUCCESSFUL;
    }

    DbgPrint("[OK] -> 0x%x.\n", pIrp);

    DbgPrint("Calling the driver..");
    status = IoCallDriver(pDevHdDevice, pIrp);

    if(status == STATUS_PENDING)
	{
	    DbgPrint("Waiting it handles our query..");
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE,	NULL);
        status = ioStatus.Status;

        if(!NT_SUCCESS(status))
        {
            DbgPrint("\n[!] Error at IoCallDriver+KeWaitForSingleObject : 0x%x.\n", status);
            return FALSE;
        }
    }

    return TRUE;
}

NTSTATUS MakeReadWriteRequest(PDEVICE_OBJECT pDev, DWORD request, PCHAR pBuf, DWORD bufSize, PLARGE_INTEGER pOffset)
{
    KEVENT event             = {0};
    PIRP pIrp                = NULL;
    IO_STATUS_BLOCK ioStatus = {0};
    NTSTATUS status          = 0;

    KeInitializeEvent(&event, NotificationEvent, FALSE);
    DbgPrint("Building IRP which will be send to the HardDisk 0 device..");
    pIrp = IoBuildSynchronousFsdRequest(request,
                                        pDev,
                                        pBuf,
                                        bufSize,
                                        pOffset,
                                        &event,
                                        &ioStatus
                                       );
    if(pIrp == NULL)
    {
        DbgPrint("\n[!] Error at IoBuildSynchronousFsdRequest().\n");
        return STATUS_UNSUCCESSFUL;
    }

    DbgPrint("[OK] -> 0x%x.\n", pIrp);

    DbgPrint("Calling the driver..");
    status = IoCallDriver(pDev, pIrp);

    if(status == STATUS_PENDING)
	{
	    DbgPrint("Waiting it handles our query..");
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE,	NULL);
        status = ioStatus.Status;

        if(!NT_SUCCESS(status))
        {
            DbgPrint("\n[!] Error at IoCallDriver+KeWaitForSingleObject : 0x%x.\n", status);
            return STATUS_UNSUCCESSFUL;
        }
    }

    return STATUS_SUCCESS;
}

VOID Unload(PDRIVER_OBJECT pDrivObj)
{
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)pDrivObj->DeviceObject->DeviceExtension;
    DbgPrint("[ Unloading.. ]\n");

    if(pDevExt->pBuf != NULL)
        ExFreePoolWithTag(pDevExt->pBuf, 't4pz');

    IoDeleteDevice(pDrivObj->DeviceObject);
    return;
}

// thanx to kraken for the hexdump function
void hexdump(unsigned char *data, unsigned int amount)
{
   unsigned int dp, p;
   const char trans[] =
      "................................ !\"#$%&'()*+,-./0123456789"
      ":;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklm"
      "nopqrstuvwxyz{|}~...................................."
      "....................................................."
      "........................................";

   for (dp = 1; dp <= amount; dp++)
   {
      DbgPrint("%02x ", data[dp-1]);
      if ((dp % 8) == 0)
         DbgPrint(" ");

      if ((dp % 16) == 0)
	  {
         DbgPrint ("| ");
         p = dp;
         for (dp -= 16; dp < p; dp++)
            DbgPrint("%c", trans[data[dp]]);
         DbgPrint("\n");
      }
   }
   if ((amount % 16) != 0) 
   {
      p = dp = 16 - (amount % 16);
      for (dp = p; dp > 0; dp--) 
	  {
         DbgPrint("   ");
         if (((dp % 8) == 0) && (p != 8))
            DbgPrint(" ");
      }
      DbgPrint (" | ");
      for (dp = (amount - (16 - p)); dp < amount; dp++)
         DbgPrint("%c", trans[data[dp]]);
   }
   DbgPrint("\n");
   return;
}
