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

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

#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>
#include <winternl.h>
#include <tchar.h>

typedef NTSTATUS (WINAPI *NtProtectVirtualMemory_T) (
    IN HANDLE ProcessHandle,
    IN OUT PVOID *BaseAddress,
    IN OUT PULONG ProtectSize,
    IN ULONG NewProtect,
    OUT PULONG OldProtect
    );

static int magicvalue=-1,passthrough=0;
static struct {
    void *addr0,*addr1;
    char *msg;
} sentinel[3];

static void WINAPI NtProtectVirtualMemory__(
    IN HANDLE ProcessHandle,
    IN OUT PVOID *BaseAddress,
    IN OUT PULONG ProtectSize,
    IN ULONG NewProtect,
    OUT PULONG OldProtect
    ) { /* return whatever next function returns */ }

static int NtProtectVirtualMemory_(
    void         *caller,
    IN HANDLE     ProcessHandle,
    IN OUT PVOID *BaseAddress,
    IN OUT PULONG ProtectSize,
    IN ULONG      NewProtect,
    OUT PULONG    OldProtect
    )
{ int   i;
  void *p;

    if (passthrough)
    { size_t page1=(size_t)caller & ~4095,
             page2=(size_t)NtProtectVirtualMemory_ & ~4095;
       passthrough=0;
       if (page1 == page2)
       {   OutputDebugString(_T("VirtualProtect Sentinel: letting myself through..."));
           return magicvalue;
       }
    }

    for (p=*BaseAddress,i=0;i<3;i++)
    {   if (p>=sentinel[i].addr0 && p<sentinel[i].addr1)
        { volatile void **ra=&caller-1;
            OutputDebugStringA(sentinel[i].msg);
            *ra=NtProtectVirtualMemory__; /* wow! */
            return STATUS_ACCESS_VIOLATION;
        }
    }

  return magicvalue;
}

static _VirtualLockup(
    HINSTANCE hmod,
    LPVOID entry,
    SIZE_T dwSize,
    DWORD flNewProtect,
    PDWORD lpflOldProtect)
{ NtProtectVirtualMemory_T  func=(NtProtectVirtualMemory_T)entry;
  SIZE_T len;
  HANDLE pid=GetCurrentProcess();

    passthrough=1;
    (*func) (pid,&entry,&dwSize,flNewProtect,lpflOldProtect);

    len=(size_t)sentinel[2].addr1-(size_t)sentinel[2].addr0;

    passthrough=1;
    /*
     * This locks up the whole module, no variable
     * modifications will be possible.
     */
    (*func) (pid,&hmod,&len,PAGE_EXECUTE_READ,lpflOldProtect);
}

BOOL WINAPI DllMain (HINSTANCE self, DWORD reason, LPVOID junk)
{ DWORD                 acc;
  HMODULE               hmod;
  unsigned char        *entry;
  MSGBOXPARAMS          m;
  IMAGE_DOS_HEADER     *dos_header;
  IMAGE_NT_HEADERS     *nt_headers;

    m.cbSize              = sizeof(m);
    m.hwndOwner           = NULL;
    m.lpszCaption         = _T("VirtualProtect Sentinel");
    m.dwStyle             = MB_OK;
    m.hInstance           = NULL;
    m.lpszIcon            = IDI_ERROR;
    m.dwContextHelpId     = 0;
    m.lpfnMsgBoxCallback  = NULL;
    m.dwLanguageId        = MAKELANGID(LANG_ENGLISH,SUBLANG_ENGLISH_US);

    switch (reason)
    { case DLL_PROCESS_ATTACH:

        DisableThreadLibraryCalls(self);

        hmod = self;
        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);
        sentinel[2].addr0=hmod;
        sentinel[2].addr1=(char *)hmod+nt_headers->OptionalHeader.SizeOfImage;
        sentinel[2].msg="VirtualProtect Sentinel: protecting myself";

        if (!(hmod=GetModuleHandle(_T("KERNEL32.DLL"))))
        {   m.lpszText = _T("KERNEL32.DLL not found?, aborting...");
            MessageBoxIndirect(&m);
            ExitProcess(0);
            return FALSE;
        }
        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);
        sentinel[1].addr0=hmod;
        sentinel[1].addr1=(char *)hmod+nt_headers->OptionalHeader.SizeOfImage;
        sentinel[1].msg="VirtualProtect Sentinel: protecting KERNEL32.DLL";

        if (!(hmod=GetModuleHandle(_T("NTDLL.DLL"))))
        {   m.lpszText = _T("NTDLL.DLL not found?, aborting...");
            MessageBoxIndirect(&m);
            ExitProcess(0);
            return FALSE;
        }
        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);
        sentinel[0].addr0=hmod;
        sentinel[0].addr1=(char *)hmod+nt_headers->OptionalHeader.SizeOfImage;
        sentinel[0].msg="VirtualProtect Sentinel: protecting NTDLL.DLL";

        entry=(unsigned char*)GetProcAddress(hmod,"NtProtectVirtualMemory");

        if (!entry || entry[0] != 0xb8)   /* mov  eax,imm32   */
        {   m.lpszText = _T("Unexpected instruction, aborting...");
            MessageBoxIndirect(&m);
            ExitProcess(0);
            return FALSE;
        }

        if (!VirtualProtect (entry,5,PAGE_EXECUTE_READWRITE,&acc))
        {   m.lpszText = _T("Unable to unlock code page, aborting...");
            MessageBoxIndirect(&m);
            ExitProcess(0);
            return FALSE;
        }

        magicvalue=*((unsigned int *)(entry+1));

        entry[0] = 0xe8;        /* call relative    */
        *((unsigned int *)(entry+1)) =
                (unsigned char *)NtProtectVirtualMemory_
                -(entry+5);
        FlushInstructionCache(GetCurrentProcess(),entry,5);

        _VirtualLockup (self,entry,5,acc,&acc);

        break;

      case DLL_PROCESS_DETACH:
        return FALSE;
        break;
    }

  return TRUE;
}