#pragma once
#include "acceptance-state.h"
#include <wincrypt.h>
#pragma comment(lib, "crypt32.lib")
namespace credential_contract {
using release_contract::Json;
static constexpr DWORD STORE_LIMIT = 1048576;
struct Wipe {
  std::vector<BYTE> bytes;
  ~Wipe() {
    if (!bytes.empty())
      SecureZeroMemory(bytes.data(), bytes.size());
  }
};
struct WipeText {
  std::string &text;
  ~WipeText() {
    if (!text.empty())
      SecureZeroMemory(text.data(), text.size());
  }
};
static Json parse(const std::string &text) {
  release_contract::Parser parser{text};
  auto result = parser.parse();
  if (parser.offset != text.size())
    throw 1;
  return result;
}
static bool tokenValid(const std::string &token) {
  if (token.size() < 16 || token.size() > 4096)
    return false;
  for (char c : token)
    if (!((c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') ||
          (c >= '0' && c <= '9') || c == '.' || c == '_' || c == '~' ||
          c == '-'))
      return false;
  return true;
}
static Json binding(const Json &request,
                    const release_contract::Assignment &assignment,
                    const Json &state) {
  if (assignment.managementOrigin.empty() ||
      release_contract::str(request.object.at("deviceId")) !=
          assignment.deviceId ||
      release_contract::str(request.object.at("origin")) !=
          assignment.managementOrigin)
    throw 1;
  auto der = release_contract::pemDer(assignment.policyPem);
  const std::string keyBytes(der.begin(), der.end());
  if (release_contract::str(request.object.at("policyKeySha256")) !=
      release_contract::digest(keyBytes))
    throw 1;
  const auto version =
      release_contract::number(request.object.at("policyVersion"));
  const auto digest = release_contract::str(request.object.at("policyDigest"));
  if (!acceptance_contract::sha(digest) || version < assignment.minimumVersion)
    throw 1;
  bool matched = false;
  for (const auto &record : state.object.at("policies").array) {
    if (release_contract::str(record.object.at("sid")) == assignment.sid &&
        release_contract::str(record.object.at("deviceId")) ==
            assignment.deviceId &&
        release_contract::str(record.object.at("policyKeySha256")) ==
            release_contract::str(request.object.at("policyKeySha256"))) {
      matched =
          release_contract::number(record.object.at("minimumVersion")) ==
              version &&
          release_contract::str(record.object.at("acceptedDigest")) == digest;
    }
  }
  if (!matched)
    throw 1;
  Json result = request;
  result.object.erase("token");
  result.object.erase("expectedPreviousSHA");
  Json sid;
  sid.kind = Json::String;
  sid.string = assignment.sid;
  result.object["sid"] = sid;
  return result;
}
static Json validateStore(const std::string &text) {
  if (text.empty() || text.size() > STORE_LIMIT)
    throw 1;
  auto value = parse(text);
  if (release_contract::canonical(value) != text ||
      !release_contract::exact(value, {"format", "records"}) ||
      release_contract::str(value.object.at("format")) !=
          "safetynet-management-credentials-v1" ||
      value.object.at("records").kind != Json::Array ||
      value.object.at("records").array.size() > 256)
    throw 1;
  std::string previous;
  for (const auto &record : value.object.at("records").array) {
    if (!release_contract::exact(record, {"sid", "deviceId", "origin",
                                          "policyVersion", "policyKeySha256",
                                          "policyDigest", "protectedToken"}) ||
        !acceptance_contract::userSid(
            release_contract::str(record.object.at("sid"))) ||
        !release_contract::managementOrigin(
            release_contract::str(record.object.at("origin"))) ||
        !acceptance_contract::sha(
            release_contract::str(record.object.at("policyKeySha256"))) ||
        !acceptance_contract::sha(
            release_contract::str(record.object.at("policyDigest"))))
      throw 1;
    release_contract::number(record.object.at("policyVersion"));
    const auto device = release_contract::str(record.object.at("deviceId"));
    if (device.empty() || device.size() > 80 ||
        !isalnum(static_cast<unsigned char>(device[0])))
      throw 1;
    for (char c : device)
      if (!isalnum(static_cast<unsigned char>(c)) && c != '.' && c != '_' &&
          c != '-')
        throw 1;
    const auto cipher =
        release_contract::str(record.object.at("protectedToken"));
    if (cipher.empty() || cipher.size() > 16384)
      throw 1;
    release_contract::decode(cipher);
    Json identity = record;
    identity.object.erase("protectedToken");
    const auto key = release_contract::canonical(identity);
    if (!previous.empty() && key <= previous)
      throw 1;
    previous = key;
  }
  return value;
}
static std::string stableIdentity(Json metadata) {
  metadata.object.erase("policyVersion");
  metadata.object.erase("policyDigest");
  metadata.object.erase("protectedToken");
  return release_contract::canonical(metadata);
}
static const Json *selectRecord(const Json &metadata,
                                const std::vector<Json> &records) {
  const auto identity = stableIdentity(metadata);
  const auto current =
      release_contract::number(metadata.object.at("policyVersion"));
  const Json *selected = nullptr;
  ULONGLONG highest = 0;
  std::set<ULONGLONG> versions;
  for (const auto &record : records) {
    if (stableIdentity(record) != identity)
      continue;
    const auto version =
        release_contract::number(record.object.at("policyVersion"));
    if (!versions.insert(version).second)
      throw 1;
    if (version > current)
      throw 1;
    if (version == current &&
        release_contract::str(record.object.at("policyDigest")) !=
            release_contract::str(metadata.object.at("policyDigest")))
      throw 1;
    if (version > highest) {
      highest = version;
      selected = &record;
    }
  }
  return selected;
}
static void replaceRecord(std::vector<Json> &records, const Json &record) {
  const auto identity = stableIdentity(record);
  records.erase(std::remove_if(records.begin(), records.end(),
                               [&](const Json &old) {
                                 return stableIdentity(old) == identity;
                               }),
                records.end());
  records.push_back(record);
  std::sort(records.begin(), records.end(), [](Json a, Json b) {
    a.object.erase("protectedToken");
    b.object.erase("protectedToken");
    return release_contract::canonical(a) < release_contract::canonical(b);
  });
}
#ifdef SAFETYNET_NATIVE_SELF_TEST
static bool selectionSelfTest() {
  auto current =
      parse("{\"deviceId\":\"synthetic\",\"origin\":\"https://"
            "management.example\",\"policyDigest\":\"" +
            std::string(64, 'b') + "\",\"policyKeySha256\":\"" +
            std::string(64, 'a') +
            "\",\"policyVersion\":8,\"sid\":\"S-1-5-21-1-2-3-1001\"}");
  auto old = current;
  old.object.at("policyVersion").integer = 7;
  old.object.at("policyDigest").string = std::string(64, 'c');
  auto oldest = old;
  oldest.object.at("policyVersion").integer = 6;
  std::vector<Json> records{oldest, old};
  if (selectRecord(current, records) != &records[1])
    return false;
  auto other = old;
  other.object.at("origin").string = "https://other.example";
  records.push_back(other);
  replaceRecord(records, current);
  if (records.size() != 2 || !selectRecord(current, records))
    return false;
  auto fails = [&](std::vector<Json> rows) {
    try {
      selectRecord(current, rows);
      return false;
    } catch (...) {
      return true;
    }
  };
  auto conflicting = current;
  conflicting.object.at("policyDigest").string = std::string(64, 'd');
  auto future = current;
  future.object.at("policyVersion").integer = 9;
  if (!fails({old, old}) || !fails({conflicting}) || !fails({future}))
    return false;
  auto wrongKey = old;
  wrongKey.object.at("policyKeySha256").string = std::string(64, 'e');
  return !selectRecord(current, {other, wrongKey});
}
#endif
static std::string load(HKEY key) {
  DWORD type = 0, size = 0;
  if (RegQueryValueExW(key, L"ManagementCredentials", nullptr, &type, nullptr,
                       &size) ||
      type != REG_BINARY || !size || size > STORE_LIMIT)
    throw 1;
  std::string value(size, '\0');
  if (RegQueryValueExW(key, L"ManagementCredentials", nullptr, &type,
                       reinterpret_cast<BYTE *>(value.data()), &size) ||
      type != REG_BINARY || size != value.size())
    throw 1;
  validateStore(value);
  return value;
}
static std::vector<BYTE> entropy(const Json &metadata) {
  const auto value =
      std::string("SafetyNet management credential v1\0",
                  sizeof("SafetyNet management credential v1\0") - 1) +
      release_contract::canonical(metadata);
  auto digest = release_contract::unhex(release_contract::digest(value));
  return digest;
}
static std::string protect(const std::string &token, const Json &metadata) {
  if (!tokenValid(token))
    throw 1;
  auto bindingBytes = entropy(metadata);
  Wipe plain;
  plain.bytes.assign(token.begin(), token.end());
  DATA_BLOB input{static_cast<DWORD>(plain.bytes.size()), plain.bytes.data()},
      additional{static_cast<DWORD>(bindingBytes.size()), bindingBytes.data()},
      output{};
  if (!CryptProtectData(
          &input, L"SafetyNet credential", &additional, nullptr, nullptr,
          CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN, &output))
    throw 1;
  std::vector<BYTE> cipher(output.pbData, output.pbData + output.cbData);
  LocalFree(output.pbData);
  return release_contract::base64url(cipher);
}
static std::string unprotect(const std::string &cipher, const Json &metadata) {
  auto encrypted = release_contract::decode(cipher),
       bindingBytes = entropy(metadata);
  DATA_BLOB input{static_cast<DWORD>(encrypted.size()), encrypted.data()},
      additional{static_cast<DWORD>(bindingBytes.size()), bindingBytes.data()},
      output{};
  if (!CryptUnprotectData(&input, nullptr, &additional, nullptr, nullptr,
                          CRYPTPROTECT_UI_FORBIDDEN, &output))
    throw 1;
  std::string token(reinterpret_cast<char *>(output.pbData), output.cbData);
  SecureZeroMemory(output.pbData, output.cbData);
  LocalFree(output.pbData);
  if (!tokenValid(token)) {
    SecureZeroMemory(token.data(), token.size());
    throw 1;
  }
  return token;
}
} // namespace credential_contract