#ifndef UNICODE
#define UNICODE
#endif
#ifndef _UNICODE
#define _UNICODE
#endif
#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#include <aclapi.h>
#include <sddl.h>
#include <array>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

namespace fs = std::filesystem;
constexpr DWORD ShareAll = FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE;
constexpr DWORD ControlAccess = READ_CONTROL | WRITE_DAC | SYNCHRONIZE | FILE_READ_ATTRIBUTES;

class Handle {
    HANDLE value_ = nullptr;
public:
    explicit Handle(HANDLE value = nullptr) : value_(value) {}
    ~Handle() { if (valid()) CloseHandle(value_); }
    Handle(const Handle&) = delete;
    Handle& operator=(const Handle&) = delete;
    Handle(Handle&& other) noexcept : value_(std::exchange(other.value_, nullptr)) {}
    HANDLE get() const { return value_; }
    bool valid() const { return value_ != nullptr && value_ != INVALID_HANDLE_VALUE; }
};

class LocalMemory {
    void* value_;
public:
    explicit LocalMemory(void* value) : value_(value) {}
    ~LocalMemory() { if (value_) LocalFree(value_); }
    LocalMemory(const LocalMemory&) = delete;
    LocalMemory& operator=(const LocalMemory&) = delete;
};

void require(bool success, const char* api) {
    const DWORD error = success ? ERROR_SUCCESS : GetLastError();
    if (!success) throw std::system_error(static_cast<int>(error), std::system_category(), api);
}

void require_status(DWORD error, const char* api) {
    if (error != ERROR_SUCCESS)
        throw std::system_error(static_cast<int>(error), std::system_category(), api);
}

void expect(bool success, DWORD error, bool wanted, const wchar_t* label, DWORD expectedError = ERROR_ACCESS_DENIED) {
    std::wcout << label << L": " << (success ? L"SUCCESS" : L"FAILURE")
               << L" (Win32 " << error << L")\n";
    if (success != wanted || (!success && error != expectedError))
        throw std::runtime_error("An API result differed from the prediction");
}

std::wstring hex(DWORD value) {
    std::wostringstream text;
    text << L"0x" << std::hex << std::uppercase << std::setw(8) << std::setfill(L'0') << value;
    return text.str();
}

void granted_access(HANDLE file, const wchar_t* label) {
    struct BasicInformation {
        ULONG Attributes;
        ACCESS_MASK GrantedAccess;
        ULONG HandleCount;
        ULONG PointerCount;
        ULONG Reserved[10];
    } information{};
    using Query = LONG (NTAPI*)(HANDLE, int, void*, ULONG, ULONG*);
    const FARPROC address = GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "NtQueryObject");
    Query query = nullptr;
    static_assert(sizeof(query) == sizeof(address));
    std::memcpy(&query, &address, sizeof(query));
    if (!query) {
        std::wcout << label << L": NtQueryObject unavailable; no measured mask\n";
        return;
    }
    ULONG returned = 0;
    const LONG status = query(file, 0, &information, sizeof(information), &returned);
    std::wcout << label << L": handle=" << file << L"; NTSTATUS=" << hex(static_cast<DWORD>(status));
    if (status >= 0) std::wcout << L"; measured GrantedAccess=" << hex(information.GrantedAccess);
    std::wcout << L"\n";
}

Handle open_file(const fs::path& path, DWORD access, const wchar_t* label, bool wanted = true) {
    HANDLE raw = CreateFileW(path.c_str(), access, ShareAll, nullptr, OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL, nullptr);
    const DWORD error = raw == INVALID_HANDLE_VALUE ? GetLastError() : ERROR_SUCCESS;
    Handle result(raw);
    std::wcout << L"Requested " << hex(access) << L"; ";
    expect(result.valid(), error, wanted, label);
    if (result.valid()) granted_access(result.get(), label);
    return result;
}

void seek_start(HANDLE file) {
    LARGE_INTEGER zero{};
    require(SetFilePointerEx(file, zero, nullptr, FILE_BEGIN) != FALSE, "SetFilePointerEx");
}

void write_marker(HANDLE file, const std::string& marker, const wchar_t* label) {
    if (marker.size() > 64) throw std::runtime_error("Marker too large");
    const std::string data = marker + std::string(64 - marker.size(), ' ');
    seek_start(file);
    DWORD written = 0;
    const BOOL success = WriteFile(file, data.data(), static_cast<DWORD>(data.size()), &written, nullptr);
    const DWORD error = success ? ERROR_SUCCESS : GetLastError();
    expect(success != FALSE, error, true, label);
    if (written != data.size()) throw std::runtime_error("Short WriteFile");
}

std::string read_marker(HANDLE file, const wchar_t* label) {
    seek_start(file);
    std::array<char, 64> data{};
    DWORD read = 0;
    const BOOL success = ReadFile(file, data.data(), static_cast<DWORD>(data.size()), &read, nullptr);
    const DWORD error = success ? ERROR_SUCCESS : GetLastError();
    expect(success != FALSE, error, true, label);
    std::string text(data.data(), read);
    while (!text.empty() && text.back() == ' ') text.pop_back();
    std::wcout << L"    Marker: " << std::wstring(text.begin(), text.end()) << L"\n";
    return text;
}

std::wstring user_sid(HANDLE token) {
    DWORD size = 0;
    if (GetTokenInformation(token, TokenUser, nullptr, 0, &size) || GetLastError() != ERROR_INSUFFICIENT_BUFFER)
        throw std::runtime_error("Unexpected TokenUser size query result");
    std::vector<unsigned char> buffer(size);
    require(GetTokenInformation(token, TokenUser, buffer.data(), size, &size) != FALSE, "TokenUser");
    LPWSTR text = nullptr;
    require(ConvertSidToStringSidW(reinterpret_cast<TOKEN_USER*>(buffer.data())->User.Sid, &text) != FALSE,
        "ConvertSidToStringSidW");
    LocalMemory allocation(text);
    return text;
}

Handle process_token() {
    HANDLE raw = nullptr;
    require(OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_DUPLICATE, &raw) != FALSE, "OpenProcessToken");
    Handle token(raw);
    TOKEN_ELEVATION elevation{};
    DWORD size = 0;
    require(GetTokenInformation(token.get(), TokenElevation, &elevation, sizeof(elevation), &size) != FALSE, "TokenElevation");
    if (elevation.TokenIsElevated) throw std::runtime_error("Use an unelevated console for the lab");
    HANDLE threadRaw = nullptr;
    const BOOL hasThreadToken = OpenThreadToken(GetCurrentThread(), TOKEN_QUERY, TRUE, &threadRaw);
    const DWORD threadError = hasThreadToken ? ERROR_SUCCESS : GetLastError();
    Handle threadToken(threadRaw);
    if (hasThreadToken || threadError != ERROR_NO_TOKEN)
        throw std::runtime_error("Start without thread impersonation");
    return token;
}

std::wstring dacl_text(PSECURITY_DESCRIPTOR descriptor) {
    LPWSTR text = nullptr;
    require(ConvertSecurityDescriptorToStringSecurityDescriptorW(descriptor, SDDL_REVISION_1,
        DACL_SECURITY_INFORMATION, &text, nullptr) != FALSE, "DACL to SDDL");
    LocalMemory allocation(text);
    return text;
}

std::string utf8(const std::wstring& text) {
    if (text.size() > static_cast<size_t>(std::numeric_limits<int>::max()))
        throw std::runtime_error("Descriptor text too long");
    const int length = static_cast<int>(text.size());
    const int size = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, text.data(), length,
        nullptr, 0, nullptr, nullptr);
    require(size != 0, "Size UTF-8 descriptor");
    std::string encoded(static_cast<size_t>(size), '\0');
    require(WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, text.data(), length,
        encoded.data(), size, nullptr, nullptr) != 0, "Encode UTF-8 descriptor");
    return encoded;
}

class DaclGuard {
    HANDLE file_;
    PSECURITY_DESCRIPTOR saved_ = nullptr;
    PACL dacl_ = nullptr;
    SECURITY_DESCRIPTOR_CONTROL control_ = 0;
    bool restored_ = false;
public:
    DaclGuard(HANDLE file, const fs::path& path) : file_(file) {
        require_status(GetSecurityInfo(file_, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION,
            nullptr, nullptr, &dacl_, nullptr, &saved_), "GetSecurityInfo");
        try {
            DWORD revision = 0;
            require(GetSecurityDescriptorControl(saved_, &control_, &revision) != FALSE, "DACL control");
            const std::wstring text = dacl_text(saved_);
            fs::path backup = path;
            backup += L".original-dacl.sddl";
            std::ofstream output(backup, std::ios::binary);
            output << utf8(text);
            if (!output) throw std::runtime_error("Could not save original DACL");
        } catch (...) {
            LocalFree(saved_);
            saved_ = nullptr;
            throw;
        }
    }
    ~DaclGuard() {
        if (!restored_) {
            try { restore(); }
            catch (const std::exception& error) { std::cerr << "DACL RESTORATION FAILED: " << error.what() << "\n"; }
        }
        if (saved_) LocalFree(saved_);
    }
    DaclGuard(const DaclGuard&) = delete;
    DaclGuard& operator=(const DaclGuard&) = delete;
    void restore() {
        if (restored_) return;
        const SECURITY_INFORMATION protection = (control_ & SE_DACL_PROTECTED)
            ? PROTECTED_DACL_SECURITY_INFORMATION : UNPROTECTED_DACL_SECURITY_INFORMATION;
        require_status(SetSecurityInfo(file_, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION | protection,
            nullptr, nullptr, dacl_, nullptr), "Restore original DACL");
        PSECURITY_DESCRIPTOR current = nullptr;
        require_status(GetSecurityInfo(file_, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION,
            nullptr, nullptr, nullptr, nullptr, &current), "Verify restored DACL");
        LocalMemory allocation(current);
        if (dacl_text(current) != dacl_text(saved_)) throw std::runtime_error("Restored DACL differs from saved SDDL");
        restored_ = true;
        std::wcout << L"Original DACL restored and verified.\n";
    }
};

void set_policy(HANDLE control, const std::wstring& sid, DWORD dataAccess) {
    const std::wstring sddl = L"D:P(A;;" + hex(ControlAccess | dataAccess) + L";;;" + sid + L")";
    PSECURITY_DESCRIPTOR descriptor = nullptr;
    require(ConvertStringSecurityDescriptorToSecurityDescriptorW(sddl.c_str(), SDDL_REVISION_1,
        &descriptor, nullptr) != FALSE, "Parse experimental SDDL");
    LocalMemory allocation(descriptor);
    BOOL present = FALSE, defaulted = FALSE;
    PACL dacl = nullptr;
    require(GetSecurityDescriptorDacl(descriptor, &present, &dacl, &defaulted) != FALSE, "Extract DACL");
    if (!present || !dacl) throw std::runtime_error("Refusing a null DACL");
    require_status(SetSecurityInfo(control, SE_FILE_OBJECT,
        DACL_SECURITY_INFORMATION | PROTECTED_DACL_SECURITY_INFORMATION,
        nullptr, nullptr, dacl, nullptr), "SetSecurityInfo");
    std::wcout << L"SetSecurityInfo: " << sddl << L"\n";
}

std::wstring file_identity(HANDLE file) {
    BY_HANDLE_FILE_INFORMATION information{};
    require(GetFileInformationByHandle(file, &information) != FALSE, "File identity");
    std::wostringstream text;
    text << std::hex << std::uppercase << information.dwVolumeSerialNumber << L":"
         << information.nFileIndexHigh << L":" << information.nFileIndexLow;
    return text.str();
}

void checkpoint(bool pause, const wchar_t* text) {
    std::wcout << text << L"\n";
    if (pause) {
        std::wcout << L"Press Enter to continue.\n";
        std::wstring line;
        std::getline(std::wcin, line);
    }
}

class ImpersonationGuard {
public:
    explicit ImpersonationGuard(HANDLE token) {
        require(ImpersonateLoggedOnUser(token) != FALSE, "ImpersonateLoggedOnUser");
    }
    ~ImpersonationGuard() {
        if (!RevertToSelf()) {
            const DWORD error = GetLastError();
            std::wcerr << L"RevertToSelf failed; terminating this lab process.\n";
            TerminateProcess(GetCurrentProcess(), error);
            std::terminate();
        }
    }
    ImpersonationGuard(const ImpersonationGuard&) = delete;
    ImpersonationGuard& operator=(const ImpersonationGuard&) = delete;
};

fs::path new_run_directory() {
    std::array<wchar_t, 32768> local{};
    const DWORD length = GetEnvironmentVariableW(L"LOCALAPPDATA", local.data(), static_cast<DWORD>(local.size()));
    if (!length || length >= local.size()) throw std::runtime_error("LOCALAPPDATA unavailable or too long");
    const fs::path base(local.data());
    std::array<wchar_t, 32768> volume{};
    require(GetVolumePathNameW(base.c_str(), volume.data(), static_cast<DWORD>(volume.size())) != FALSE, "GetVolumePathNameW");
    std::array<wchar_t, 32> format{};
    require(GetVolumeInformationW(volume.data(), nullptr, 0, nullptr, nullptr, nullptr,
        format.data(), static_cast<DWORD>(format.size())) != FALSE, "GetVolumeInformationW");
    if (GetDriveTypeW(volume.data()) != DRIVE_FIXED || std::wstring(format.data()) != L"NTFS")
        throw std::runtime_error("Use fixed local NTFS storage");
    fs::path root = base / L"WindowsHandleLab" / L"data";
    for (fs::path ancestor = root; !ancestor.empty(); ancestor = ancestor.parent_path()) {
        const DWORD attributes = GetFileAttributesW(ancestor.c_str());
        if (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_REPARSE_POINT))
            throw std::runtime_error("Reparse points are outside the lab scope");
        if (ancestor == ancestor.parent_path()) break;
    }
    fs::create_directories(root);
    const fs::path run = root / (L"cpp-" + std::to_wstring(GetCurrentProcessId()) + L"-" + std::to_wstring(GetTickCount64()));
    if (!fs::create_directory(run)) throw std::runtime_error("Run directory already exists");
    return run;
}

int wmain(int argc, wchar_t** argv) {
    bool pause = false, impersonation = false;
    for (int i = 1; i < argc; ++i) {
        const std::wstring argument(argv[i]);
        if (argument == L"--pause") pause = true;
        else if (argument == L"--impersonation") impersonation = true;
        else if (argument == L"--help") {
            std::wcout << L"HandleLab.exe [--pause] [--impersonation]\n"
                << L"Creates a new file under LOCALAPPDATA\\WindowsHandleLab\\data; restores its original DACL.\n";
            return 0;
        } else { std::wcerr << L"Unknown option: " << argument << L"\n"; return 2; }
    }
    try {
        Handle token = process_token();
        const std::wstring sid = user_sid(token.get());
        const fs::path run = new_run_directory();
        const fs::path path = run / L"native-experiment.txt";
        std::wcout << L"PID: " << GetCurrentProcessId() << L"; user SID: " << sid
            << L"\nFile: " << path.c_str() << L"\n";
        checkpoint(pause, L"Start observation before creating the file.");
        {
            Handle seed(CreateFileW(path.c_str(), GENERIC_READ | GENERIC_WRITE, ShareAll, nullptr,
                CREATE_NEW, FILE_ATTRIBUTE_NORMAL, nullptr));
            require(seed.valid(), "CREATE_NEW disposable file");
            write_marker(seed.get(), "NATIVE-BEFORE-DACL-CHANGE", L"Seed WriteFile");
        }
        Handle control = open_file(path, ControlAccess, L"control handle without data rights");
        DaclGuard restore(control.get(), path);
        const std::wstring identity = file_identity(control.get());
        set_policy(control.get(), sid, FILE_GENERIC_READ | FILE_GENERIC_WRITE);
        {
            Handle reader = open_file(path, GENERIC_READ, L"retained reader");
            Handle writer = open_file(path, GENERIC_WRITE, L"retained writer");
            read_marker(reader.get(), L"ReadFile before DACL change");
            set_policy(control.get(), sid, 0);
            checkpoint(pause, L"DACL now allows control rights but no file data access.");
            const std::string marker = "NATIVE-AFTER-" + std::to_string(GetTickCount64());
            write_marker(writer.get(), marker, L"WriteFile through retained writer after DACL change");
            require(FlushFileBuffers(writer.get()) != FALSE, "FlushFileBuffers");
            if (read_marker(reader.get(), L"fresh ReadFile through retained reader") != marker)
                throw std::runtime_error("Post-change marker does not match");
            if (file_identity(reader.get()) != identity || file_identity(writer.get()) != identity)
                throw std::runtime_error("File identity changed");
            granted_access(reader.get(), L"reader mask after DACL change");
            granted_access(writer.get(), L"writer mask after DACL change");
            Handle deniedRead = open_file(path, GENERIC_READ, L"fresh reader after DACL change", false);
            Handle deniedWrite = open_file(path, GENERIC_WRITE, L"fresh writer after DACL change", false);
        }
        if (impersonation) {
            set_policy(control.get(), sid, FILE_GENERIC_READ | FILE_GENERIC_WRITE);
            Handle old = open_file(path, GENERIC_READ | GENERIC_WRITE, L"handle acquired before impersonation");
            std::array<unsigned char, SECURITY_MAX_SID_SIZE> nullSid{};
            DWORD size = static_cast<DWORD>(nullSid.size());
            require(CreateWellKnownSid(WinNullSid, nullptr, nullSid.data(), &size) != FALSE, "CreateWellKnownSid");
            SID_AND_ATTRIBUTES restriction{nullSid.data(), 0};
            HANDLE restrictedRaw = nullptr;
            require(CreateRestrictedToken(token.get(), DISABLE_MAX_PRIVILEGE, 0, nullptr,
                0, nullptr, 1, &restriction, &restrictedRaw) != FALSE, "CreateRestrictedToken");
            Handle restricted(restrictedRaw);
            if (!IsTokenRestricted(restricted.get()) || user_sid(restricted.get()) != sid)
                throw std::runtime_error("Restricted token does not match the experiment assumptions");
            checkpoint(pause, L"Next: impersonate a token with the same user SID and S-1-0-0 as a restricting SID.");
            {
                ImpersonationGuard revert(restricted.get());
                Handle denied = open_file(path, GENERIC_READ, L"fresh open under restricted impersonation", false);
                write_marker(old.get(), "WRITE-WHILE-IMPERSONATING", L"retained handle writes while impersonating");
                read_marker(old.get(), L"retained handle reads while impersonating");
            }
            Handle fresh = open_file(path, GENERIC_READ, L"fresh open after RevertToSelf");
            std::wcout << L"Only the thread's token changed in the impersonation experiment.\n";
        }
        restore.restore();
        std::wcout << L"All native checks matched their predictions. Disposable file retained for inspection.\n";
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "Lab failed: " << error.what() << "\n";
        return 1;
    }
}