#include "stdafx.h"
#include <atlbase.h>
#include "shlwapi.h"
#include "detours.h"
#include "windows.h"
#include <tlhelp32.h>
#include <iostream.h>
#include <fstream.h>
#include <sstream>
#include <string>
#include <io.h>
#include <winnt.h>
using namespace std;;

typedef LONG NTSTATUS;

#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >= 0)

typedef enum _FILE_INFORMATION_CLASS {
    FileDirectoryInformation         = 1,
        FileFullDirectoryInformation,   // 2
        FileBothDirectoryInformation,   // 3
        FileBasicInformation,           // 4 wdm
        FileStandardInformation,        // 5 wdm
        FileInternalInformation,        // 6
        FileEaInformation,              // 7
        FileAccessInformation,          // 8
        FileNameInformation,            // 9
        FileRenameInformation,          // 10
        FileLinkInformation,            // 11
        FileNamesInformation,           // 12
        FileDispositionInformation,     // 13
        FilePositionInformation,        // 14 wdm
        FileFullEaInformation,          // 15
        FileModeInformation,            // 16
        FileAlignmentInformation,       // 17
        FileAllInformation,             // 18
        FileAllocationInformation,      // 19
        FileEndOfFileInformation,       // 20 wdm
        FileAlternateNameInformation,   // 21
        FileStreamInformation,          // 22
        FilePipeInformation,            // 23
        FilePipeLocalInformation,       // 24
        FilePipeRemoteInformation,      // 25
        FileMailslotQueryInformation,   // 26
        FileMailslotSetInformation,     // 27
        FileCompressionInformation,     // 28
        FileObjectIdInformation,        // 29
        FileCompletionInformation,      // 30
        FileMoveClusterInformation,     // 31
        FileQuotaInformation,           // 32
        FileReparsePointInformation,    // 33
        FileNetworkOpenInformation,     // 34
        FileAttributeTagInformation,    // 35
        FileTrackingInformation,        // 36
        FileIdBothDirectoryInformation, // 37
        FileIdFullDirectoryInformation, // 38
        FileValidDataLengthInformation, // 39
        FileShortNameInformation,       // 40
        FileMaximumInformation
} FILE_INFORMATION_CLASS, *PFILE_INFORMATION_CLASS;

typedef struct _FILE_BOTH_DIR_INFORMATION {
    ULONG             NextEntryOffset;
    ULONG             FileIndex;
    LARGE_INTEGER     CreationTime;
    LARGE_INTEGER     LastAccessTime;
    LARGE_INTEGER     LastWriteTime;
    LARGE_INTEGER     ChangeTime;
    LARGE_INTEGER     EndOfFile;
    LARGE_INTEGER     AllocationSize;
    ULONG             FileAttributes;
    ULONG             FileNameLength;
    ULONG             EaSize;
    CCHAR             ShortNameLength;
    WCHAR             ShortName[12];
    WCHAR             FileName[1];
} FILE_BOTH_DIR_INFORMATION, *PFILE_BOTH_DIR_INFORMATION;

typedef struct _STRING {
    USHORT Length;
    USHORT MaximumLength;
    PCHAR Buffer;
} UNICODE_STRING,*PUNICODE_STRING;

typedef struct _IO_STATUS_BLOCK{  
    DWORD Status;  
    ULONG Information;  
} IO_STATUS_BLOCK, *PIO_STATUS_BLOCK;

typedef VOID (NTAPI *PIO_APC_ROUTINE)(PVOID ApcContext, PIO_STATUS_BLOCK IoStatusBlock, ULONG Reserved);

typedef NTSTATUS (*ZWQUERYDIRECTORYFILE)(
    IN  HANDLE FileHandle,
    IN  HANDLE Event OPTIONAL,
    IN  PIO_APC_ROUTINE ApcRoutine OPTIONAL,
    IN  PVOID ApcContext OPTIONAL,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    OUT PVOID FileInformation,
    IN  ULONG Length,
    IN  FILE_INFORMATION_CLASS FileInformationClass,
    IN  BOOLEAN ReturnSingleEntry,
    IN  PUNICODE_STRING FileName OPTIONAL,
    IN  BOOLEAN RestartScan
);

//hook zwQueryDirectoryFile
NTSTATUS WINAPI CopyZwQueryDirectoryFile(                        
    IN  HANDLE FileHandle,
    IN  HANDLE Event OPTIONAL,
    IN  PIO_APC_ROUTINE ApcRoutine OPTIONAL,
    IN  PVOID ApcContext OPTIONAL,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    OUT PVOID FileInformation,
    IN  ULONG Length,
    IN  FILE_INFORMATION_CLASS FileInformationClass,
    IN  BOOLEAN ReturnSingleEntry,
    IN  PUNICODE_STRING FileName OPTIONAL,
    IN  BOOLEAN RestartScan
);
NTSTATUS WINAPI MyZwQueryDirectoryFile(                        
     IN  HANDLE FileHandle,
     IN  HANDLE Event OPTIONAL,
     IN  PIO_APC_ROUTINE ApcRoutine OPTIONAL,
     IN  PVOID ApcContext OPTIONAL,
     OUT PIO_STATUS_BLOCK IoStatusBlock,
     OUT PVOID FileInformation,
     IN  ULONG Length,
     IN  FILE_INFORMATION_CLASS FileInformationClass,
     IN  BOOLEAN ReturnSingleEntry,
     IN  PUNICODE_STRING FileName OPTIONAL,
     IN  BOOLEAN RestartScan
     )
{
    return 0;
}

BOOL APIENTRY DllMain(HANDLE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
{
    HMODULE hModuleNT;
    ZWQUERYDIRECTORYFILE* zwQueryDirectoryFile;

    switch (ul_reason_for_call)
    {
        case DLL_PROCESS_ATTACH:
            hModuleNT = LoadLibrary("ntdll.dll");
            zwQueryDirectoryFile = (ZWQUERYDIRECTORYFILE*)GetProcAddress(hModuleNT, "ZwQueryDirectoryFile");
            DetourFunctionWithEmptyTrampoline((PBYTE)zwQueryDirectoryFile,(PBYTE)CopyZwQueryDirectoryFile,(PBYTE)MyZwQueryDirectoryFile);
            break;
        case DLL_THREAD_ATTACH:
            break;
        case DLL_THREAD_DETACH:
            break;
        case DLL_PROCESS_DETACH:
            DetourRemoveWithTrampoline((PBYTE)CopyZwQueryDirectoryFile, (PBYTE)MyZwQueryDirectoryFile);
            break;
    }
    return TRUE;
}