/*
 * Copyright (c) 2010 Andy Polyakov <appro@fy.chalmers.se>
 *
 * Build with:
 *
 * cl -Ox -GD -GF -GS- -Zl -MD -LD statpatch.c kernel32.lib ntdll.lib
 *
 * See http://fy.chalmers.se/~appro/nt/DLL_PRELOAD/ for further details.
 *
 */
#ifndef _DLL
#error "_DLL is not defined."
#endif

#ifdef _WIN64
#pragma comment(linker,"/entry:DllMain")
#pragma comment(linker,"/merge:.rdata=.text")
#else
#pragma comment(linker,"/entry:DllMain@12")
#pragma comment(linker,"/section:.text,erw")
#pragma comment(linker,"/merge:.rdata=.text")
#pragma comment(linker,"/merge:.data=.text")
#endif

#define UNICODE
#define _UNICODE

#if defined(WIN32) && !defined(_WIN32)
#define _WIN32
#endif

#define _WIN32_WINNT 0x0500

#include <windows.h>
#include <winbase.h>
#include <winnt.h>
#ifdef _WIN64
/* October 2002 Platform SDK is screwed up */
#define _RUNTIME_FUNCTION _RUNTIME_FUNCTION_
#define RUNTIME_FUNCTION  RUNTIME_FUNCTION_
#define PRUNTIME_FUNCTION PRUNTIME_FUNCTION_
#endif
#include <winternl.h>
#include <tchar.h>

#ifndef STATUS_SUCCESS
#define STATUS_SUCCESS			(0x0)
#define STATUS_ACCESS_DENIED		(0xc0000022L)
#endif

#ifndef FILE_DIRECTORY_FILE
/* Excerpt from <ddk>/inc/wdm.h */
#define FILE_DIRECTORY_FILE                     0x00000001
#define FILE_WRITE_THROUGH                      0x00000002
#define FILE_SEQUENTIAL_ONLY                    0x00000004
#define FILE_NO_INTERMEDIATE_BUFFERING          0x00000008

#define FILE_SYNCHRONOUS_IO_ALERT               0x00000010
#define FILE_SYNCHRONOUS_IO_NONALERT            0x00000020
#define FILE_NON_DIRECTORY_FILE                 0x00000040
#define FILE_CREATE_TREE_CONNECTION             0x00000080

#define FILE_COMPLETE_IF_OPLOCKED               0x00000100
#define FILE_NO_EA_KNOWLEDGE                    0x00000200
#define FILE_OPEN_FOR_RECOVERY                  0x00000400
#define FILE_RANDOM_ACCESS                      0x00000800

#define FILE_DELETE_ON_CLOSE                    0x00001000
#define FILE_OPEN_BY_FILE_ID                    0x00002000
#define FILE_OPEN_FOR_BACKUP_INTENT             0x00004000
#define FILE_NO_COMPRESSION                     0x00008000
#endif

/* from <ddk>/inc/ntdef.h */
#ifndef OBJ_CASE_INSENSITIVE
#define OBJ_CASE_INSENSITIVE            (0x00000040L)
#endif
#ifndef InitializeObjectAttributes
#define InitializeObjectAttributes( p, n, a, r, s ) {   \
    (p)->Length = sizeof( OBJECT_ATTRIBUTES );          \
    (p)->RootDirectory = r;                             \
    (p)->Attributes = a;                                \
    (p)->ObjectName = n;                                \
    (p)->SecurityDescriptor = s;                        \
    (p)->SecurityQualityOfService = NULL;               \
    }
#endif

/* from "Windows NT Native API" */
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;

#ifdef _DEBUG
static VOID DebugOutputA (const char *fmt,...)
{ va_list argv;
  char buf[256];

  va_start(argv,fmt);
    _vsnprintf (buf,sizeof(buf)/sizeof(buf[0])-1,fmt,argv);
    buf[sizeof(buf)/sizeof(buf[0])-1]='\0';
    OutputDebugStringA (buf);
  va_end(argv);
}
#endif

NTSYSAPI NTSTATUS NTAPI ZwClose(
    IN HANDLE FileHandle
    );

static IMAGE_THUNK_DATA *__NtCreateFile=NULL,*__NtCreateFile7=NULL;

NTSYSAPI NTSTATUS NTAPI ZwCreateFile(
    OUT PHANDLE FileHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    IN PLARGE_INTEGER AllocationSize OPTIONAL,
    IN ULONG FileAttributes,
    IN ULONG ShareAccess,
    IN ULONG CreateDisposition,
    IN ULONG CreateOptions,
    IN PVOID EaBuffer OPTIONAL,
    IN ULONG EaLength
    );

#ifdef _DEBUG
static NTSTATUS WINAPI NtCreateFile_ (
    OUT PHANDLE FileHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    IN PLARGE_INTEGER AllocationSize OPTIONAL,
    IN ULONG FileAttributes,
    IN ULONG ShareAccess,
    IN ULONG CreateDisposition,
    IN ULONG CreateOptions,
    IN PVOID EaBuffer OPTIONAL,
    IN ULONG EaLength
    )
{ NTSTATUS ret;

    ret = ZwCreateFile(
            FileHandle,
            DesiredAccess,
            ObjectAttributes,
            IoStatusBlock,
            AllocationSize,
            FileAttributes,
            ShareAccess,
            CreateDisposition,
            CreateOptions,
            EaBuffer,
            EaLength);

    DebugOutputA("create(\"%.*S\",%x,%x,%x)=%x",
		ObjectAttributes->ObjectName->Length/2,
    		ObjectAttributes->ObjectName->Buffer,
    		DesiredAccess,ShareAccess,CreateOptions,ret);

  return ret;
}
#endif

static IMAGE_THUNK_DATA *__NtOpenFile=NULL,*__NtOpenFile7=NULL;

NTSYSAPI NTSTATUS NTAPI ZwOpenFile(
    OUT PHANDLE FileHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    IN ULONG ShareAccess,
    IN ULONG OpenOptions
    );

#ifdef _DEBUG
static NTSTATUS WINAPI NtOpenFile_ (
    OUT PHANDLE FileHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    IN ULONG ShareAccess,
    IN ULONG OpenOptions
    )
{ NTSTATUS ret;

    ret = ZwOpenFile(
            FileHandle,
            DesiredAccess,
            ObjectAttributes,
            IoStatusBlock,
            ShareAccess,
            OpenOptions);

    DebugOutputA("open(\"%.*S\",%x,%x,%x)=%x",
		ObjectAttributes->ObjectName->Length/2,
    		ObjectAttributes->ObjectName->Buffer,
    		DesiredAccess,ShareAccess,OpenOptions,ret);

  return ret;
}
#endif

static IMAGE_THUNK_DATA *__NtQueryDirectoryFile=NULL,*__NtQueryDirectoryFile7=NULL;

NTSYSAPI NTSTATUS NTAPI 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 FileInformationLength,
    IN FILE_INFORMATION_CLASS FileInformationClass,
    IN BOOLEAN ReturnSingleEntry,
    IN PUNICODE_STRING FileName OPTIONAL,
    IN BOOLEAN RestartScan
    );

static NTSTATUS WINAPI NtQueryDirectoryFile_ (
    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 FileInformationLength,
    IN FILE_INFORMATION_CLASS FileInformationClass,
    IN BOOLEAN ReturnSingleEntry,
    IN PUNICODE_STRING FileName OPTIONAL,
    IN BOOLEAN RestartScan
    )
{ NTSTATUS ret;

    ret = ZwQueryDirectoryFile(
            FileHandle,
            Event,
            ApcRoutine,
            ApcContext,
            IoStatusBlock,
            FileInformation,
            FileInformationLength,
            FileInformationClass,
            ReturnSingleEntry,
            FileName,
            RestartScan);

    if (ret==STATUS_ACCESS_DENIED && ReturnSingleEntry && FileInformationClass==3)
    {	HANDLE h;
        OBJECT_ATTRIBUTES oa;
        IO_STATUS_BLOCK io;

#ifdef _DEBUG
	DebugOutputA(">re-open \"%.*S\"",FileName->Length/2,FileName->Buffer);
#endif
        InitializeObjectAttributes(&oa,FileName,OBJ_CASE_INSENSITIVE,FileHandle,NULL);

        if (ZwOpenFile(&h,SYNCHRONIZE|FILE_READ_ATTRIBUTES|FILE_LIST_DIRECTORY,
        	&oa,&io,
        	FILE_SHARE_DELETE|FILE_SHARE_WRITE|FILE_SHARE_READ,
        	FILE_OPEN_FOR_BACKUP_INTENT|FILE_SYNCHRONOUS_IO_NONALERT|FILE_DIRECTORY_FILE)
            == STATUS_SUCCESS)
        {   /* Query . instead */
	    ret = ZwQueryDirectoryFile(
                        h,
                        Event,
                        ApcRoutine,
                        ApcContext,
                        IoStatusBlock,
                        FileInformation,
                        FileInformationLength,
                        FileInformationClass,
                        ReturnSingleEntry,
                        NULL,
                        RestartScan);
            if (ret==STATUS_SUCCESS)
            {	FILE_BOTH_DIR_INFORMATION *inf=(FILE_BOTH_DIR_INFORMATION *)FileInformation;
		ULONG len = FileName->Length;
#ifdef _DEBUG
		DebugOutputA(">fixing FileInformation->FileName");
#endif
		if (len > (FileInformationLength-sizeof(*inf)))
		    len=FileInformationLength-sizeof(*inf);
		inf->FileNameLength = len;
		memcpy(inf->FileName,FileName->Buffer,len);
		inf->ShortNameLength=0;
	    }
            ZwClose(h);
	}
    }
#ifdef _DEBUG
    DebugOutputA("query(\"%*S\",%d,%d)=%x",
        		FileName ? FileName->Length/2 : 6,
        		FileName ? FileName->Buffer : L"(null)",
        		FileInformationLength,FileInformationClass,ret);
#endif
  return ret;
}

static int _lstricmp(const char *s1, const char *s2)
{ char c1,c2;
  int  ret;

    while (c1=*s1, c2=*s2, c1&&c2)
    {   c1|=0x20, c2|=0x20; /* lower the case */
        if (ret=c1-c2) return ret;
        s1++, s2++;
    }

  return c1-c2;
}

BOOL WINAPI DllMain (HINSTANCE h, DWORD reason, LPVOID junk)
{ DWORD                     acc;
  HMODULE                   hmod;
  IMAGE_DOS_HEADER         *dos_header;
  IMAGE_NT_HEADERS         *nt_headers;
  IMAGE_DATA_DIRECTORY     *dir;
  IMAGE_IMPORT_DESCRIPTOR  *idesc;
  IMAGE_THUNK_DATA         *thunk;
  static void              *page=NULL,*page7=NULL;
  static size_t             plen=0,plen7=0;

    switch (reason)
    { case DLL_PROCESS_ATTACH:

        DisableThreadLibraryCalls(h);

        if (!(hmod=GetModuleHandle(_T("NTDLL.DLL"))))
        {   OutputDebugString(_T("NTDLL.DLL not found?"));
            return FALSE;
        }

        if (!(hmod=GetModuleHandle(_T("KERNEL32.DLL"))))
        {   OutputDebugString(_T("KERNEL32.DLL not found?"));
            return FALSE;
        }

        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);

        dir=&nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        idesc=(IMAGE_IMPORT_DESCRIPTOR *)((char *)hmod + dir->VirtualAddress);
        while (idesc->Name)
        {   if (!_lstricmp((char *)hmod+idesc->Name,"NTDLL.DLL")) break;
            idesc++;
        }

        if (!idesc->Name)
        {   OutputDebugString(_T("Can't locate NTDLL.DLL import descriptor in KERNEL32"));
            return FALSE;
        }

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwCreateFile;
             thunk++) ;
        __NtCreateFile=thunk;

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwOpenFile;
             thunk++) ;
        __NtOpenFile=thunk;

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwQueryDirectoryFile;
             thunk++) ;
        __NtQueryDirectoryFile=thunk;

        if ((void *)__NtOpenFile > (void *)__NtCreateFile)
            page=__NtCreateFile,
            plen=(size_t)__NtOpenFile-(size_t)page+sizeof(void(*)());
        else
            page=__NtOpenFile,
            plen=(size_t)__NtCreateFile-(size_t)page+sizeof(void(*)());

	if ((void *)__NtQueryDirectoryFile > (void *)((BYTE*)page+plen))
	    plen=(size_t)__NtQueryDirectoryFile-(size_t)page+sizeof(void(*)());
	else if ((void *)__NtQueryDirectoryFile < page)
	    page=__NtQueryDirectoryFile,
	    plen+=(size_t)page-(size_t)__NtQueryDirectoryFile;

        if (!VirtualProtect (page,plen,PAGE_EXECUTE_READWRITE,&acc))
        {   OutputDebugString(_T("Unable to unlock KERNEL32 Thunk Table"));
            return FALSE;
        }
#ifdef _DEBUG
        __NtOpenFile->u1.Function=(size_t)NtOpenFile_;
        __NtCreateFile->u1.Function=(size_t)NtCreateFile_;
#endif
        __NtQueryDirectoryFile->u1.Function=(size_t)NtQueryDirectoryFile_;
        VirtualProtect (page,plen,acc,&acc);

	/*
	 * On Windows 7 most native calls go through KernelBase.dll
	 */

        if (!(hmod=GetModuleHandle(_T("KernelBase.dll"))))
        {   OutputDebugString(_T("KernelBase.dll not found..."));
            return TRUE;
        }

        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);

        dir=&nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        idesc=(IMAGE_IMPORT_DESCRIPTOR *)((char *)hmod + dir->VirtualAddress);
        while (idesc->Name)
        {   if (!_lstricmp((char *)hmod+idesc->Name,"NTDLL.DLL")) break;
            idesc++;
        }

        if (!idesc->Name)
        {   OutputDebugString(_T("Can't locate NTDLL.DLL import descriptor in KernelBase"));
            return TRUE;
        }

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwCreateFile;
             thunk++) ;
        __NtCreateFile7=thunk;

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwOpenFile;
             thunk++) ;
        __NtOpenFile7=thunk;

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwQueryDirectoryFile;
             thunk++) ;
        __NtQueryDirectoryFile7=thunk;

        if ((void *)__NtOpenFile7 > (void *)__NtCreateFile7)
            page7=__NtCreateFile7,
            plen7=(size_t)__NtOpenFile7-(size_t)page7+sizeof(void(*)());
        else
            page7=__NtOpenFile7,
            plen7=(size_t)__NtCreateFile7-(size_t)page7+sizeof(void(*)());

	if ((void *)__NtQueryDirectoryFile7 > (void *)((BYTE*)page7+plen7))
	    plen7=(size_t)__NtQueryDirectoryFile7-(size_t)page7+sizeof(void(*)());
	else if ((void *)__NtQueryDirectoryFile7 < page7)
	    page7=__NtQueryDirectoryFile7,
	    plen7+=(size_t)page7-(size_t)__NtQueryDirectoryFile7;

        if (!VirtualProtect (page7,plen7,PAGE_EXECUTE_READWRITE,&acc))
        {   OutputDebugString(_T("Unable to unlock KernelBase Thunk Table"));
            return FALSE;
        }
#ifdef _DEBUG
        __NtOpenFile7->u1.Function=(size_t)NtOpenFile_;
        __NtCreateFile7->u1.Function=(size_t)NtCreateFile_;
#endif
        __NtQueryDirectoryFile7->u1.Function=(size_t)NtQueryDirectoryFile_;
        VirtualProtect (page7,plen7,acc,&acc);

        break;

      case DLL_PROCESS_DETACH:

        VirtualProtect (page,plen,PAGE_EXECUTE_READWRITE,&acc);
#ifdef _DEBUG
        __NtOpenFile->u1.Function=(size_t)ZwOpenFile;
        __NtCreateFile->u1.Function=(size_t)ZwCreateFile;
#endif
        __NtQueryDirectoryFile->u1.Function=(size_t)ZwQueryDirectoryFile;
        VirtualProtect (page,plen,acc,&acc);

	if (page7==NULL) return TRUE;

        VirtualProtect (page7,plen7,PAGE_EXECUTE_READWRITE,&acc);
#ifdef _DEBUG
        __NtOpenFile7->u1.Function=(size_t)ZwOpenFile;
        __NtCreateFile7->u1.Function=(size_t)ZwCreateFile;
#endif
        __NtQueryDirectoryFile7->u1.Function=(size_t)ZwQueryDirectoryFile;
        VirtualProtect (page7,plen7,acc,&acc);

        break;
    }

  return TRUE;
}