#include "util.h"

#include <STDIO.H>

BOOL WriteMemory(LPVOID destination, LPVOID source, size_t length, LPVOID oldData)
{
  DWORD oldProtec;

  if (VirtualProtect(destination, length, PAGE_EXECUTE_READWRITE, &oldProtec)) {
    // Read data out if requested
    if (oldData) {
      memcpy(oldData, destination, length);
    }

    // Write data if provided
    if (source) {
      memcpy(destination, source, length);
    }

    // Restore original protection
    VirtualProtect(destination, length, oldProtec, &oldProtec);

    return TRUE;
  } else {
    return FALSE;
  }
}

SIZE_T SearchReplacePattern(LPVOID imageBase, LPCVOID search, LPCVOID replace, SIZE_T count, BOOL only_once)
{
  SIZE_T instances = 0;

  HANDLE process = GetCurrentProcess();

  MEMORY_BASIC_INFORMATION mbi = {0};

  // Loop through memory pages
  UINT_PTR addr = (UINT_PTR)imageBase;
  while ((!instances || !only_once) && VirtualQueryEx(process, (LPVOID)addr, &mbi, sizeof(mbi)) && mbi.AllocationBase == imageBase) {
    if (mbi.State == MEM_COMMIT && mbi.Protect != PAGE_NOACCESS) {
      DWORD oldProtec;

      // Try to gain access to this memory page
      if (VirtualProtect(mbi.BaseAddress, mbi.RegionSize, PAGE_EXECUTE_READWRITE, &oldProtec)) {
        // Loop through every byte to find the pattern
        SIZE_T maxOffset = mbi.RegionSize - count;
        for (SIZE_T i=0; i<maxOffset; i++) {
          LPVOID offset = (LPVOID)((UINT_PTR)(mbi.BaseAddress)+i);

          if (!memcmp(offset, search, count)) {
            // Found pattern, overwrite it
            memcpy(offset, replace, count);
            instances++;

            if (only_once) {
              break;
            }
          }
        }

        // Restore original permissions
        VirtualProtect(mbi.BaseAddress, mbi.RegionSize, oldProtec, &oldProtec);
      }
    }

    addr += mbi.RegionSize;
  }

  return instances;
}

LPVOID OverwriteImport(LPVOID imageBase, LPCSTR overrideFunction, LPVOID override)
{
  // Get DOS header
  PIMAGE_DOS_HEADER dosHeaders = (PIMAGE_DOS_HEADER)imageBase;

  // Retrieve NT header
  PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)((UINT_PTR)imageBase + dosHeaders->e_lfanew);

  // Retrieve imports directory
  IMAGE_DATA_DIRECTORY importsDirectory = ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

  // Retrieve import descriptor
  PIMAGE_IMPORT_DESCRIPTOR importDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)((UINT_PTR)imageBase + importsDirectory.VirtualAddress);

  while (importDescriptor->Name != NULL) {
    LPCSTR libraryName = (LPCSTR)importDescriptor->Name + (UINT_PTR)imageBase;
    HMODULE library = LoadLibraryA(libraryName);

    if (library) {
      PIMAGE_THUNK_DATA originalFirstThunk = (PIMAGE_THUNK_DATA)((UINT_PTR)imageBase + importDescriptor->OriginalFirstThunk);
      PIMAGE_THUNK_DATA firstThunk = (PIMAGE_THUNK_DATA)((UINT_PTR)imageBase + importDescriptor->FirstThunk);

      while (originalFirstThunk->u1.AddressOfData != NULL) {
        PIMAGE_IMPORT_BY_NAME functionName = (PIMAGE_IMPORT_BY_NAME)((UINT_PTR)imageBase + (UINT_PTR)originalFirstThunk->u1.AddressOfData);

        if (!strcmp((const char*)functionName->Name, overrideFunction)) {
          LPVOID originalFunction = (LPVOID)firstThunk->u1.Function;
          firstThunk->u1.Function = (PDWORD)override;

          // Return original function and end loop here
          printf("Hooked %s\n", overrideFunction);
          return originalFunction;
        }
        ++originalFirstThunk;
        ++firstThunk;
      }
    }

    importDescriptor++;
  }

  return NULL;
}

LPVOID OverwriteVirtualTable(LPVOID object, SIZE_T methodIndex, LPVOID overrideFunction)
{
  LPVOID *vtable = ((LPVOID**)(object))[0];
  LPVOID &functionAddress = vtable[methodIndex];
  LPVOID originalFunction = functionAddress;

  if (overrideFunction) {
    functionAddress = overrideFunction;
    printf("Hooked vtable %p+%lu\n", vtable, methodIndex);
  }

  return originalFunction;
}

BOOL OverwriteCall(LPVOID destination, LPVOID localCall)
{
  char callInst[5];

  callInst[0] = '\xE8';

  *(DWORD*)(&callInst[1]) = (DWORD)localCall - ((DWORD)destination + 5);

  return WriteMemory(destination, callInst, 5);
}