// .\out\Release\tiny.exe --v=1 --enable-logging=stderr
// --dll=D:\ghidra\enclave\dlls\prefs_enclave_x64.dll
// --write-data=32 --read-data=32

// This is a testing tiny main.

// .\out\Release\tiny.exe --v=1 --enable-logging=stderr
// --dll=D:\ghidra\enclave\dlls\prefs_enclave_x64.dll

#include <limits>

#include <windows.h>

#include <enclaveapi.h>

#include "base/command_line.h"
#include "base/files/file_path.h"
#include "base/logging.h"
#include "base/strings/string_number_conversions.h"

namespace switches {
constexpr char kEnableLogging[] = "enable-logging";
constexpr char kLogFile[] = "log-file";
// Sets the minimum log level. Valid values are from 0 to 3:
// INFO = 0, WARNING = 1, LOG_ERROR = 2, LOG_FATAL = 3.
constexpr char kLoggingLevel[] = "log-level";
// enclave dll path
constexpr char kDll[] = "dll";
// --data=string to show round-trip
constexpr char kSealData[] = "data";
// --read-enclave to get data from offset in enclave
constexpr char kReadData[] = "read-enclave";
// --write-enclave overwrite at offset in enclave
constexpr char kWriteData[] = "write-enclave";
}  // namespace switches

namespace {

bool InitLoggingFromCommandLine(const base::CommandLine* command_line) {
  logging::LoggingSettings settings;
  settings.log_file_path = L"c:/temp/tiny-debug.log";

  if (command_line->GetSwitchValueASCII(switches::kEnableLogging) == "stderr") {
    settings.logging_dest = logging::LOG_TO_STDERR;
  }
  if (command_line->HasSwitch(switches::kLogFile)) {
    settings.logging_dest |= logging::LOG_TO_FILE;
    settings.log_file_path =
        command_line->GetSwitchValueNative(switches::kLogFile).c_str();
    settings.delete_old = logging::DELETE_OLD_LOG_FILE;
  }
  logging::SetLogItems(true /* Process ID */, true /* Thread ID */,
                       true /* Timestamp */, false /* Tick count */);

  logging::InitLogging(settings);

  if (command_line->HasSwitch(switches::kLoggingLevel) &&
      logging::GetMinLogLevel() >= 0) {
    std::string log_level =
        command_line->GetSwitchValueASCII(switches::kLoggingLevel);
    int level = 0;
    if (base::StringToInt(log_level, &level) && level >= 0 &&
        level < logging::LOGGING_NUM_SEVERITIES) {
      logging::SetMinLogLevel(level);
    } else {
      DLOG(WARNING) << "Bad log level: " << log_level;
    }
  }
  return true;
}

typedef void* INIT_TAG;

typedef struct _prefs_init {
  char* init_name;
} PREFS_INIT;

typedef struct _seal_args {
  INIT_TAG config_ll;
  unsigned char* data_to_seal;
  unsigned char* protectedBlob;
  DWORD sz_data_to_seal;
  DWORD sz_protected_blob_size;
} SEAL_ARGS;

typedef struct _unseal_args {
  INIT_TAG config_ll;
  unsigned char* protected_blob;
  unsigned char* unsealed_data;
  DWORD sz_protected_blob;
  // not sure about these fields exactly
  DWORD unsealed_size;
  DWORD unsealed_size_max;
} UNSEAL_ARGS;

}  // namespace

int main(int argc, char** argv) {
  base::CommandLine::Init(argc, argv);
  InitLoggingFromCommandLine(base::CommandLine::ForCurrentProcess());

  VLOG(1) << "Verbose logging with --v=1 --enable-logging=stderr";

  // Faffing with enclaves
  ENCLAVE_CREATE_INFO_VBS create_info{};
  create_info.Flags = 0;
  create_info.OwnerID[0] = 1;  // 31 zero bytes for now?

  LPVOID enclave_addr = CreateEnclave(
      GetCurrentProcess(),
      nullptr,     // let Windows decide base
      0x10000000,  // Must be multiple of 2M in size - not committed initially?
      0,           // not used for VBS enclaves
      ENCLAVE_TYPE_VBS, &create_info, sizeof(ENCLAVE_CREATE_INFO_VBS),
      nullptr  // not used for VBS enclaves
  );

  VLOG(1) << "Create enclave " << enclave_addr << " gle: " << GetLastError();
  if (!enclave_addr) {
    return 1;
  }

  base::FilePath enclave_dll =
      base::CommandLine::ForCurrentProcess()->GetSwitchValuePath(
          switches::kDll);
  BOOL ret = LoadEnclaveImageW(enclave_addr, enclave_dll.value().c_str());

  DWORD gle = GetLastError();
  VLOG(1) << "Loaded image ret " << ret << " dll " << enclave_dll
          << " gle: " << gle;
  if (!ret) {
    return 2;
  }

  ENCLAVE_INIT_INFO_VBS init_info{};
  init_info.Length = sizeof(init_info);
  init_info.ThreadCount = 1;

  // This can fail (87) if Create params were wrong (e.g. size)
  ret = InitializeEnclave(GetCurrentProcess(), enclave_addr, &init_info,
                          sizeof(init_info), nullptr);
  gle = GetLastError();

  VLOG(1) << "Init Enclave " << ret << " tc " << init_info.ThreadCount
          << " gle: " << gle;
  if (!ret) {
    return 3;
  }

  // ordinal hint RVA      name
  //       1    0 00001080 Init
  //       2    1 000011E0 SealSettings
  //       3    2 00001380 UnsealSettings

  // Load and initialize the enclave.

  LPENCLAVE_ROUTINE init_fn =
      (LPENCLAVE_ROUTINE)GetProcAddress((HMODULE)enclave_addr, "Init");
  gle = GetLastError();
  VLOG(1) << "Init addr: " << (LPVOID)init_fn << " gle " << gle;
  if (!init_fn) {
    return 4;
  }
  LPENCLAVE_ROUTINE seal_fn =
      (LPENCLAVE_ROUTINE)GetProcAddress((HMODULE)enclave_addr, "SealSettings");
  gle = GetLastError();
  VLOG(1) << "Seal addr: " << (LPVOID)seal_fn << " gle " << gle;
  if (!seal_fn) {
    return 4;
  }
  LPENCLAVE_ROUTINE laes_fn = (LPENCLAVE_ROUTINE)GetProcAddress(
      (HMODULE)enclave_addr, "UnsealSettings");
  gle = GetLastError();
  VLOG(1) << "laeS addr: " << (LPVOID)laes_fn << " gle " << gle;
  if (!laes_fn) {
    return 4;
  }

  char init_name[] = "testtest";
  PREFS_INIT init_args{};
  init_args.init_name = init_name;

  INIT_TAG llconfig = 0;
  ret = CallEnclave(init_fn, &init_args, TRUE, (void**)&llconfig);

  gle = GetLastError();
  VLOG(1) << "Call Init " << ret << " ll " << llconfig << " gle " << gle;

  // (demonstrates a normal seal/unseal operation - not a bug)
  if (base::CommandLine::ForCurrentProcess()->HasSwitch(switches::kSealData)) {
    // try seal/unseal
    SEAL_ARGS seal_args{};
    seal_args.config_ll = llconfig;
    std::string seal_arg =
        base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
            switches::kSealData);

    seal_args.data_to_seal = (unsigned char*)seal_arg.c_str();
    seal_args.sz_data_to_seal = seal_arg.size();
    seal_args.protectedBlob = nullptr;
    seal_args.sz_protected_blob_size = 0;

    DWORD szNeeded = 0;
    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(0) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> protectedBlob;
    protectedBlob.resize(szNeeded);
    seal_args.protectedBlob = protectedBlob.data();
    seal_args.sz_protected_blob_size = szNeeded;

    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(sz) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> unsealed;
    unsealed.resize(seal_arg.size());

    UNSEAL_ARGS unseal{};
    unseal.config_ll = llconfig;
    unseal.protected_blob = protectedBlob.data();
    unseal.sz_protected_blob = protectedBlob.size();
    unseal.unsealed_size = unsealed.size();
    unseal.unsealed_size_max = unsealed.size();
    unseal.unsealed_data = unsealed.data();

    DWORD intRet = 0;
    ret = CallEnclave(laes_fn, &unseal, TRUE, (void**)&intRet);
    std::string unsealed_string((char*)unsealed.data(), unsealed.size());
    VLOG(1) << "Call Unseal(sz) " << ret << " data: " << unsealed_string;
  }

  // Begin odd things.
  // BUG? llconfig leaks a valid address inside the enclave, so we can write
  // near it.
  //--write-data
  if (base::CommandLine::ForCurrentProcess()->HasSwitch(switches::kWriteData)) {
    std::vector<unsigned char> writeme;
    writeme.push_back(0xde);
    writeme.push_back(0xad);
    writeme.push_back(0xbe);
    writeme.push_back(0xef);
    writeme.push_back(0xde);
    writeme.push_back(0xad);
    writeme.push_back(0xbe);
    writeme.push_back(0xef);

    std::wstring offset_str =
        base::CommandLine::ForCurrentProcess()->GetSwitchValueNative(
            switches::kWriteData);
    size_t offset = 0;
    base::StringToSizeT(offset_str, &offset);

    std::vector<unsigned char> sealed_internal_data;
    SEAL_ARGS seal_args{};
    seal_args.config_ll = llconfig;
    seal_args.data_to_seal = writeme.data();
    seal_args.sz_data_to_seal = writeme.size();
    seal_args.protectedBlob = nullptr;
    seal_args.sz_protected_blob_size = 0;

    DWORD szNeeded = 0;
    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(0) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> protectedBlob;
    protectedBlob.resize(szNeeded);
    seal_args.protectedBlob = protectedBlob.data();
    seal_args.sz_protected_blob_size = szNeeded;

    // This seals the data we want to write to a blob outside the enclave.
    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(sz) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> unsealed;
    unsealed.resize(0x20);

    UNSEAL_ARGS unseal{};
    unseal.config_ll = llconfig;
    unseal.protected_blob = protectedBlob.data();
    unseal.sz_protected_blob = protectedBlob.size();
    unseal.unsealed_size = writeme.size();
    unseal.unsealed_size_max = writeme.size();
    unseal.unsealed_data = (unsigned char*)llconfig + offset;

    VLOG(1) << "Unseal to " << (void*)unseal.unsealed_data;
    DWORD intRet = 0;
    // BUG: we unseal to an address of our chosing in the enclave.
    ret = CallEnclave(laes_fn, &unseal, TRUE, (void**)&intRet);
    VLOG(1) << "Call Unseal(sz) " << ret;
  }

  //--read-enclave=offset
  if (base::CommandLine::ForCurrentProcess()->HasSwitch(switches::kReadData)) {
    // 1. Read from enclave
    std::wstring offset_str =
        base::CommandLine::ForCurrentProcess()->GetSwitchValueNative(
            switches::kReadData);
    size_t offset = 0;
    base::StringToSizeT(offset_str, &offset);

    std::vector<unsigned char> sealed_internal_data;
    SEAL_ARGS seal_args{};
    seal_args.config_ll = llconfig;
    seal_args.data_to_seal = (unsigned char*)llconfig + offset;
    seal_args.sz_data_to_seal = 0x08;
    seal_args.protectedBlob = nullptr;
    seal_args.sz_protected_blob_size = 0;

    DWORD szNeeded = 0;
    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(0) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> protectedBlob;
    protectedBlob.resize(szNeeded);
    seal_args.protectedBlob = protectedBlob.data();
    seal_args.sz_protected_blob_size = szNeeded;

    // BUG - sealing data inside the enclave to a blob outside the enclave.
    VLOG(1) << "Seal from " << (void*)seal_args.data_to_seal;
    ret = CallEnclave(seal_fn, &seal_args, TRUE, (void**)&szNeeded);
    gle = GetLastError();
    VLOG(1) << "Call Seal(sz) " << ret << " sz " << szNeeded << " gle " << gle;

    std::vector<unsigned char> unsealed;
    unsealed.resize(0x08);

    UNSEAL_ARGS unseal{};
    unseal.config_ll = llconfig;
    unseal.protected_blob = protectedBlob.data();
    unseal.sz_protected_blob = protectedBlob.size();
    unseal.unsealed_size = unsealed.size();
    unseal.unsealed_size_max = unsealed.size();
    unseal.unsealed_data = unsealed.data();

    DWORD intRet = 0;
    // unseals the data so we can see it.
    ret = CallEnclave(laes_fn, &unseal, TRUE, (void**)&intRet);
    VLOG(1) << "Call Unseal(sz) " << ret
            << " data: " << base::HexEncode(unsealed);
  }

  return enclave_addr ? 1 : 0;
}