/********************************************************************
  created:  2013/02/26
  created:  25:2:2013   20:29
  filename:   h:\TsFuzzerk\TsFuzzerk\TsFuzzerkFuzzHandler.cpp
  file path:  h:\TsFuzzerk\TsFuzzerk
  file base:  TsFuzzerkFuzzHandler
  file ext:  cpp
  author:    tishion@163.com
  
  purpose:  The implementation of fuzz handler
*********************************************************************/
#include "TsFuzzerk.h"
#include "SysTypedef.h"
#include "TsFuzzerkFuzzItemList.h"

extern PVOID
g_pNtDeviceIoControlFile;

extern CTsFuzzerkFuzzItemList* 
g_pFuzzItemList;

/*
 * A stub to jmp to the original NtDeviceIoControlFile function after we
 * made it into our fake function.
 */
NTSTATUS __declspec(naked) NtDeviceIoControlFile_Stub(
  IN HANDLE  FileHandle,
  IN HANDLE  Event,
  IN PIO_APC_ROUTINE  ApcRoutine,
  IN PVOID  ApcContext,
  OUT PIO_STATUS_BLOCK  IoStatusBlock,
  IN ULONG  IoControlCode,
  IN PVOID  InputBuffer,
  IN ULONG  InputBufferLength,
  OUT PVOID  OutputBuffer,
  IN ULONG  OutputBufferLength)
{
  _asm
  {
    mov    edi, edi
    push  ebp
    mov    ebp, esp
    mov    eax, g_pNtDeviceIoControlFile
    add    eax, 5
    jmp    eax 
  }
}


#define _MAX_NAME_STRING 1024/sizeof(WCHAR)  
typedef struct _NAME_STRING {  
  UNICODE_STRING Name;  
  WCHAR     Buffer[_MAX_NAME_STRING];  
} NAME_STRING, *PNAME_STRING;   


//PsGetProcessPeb 取得指定EPROCESS的Peb地址
typedef PPEB (__stdcall *FUNP_PSGETPROCESSPEB)(PEPROCESS pEProcess);
PPEB PsGetProcessPebWrapper(PEPROCESS pEProcess)
{
  UNICODE_STRING ustrFunName;
  FUNP_PSGETPROCESSPEB fpPsGetProcessPeb = NULL;

  RtlInitUnicodeString(&ustrFunName, (PWCHAR)L"PsGetProcessPeb");
  fpPsGetProcessPeb = (FUNP_PSGETPROCESSPEB)MmGetSystemRoutineAddress(&ustrFunName);
  if (fpPsGetProcessPeb == NULL) 
  {
    return NULL;
  }

  return fpPsGetProcessPeb(pEProcess);
}

#define STRING_POOL_BUFFER_LEN  512*1024
CHAR g_Format[] = {
  "[=======================================TsFuzzer=======================================]\n" \
  "\tProcess:%wZ PID:0x%08X\n" \
  "\tDevice:%wZ Driver:%wZ\n" \
  "\tIOCTL Code:0x%08X Method:%s\n" \
  "\t++++++Parameters Before Fuzz+++++\n" \
  "\tIBuffer:0x%08X Length:0x%08X\n" \
  "\tOBuffer:0x%08X Length:0x%08X\n" \
  "\tInput Data:"
};

CHAR g_Format_after[] = {
  "\t------Parameters After Fuzz------\n" \
  "\tIBuffer:0x%08X Length:0x%08X\n" \
  "\tOBuffer:0x%08X Length:0x%08X\n" \
  "\tInput Data:"
};

PCHAR g_IoMethod[] = {
  "METHOD_BUFFERED",         // 0
  "METHOD_IN_DIRECT",              // 1
  "METHOD_OUT_DIRECT",             // 2
  "METHOD_NEITHER",                // 3
};

NTSTATUS Fake_NtDeviceIoControlFile(
                  IN HANDLE  FileHandle,
                  IN HANDLE  Event,
                  IN PIO_APC_ROUTINE  ApcRoutine,
                  IN PVOID  ApcContext,
                  OUT PIO_STATUS_BLOCK  IoStatusBlock,
                  IN ULONG  IoControlCode,
                  IN PVOID  InputBuffer,
                  IN ULONG  InputBufferLength,
                  OUT PVOID  OutputBuffer,
                  IN ULONG  OutputBufferLength)
{
  NTSTATUS status = STATUS_SUCCESS;
  //KdPrint(("[TsFuzzerk]:++++Fake_NtDeviceIoControlFile()\n"));

  ULONG ulControlCode = 0;
  ULONG ulPid = 0;

  NAME_STRING ImageName;
  ImageName.Name.Buffer = ImageName.Buffer;  
  ImageName.Name.Length = 0;  
  ImageName.Name.MaximumLength = sizeof(ImageName.Buffer);  

  NAME_STRING devName;
  devName.Name.Buffer = devName.Buffer;  
  devName.Name.Length = 0;  
  devName.Name.MaximumLength = sizeof(devName.Buffer);  

  NAME_STRING drvPath;
  drvPath.Name.Buffer = drvPath.Buffer;  
  drvPath.Name.Length = 0;  
  drvPath.Name.MaximumLength = sizeof(drvPath.Buffer);  

  // 获取DeviceName DriverName DriverPath
  PFILE_OBJECT pFileObject = NULL;

  ULONG ulReturnedLength = 0;

  status = ObReferenceObjectByHandle(
    FileHandle, 
    STANDARD_RIGHTS_ALL, 
    *IoFileObjectType, 
    KernelMode, 
    (PVOID*)&pFileObject, 
    NULL);

  if (NT_SUCCESS(status))
  {
    status = ObQueryNameString(
      pFileObject->DeviceObject, 
      (POBJECT_NAME_INFORMATION)&devName, 
      devName.Name.MaximumLength, 
      &ulReturnedLength);

    if (!NT_SUCCESS(status))
    {
      RtlStringCchPrintfW(devName.Buffer, _MAX_NAME_STRING, L"Unknown");
      RtlInitUnicodeString(&(devName.Name), devName.Buffer);
    }

    PLDR_DATA_TABLE_ENTRY pLdr = (PLDR_DATA_TABLE_ENTRY)
      (pFileObject->DeviceObject->DriverObject->DriverSection);
    if (NULL != pLdr && NULL != pLdr->FullDllName.Buffer && 0 != pLdr->FullDllName.Length)
    {
      RtlCopyUnicodeString(&(drvPath.Name), &(pLdr->FullDllName));
    }
    else
    {
      RtlStringCchPrintfW(drvPath.Buffer, _MAX_NAME_STRING, L"Unknown");
      RtlInitUnicodeString(&(drvPath.Name), drvPath.Buffer);
    }

    ObDereferenceObject(pFileObject);
  }

  PFUZZ_ITEM pFuzzListItem = g_pFuzzItemList->GetItemByDeviceName(&(devName.Name));

  if (NULL != pFuzzListItem)
  {
    // 获取IoControlCode
    ulControlCode = IoControlCode;

    if ((0 == pFuzzListItem->IoControlCode) || (ulControlCode == pFuzzListItem->IoControlCode))
    {
      // 获取进程镜像文件路径和PID
      ulPid = (ULONG)PsGetCurrentProcessId();

      PPEB pPeb = NULL;
      PEPROCESS pEprocess = NULL;
      pEprocess = PsGetCurrentProcess();
      if (NULL != pEprocess)
      {
        pPeb = PsGetProcessPebWrapper(pEprocess);
        if (NULL != pPeb)
        {
          RtlCopyUnicodeString(&(ImageName.Name), &(pPeb->ProcessParameters->ImagePathName));
        }
        else
        {
          RtlStringCchPrintfW(ImageName.Buffer, _MAX_NAME_STRING, L"Unknown");
          RtlInitUnicodeString(&(ImageName.Name), ImageName.Buffer);
        }
      }

      PCHAR pszStringBuffer = (PCHAR)ExAllocatePoolWithTag(
        NonPagedPool, STRING_POOL_BUFFER_LEN, 'FUBS');

      if (NULL != pszStringBuffer)
      {
        RtlStringCchPrintfA(pszStringBuffer, STRING_POOL_BUFFER_LEN, g_Format,
          &(ImageName.Name), ulPid, 
          &(devName.Name), &(drvPath.Name), 
          ulControlCode, g_IoMethod[METHOD_FROM_CTL_CODE(ulControlCode)],
          InputBuffer, InputBufferLength,
          OutputBuffer, OutputBufferLength);
      }

      ULONG ulDataStringLength = (InputBufferLength * 4);
      ulDataStringLength = ulDataStringLength + (((InputBufferLength / 0x10) + 1) * 0x10);

      PCHAR pszDataBuffer = (PCHAR)ExAllocatePoolWithTag(
        NonPagedPool, ulDataStringLength, 'FUBD');

      if (pszDataBuffer != NULL)
      {
        pszDataBuffer[0] = '\0';

        CHAR ByteData[] = {"00  "};
        CHAR AddressData[] = {"00000000  "};

        PUCHAR pInputBuffer = (PUCHAR)InputBuffer;

        for (ULONG i = 0; i < InputBufferLength; i++)
        {
          if (0 == i % 0x10)
          {
            RtlStringCchCatA(pszDataBuffer, ulDataStringLength, "\n\t");
            RtlStringCchPrintfA(AddressData, sizeof(AddressData), "%08x  ", (ULONG)InputBuffer + i);
            RtlStringCchCatA(pszDataBuffer, ulDataStringLength, AddressData);
          }

          RtlStringCchPrintfA(ByteData, sizeof(ByteData), "%02X ", pInputBuffer[i]);
          RtlStringCchCatA(pszDataBuffer, ulDataStringLength, ByteData);
        }

        RtlStringCchCatA(pszDataBuffer, ulDataStringLength, "\n");
      }

      DbgPrint(pszStringBuffer);
      DbgPrint(pszDataBuffer);


      // 在此处进行参数Fuzz
      if (FuzzMethod_FillWithRandomData == pFuzzListItem->fmInputBuffer)
      {
        LARGE_INTEGER liSystemTime;
        KeQuerySystemTime(&liSystemTime);
        for (ULONG i = 0; i < InputBufferLength / sizeof(ULONG); i++)
        {
          ((PULONG)InputBuffer)[i] = RtlRandomEx(&(liSystemTime.LowPart));
        }
      }
      else if (FuzzMethod_FillWithFixedDWord == pFuzzListItem->fmInputBuffer)
      {
        for (ULONG i = 0; i < InputBufferLength / sizeof(ULONG); i++)
        {
          ((PULONG)InputBuffer)[i] = pFuzzListItem->FixedDWord;
        }
      }

      if (pszDataBuffer != NULL)
      {
        pszDataBuffer[0] = '\0';

        CHAR ByteData[] = {"00  "};
        CHAR AddressData[] = {"00000000  "};

        PUCHAR pInputBuffer = (PUCHAR)InputBuffer;

        for (ULONG i = 0; i < InputBufferLength; i++)
        {
          if (0 == i % 0x10)
          {
            RtlStringCchCatA(pszDataBuffer, ulDataStringLength, "\n\t");
            RtlStringCchPrintfA(AddressData, sizeof(AddressData), "%08x  ", (ULONG)InputBuffer + i);
            RtlStringCchCatA(pszDataBuffer, ulDataStringLength, AddressData);
          }

          RtlStringCchPrintfA(ByteData, sizeof(ByteData), "%02X ", pInputBuffer[i]);
          RtlStringCchCatA(pszDataBuffer, ulDataStringLength, ByteData);
        }

        RtlStringCchCatA(pszDataBuffer, ulDataStringLength, "\n");
      }

      if (FuzzMethod_ReplaceByRandomLength == pFuzzListItem->fmInputLength)
      {
        LARGE_INTEGER liSystemTime;
        KeQuerySystemTime(&liSystemTime);
        InputBufferLength = RtlRandomEx(&(liSystemTime.LowPart));
      }
      else if (FuzzMethod_ReplaceByFixedLength == pFuzzListItem->fmInputLength)
      {
        InputBufferLength = pFuzzListItem->FixedLength;
      }

      if (NULL != pszStringBuffer)
      {
        RtlStringCchPrintfA(pszStringBuffer, STRING_POOL_BUFFER_LEN, g_Format_after,
          InputBuffer, InputBufferLength,
          OutputBuffer, OutputBufferLength);
      }

      DbgPrint(pszStringBuffer);
      DbgPrint(pszDataBuffer);

      if (NULL != pszStringBuffer)
      {
        ExFreePoolWithTag(pszStringBuffer, 'FUBS');
      }

      if (NULL != pszDataBuffer)
      {
        ExFreePoolWithTag(pszDataBuffer, 'FUBD');
      }
    }
  }

  status = NtDeviceIoControlFile_Stub(
    FileHandle,
    Event,
    ApcRoutine,
    ApcContext,
    IoStatusBlock,
    IoControlCode,
    InputBuffer,
    InputBufferLength,
    OutputBuffer,
    OutputBufferLength);

  //KdPrint(("[TsFuzzerk]:----Fake_NtDeviceIoControlFile()\n"));

  return status;
}

[课程]Linux pwn 探索篇！