#pragma once

namespace workspace_state {
struct PrivateDacl {
  PSECURITY_DESCRIPTOR descriptor = nullptr;
  PACL acl = nullptr;
  PrivateDacl() {
    Handle token;
    require(OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token.value) != FALSE, "Open workspace seal token");
    DWORD bytes = 0; GetTokenInformation(token.value, TokenUser, nullptr, 0, &bytes);
    std::vector<BYTE> user(bytes);
    require(GetTokenInformation(token.value, TokenUser, user.data(), bytes, &bytes) != FALSE, "Read workspace seal owner");
    LPWSTR sid = nullptr;
    require(ConvertSidToStringSidW(reinterpret_cast<TOKEN_USER*>(user.data())->User.Sid, &sid) != FALSE, "Encode workspace seal owner");
    const std::wstring sddl = L"D:P(A;OICI;FA;;;SY)(A;OICI;FA;;;" + std::wstring(sid) + L")";
    LocalFree(sid);
    require(ConvertStringSecurityDescriptorToSecurityDescriptorW(sddl.c_str(), SDDL_REVISION_1, &descriptor, nullptr) != FALSE, "Build workspace private DACL");
    BOOL present = FALSE, defaulted = FALSE;
    require(GetSecurityDescriptorDacl(descriptor, &present, &acl, &defaulted) != FALSE && present && acl, "Read workspace private DACL");
  }
  ~PrivateDacl() { if (descriptor) LocalFree(descriptor); }
  void apply(HANDLE handle) const {
    security(SetSecurityInfo(handle, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION | PROTECTED_DACL_SECURITY_INFORMATION,
      nullptr, nullptr, acl, nullptr), "Seal workspace handle DACL");
  }
};
static void seal(HANDLE handle, const PrivateDacl& dacl, size_t& count, unsigned int depth = 0) {
  const auto info = workspace_read::information(handle);
  if (++count > 10001 || depth > 64) throw std::runtime_error("Workspace seal entry or nesting limit");
  dacl.apply(handle);
  if (!(info.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) return;
  std::vector<BYTE> buffer(65536);
  for (;;) {
    if (!GetFileInformationByHandleEx(handle, FileIdBothDirectoryInfo, buffer.data(), static_cast<DWORD>(buffer.size()))) {
      if (GetLastError() == ERROR_NO_MORE_FILES) break;
      require(false, "Enumerate workspace seal directory");
    }
    size_t offset = 0;
    for (;;) {
      const auto entry = reinterpret_cast<FILE_ID_BOTH_DIR_INFO*>(buffer.data() + offset);
      if (entry->FileNameLength > buffer.size() - offset - offsetof(FILE_ID_BOTH_DIR_INFO, FileName)) throw std::runtime_error("Invalid workspace seal directory record");
      const std::wstring name(entry->FileName, entry->FileNameLength / sizeof(wchar_t));
      if (name != L"." && name != L"..") {
        if (!workspace_read::component(name) || entry->FileAttributes & FILE_ATTRIBUTE_REPARSE_POINT) throw std::runtime_error("Workspace seal link or invalid name denied");
        Handle child(workspace_write::openRelative(handle, name, GENERIC_READ | READ_CONTROL | WRITE_DAC, false));
        require(child.value != INVALID_HANDLE_VALUE, "Open workspace seal child");
        seal(child.value, dacl, count, depth + 1);
      }
      if (entry->NextEntryOffset == 0) break;
      if (entry->NextEntryOffset < sizeof(FILE_ID_BOTH_DIR_INFO) || entry->NextEntryOffset > buffer.size() - offset - sizeof(FILE_ID_BOTH_DIR_INFO)) throw std::runtime_error("Invalid workspace seal directory offset");
      offset += entry->NextEntryOffset;
    }
  }
}
static int run(const std::wstring& path) {
  if (inContainer(GetCurrentProcess())) throw std::runtime_error("Workspace sealing is host-only");
  workspace_read::Lifetime lifetime(120000);
  workspace_write::Root root(path);
  if (root.held.size() < 2) throw std::runtime_error("A volume root cannot be a managed workspace");
  const auto names = workspace_read::parts(path.substr(3));
  Handle target(workspace_write::openRelative(root.held[root.held.size() - 2]->value, names.back(), GENERIC_READ | READ_CONTROL | WRITE_DAC, true));
  require(target.value != INVALID_HANDLE_VALUE, "Open workspace seal root");
  if (workspace_read::identity(target.value) != root.rootIdentity) throw std::runtime_error("Workspace seal root changed");
  PrivateDacl dacl; size_t count = 0;
  seal(target.value, dacl, count);
  workspace_write::reply("{\"protocol\":1,\"sealed\":true,\"rootIdentity\":\"" + root.rootIdentity + "\",\"entries\":" + std::to_string(count) + "}");
  return 0;
}
}