#include <windows.h>

#include "sddl.h"
#include <string>
#include <assert.h>
#include <WtsApi32.h>
#include <windows.h>
#include <shlwapi.h>
#include <strsafe.h>
#include <codecvt>
#include <shlobj.h> 

#pragma comment(lib, "wtsapi32.lib")
#pragma comment(lib, "Shlwapi.lib")

class MyDataObject : public IDataObject {
public:
    // IUnknown
    STDMETHODIMP QueryInterface(REFIID riid, void** ppvObj);
    STDMETHODIMP_(ULONG) AddRef();
    STDMETHODIMP_(ULONG) Release();

    // IDataObject
    STDMETHODIMP GetData(FORMATETC* pfe, STGMEDIUM* pmed);
    STDMETHODIMP GetDataHere(FORMATETC* pfe, STGMEDIUM* pmed);
    STDMETHODIMP QueryGetData(FORMATETC* pfe);
    STDMETHODIMP GetCanonicalFormatEtc(FORMATETC* pfeIn,
        FORMATETC* pfeOut);
    STDMETHODIMP SetData(FORMATETC* pfe, STGMEDIUM* pmed,
        BOOL fRelease);
    STDMETHODIMP EnumFormatEtc(DWORD dwDirection,
        LPENUMFORMATETC* ppefe);
    STDMETHODIMP DAdvise(FORMATETC* pfe, DWORD grfAdv,
        IAdviseSink* pAdvSink, DWORD* pdwConnection);
    STDMETHODIMP DUnadvise(DWORD dwConnection);
    STDMETHODIMP EnumDAdvise(LPENUMSTATDATA* ppefe);

    MyDataObject();

private:
    enum {
        DATA_FILEGROUPDESCRIPTOR,
        DATA_FILECONTENTS,
        DATA_NUM,
        DATA_INVALID = -1
    };

    int GetDataIndex(const FORMATETC* pfe);
private:
    ULONG refCnt_;
    FORMATETC m_rgfe[DATA_NUM];
};

void SetFORMATETC(FORMATETC* pfe, UINT cf, TYMED tymed = TYMED_HGLOBAL, LONG lindex = -1,
    DWORD dwAspect = DVASPECT_CONTENT, DVTARGETDEVICE* ptd = NULL) {
    pfe->cfFormat = (CLIPFORMAT)cf;
    pfe->tymed = tymed;
    pfe->lindex = lindex;
    pfe->dwAspect = dwAspect;
    pfe->ptd = ptd;
}

MyDataObject::MyDataObject() : refCnt_(1) {
    SetFORMATETC(&m_rgfe[DATA_FILEGROUPDESCRIPTOR],
        RegisterClipboardFormat(CFSTR_FILEDESCRIPTOR));
    SetFORMATETC(&m_rgfe[DATA_FILECONTENTS],
        RegisterClipboardFormat(CFSTR_FILECONTENTS),
        TYMED_ISTREAM, /* lindex */ 0);
}

HRESULT MyDataObject::QueryInterface(REFIID riid, void** ppv) {
    IUnknown* punk = NULL;
    if (riid == IID_IUnknown) {
        punk = static_cast<IUnknown*>(this);
    } else if (riid == IID_IDataObject) {
        punk = static_cast<IDataObject*>(this);
    }

    *ppv = punk;
    if (punk) {
        punk->AddRef();
        return S_OK;
    } else {
        return E_NOINTERFACE;
    }
}

ULONG MyDataObject::AddRef() {
    return ++refCnt_;
}

ULONG MyDataObject::Release() {
    ULONG cRef = --refCnt_;
    if (cRef == 0)
        delete this;
    return cRef;
}

int MyDataObject::GetDataIndex(const FORMATETC* pfe) {
    for (int i = 0; i < ARRAYSIZE(m_rgfe); i++) {
        if (pfe->cfFormat == m_rgfe[i].cfFormat &&
            (pfe->tymed & m_rgfe[i].tymed) &&
            pfe->dwAspect == m_rgfe[i].dwAspect &&
            pfe->lindex == m_rgfe[i].lindex) {
            return i;
        }
    }
    return DATA_INVALID;
}

HRESULT CreateHGlobalFromBlob(const void* pvData, SIZE_T cbData, UINT uFlags, HGLOBAL* phglob) {
    HGLOBAL hglob = GlobalAlloc(uFlags, cbData);
    if (hglob) {
        void* pvAlloc = GlobalLock(hglob);
        if (pvAlloc) {
            CopyMemory(pvAlloc, pvData, cbData);
            GlobalUnlock(hglob);
        }
        else {
            GlobalFree(hglob);
            hglob = NULL;
        }
    }
    *phglob = hglob;
    return hglob ? S_OK : E_OUTOFMEMORY;
}

HRESULT MyDataObject::QueryGetData(FORMATETC* pfe) {
    return GetDataIndex(pfe) == DATA_INVALID ? S_FALSE : S_OK;
}

HRESULT MyDataObject::EnumFormatEtc(DWORD dwDirection, LPENUMFORMATETC* ppefe) {
    if (dwDirection == DATADIR_GET) {
        return SHCreateStdEnumFmtEtc(ARRAYSIZE(m_rgfe), m_rgfe, ppefe);
    }
    *ppefe = NULL;
    return E_NOTIMPL;
}

HRESULT MyDataObject::GetDataHere(FORMATETC* pfe, STGMEDIUM* pmed) {
    return E_NOTIMPL;
}

HRESULT MyDataObject::GetCanonicalFormatEtc(FORMATETC* pfeIn, FORMATETC* pfeOut) {
    *pfeOut = *pfeIn;
    pfeOut->ptd = NULL;
    return DATA_S_SAMEFORMATETC;
}

HRESULT MyDataObject::SetData(FORMATETC* pfe, STGMEDIUM* pmed, BOOL fRelease) {
    return E_NOTIMPL;
}

HRESULT MyDataObject::DAdvise(FORMATETC* pfe, DWORD grfAdv, IAdviseSink* pAdvSink, DWORD* pdwConnection) {
    return OLE_E_ADVISENOTSUPPORTED;
}

HRESULT MyDataObject::DUnadvise(DWORD dwConnection) {
    return OLE_E_ADVISENOTSUPPORTED;
}

HRESULT MyDataObject::EnumDAdvise(LPENUMSTATDATA* ppefe) {
    return OLE_E_ADVISENOTSUPPORTED;
}

HRESULT MyDataObject::GetData(FORMATETC* pfe, STGMEDIUM* pmed) {
    ZeroMemory(pmed, sizeof(*pmed));
    if (GetDataIndex(pfe) != -1)
        printf("IDataObject::GetData call with format index %d\n", GetDataIndex(pfe));
    switch (GetDataIndex(pfe)) {
    case DATA_FILEGROUPDESCRIPTOR: {
        FILEGROUPDESCRIPTOR fgd;
        ZeroMemory(&fgd, sizeof(fgd));
        fgd.cItems = 1;
        StringCchCopy(fgd.fgd[0].cFileName, ARRAYSIZE(fgd.fgd[0].cFileName), TEXT("Dummy"));
        pmed->tymed = TYMED_HGLOBAL;
        return CreateHGlobalFromBlob(&fgd, sizeof(fgd),
            GMEM_MOVEABLE, &pmed->hGlobal);
    }
    case DATA_FILECONTENTS:
        pmed->tymed = TYMED_ISTREAM;
        pmed->pstm = SHOpenRegStream(HKEY_LOCAL_MACHINE,
            TEXT("Hardware\\Description\\System\\CentralProcessor\\0"),
            TEXT("~MHz"), STGM_READ);
        if (pmed->pstm) {
            LARGE_INTEGER liZero = { 0, 0 };
            pmed->pstm->Seek(liZero, STREAM_SEEK_END, NULL);
        }
        return pmed->pstm ? S_OK : E_FAIL;
    }
    return DV_E_FORMATETC;
}

BOOL IsSystemSid(PSID sid) {
    return ::IsWellKnownSid(sid, WinLocalSystemSid);
}

HANDLE OpenSystemProcessToken() {
    PWTS_PROCESS_INFO pInfo;
    DWORD count;
    if (!::WTSEnumerateProcesses(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pInfo, &count)) {
        printf("Error enumerating processes (error=%ld)\n", ::GetLastError());
        return nullptr;
    }

    HANDLE hToken{};
    for (DWORD i = 0; i < count && !hToken; i++) {
        if (pInfo[i].SessionId == 0 && IsSystemSid(pInfo[i].pUserSid)) {
            auto hProcess = ::OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pInfo[i].ProcessId);
            if (hProcess) {
                ::OpenProcessToken(hProcess, TOKEN_DUPLICATE | TOKEN_ASSIGN_PRIMARY | TOKEN_QUERY | TOKEN_IMPERSONATE, &hToken);
                ::CloseHandle(hProcess);
            }
        }
    }

    ::WTSFreeMemory(pInfo);
    return hToken;
}

BOOL SetPrivilege(HANDLE hToken, PCTSTR lpszPrivilege, bool bEnablePrivilege) {
    TOKEN_PRIVILEGES tp;
    LUID luid;

    if (!::LookupPrivilegeValue(nullptr, lpszPrivilege, &luid))
        return FALSE;

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    if (bEnablePrivilege)
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    else
        tp.Privileges[0].Attributes = 0;

    // Enable the privilege or disable all privileges.
    if (!::AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr)) {
        return FALSE;
    }

    if (::GetLastError() == ERROR_NOT_ALL_ASSIGNED)
        return FALSE;

    return TRUE;
}

BOOL EnableDebugPrivilege(void) {
    HANDLE hToken;
    BOOL result;
    if (!::OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        return FALSE;
    }
    result = SetPrivilege(hToken, SE_DEBUG_NAME, TRUE);
    ::CloseHandle(hToken);
    return result;
}

std::wstring string_to_wstring(const std::string& str) {
    using convert_typeX = std::codecvt_utf8<wchar_t>;
    std::wstring_convert<convert_typeX, wchar_t> converterX;
    return converterX.from_bytes(str);
}

//https://github.com/zodiacon/sysrun
int runAsSystem(int argc, const char* argv[]) {
    if (FALSE == EnableDebugPrivilege()) {
        printf("EnableDebugPrivilege failed: Access denied\n");
        return -1;
    }

    auto hToken = OpenSystemProcessToken();
    if (!hToken) {
        printf("OpenSystemProcessToken failed: Access denied\n");
        return -1;
    }

    HANDLE hDupToken, hPrimary;
    ::DuplicateTokenEx(hToken, TOKEN_DUPLICATE | TOKEN_IMPERSONATE | TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY | TOKEN_ADJUST_PRIVILEGES,
        nullptr, SecurityImpersonation, TokenImpersonation, &hDupToken);
    ::DuplicateTokenEx(hToken, TOKEN_ALL_ACCESS, nullptr, SecurityImpersonation, TokenPrimary, &hPrimary);
    ::CloseHandle(hToken);

    if (hDupToken == nullptr) {
        printf("Failed to create token (error=%ld)\n", ::GetLastError());
        return -1;
    }

    STARTUPINFO si = { sizeof(si) };
    wchar_t lpDesktop[16] = L"Winsta0\\default";
    si.lpDesktop = lpDesktop;

    PROCESS_INFORMATION pi;

    BOOL impersonated = ::SetThreadToken(nullptr, hDupToken);
    assert(impersonated);
    if (!impersonated) {
        printf("SetThreadToken failed: Access denied\n");
        return -1;
    }

    HANDLE hCurrentToken;
    DWORD session = 0, len = sizeof(session);
    ::OpenProcessToken(::GetCurrentProcess(), TOKEN_ALL_ACCESS, &hCurrentToken);
    ::GetTokenInformation(hCurrentToken, TokenSessionId, &session, len, &len);
    ::CloseHandle(hCurrentToken);

    if (!SetPrivilege(hDupToken, SE_ASSIGNPRIMARYTOKEN_NAME, TRUE) ||
        !SetPrivilege(hDupToken, SE_INCREASE_QUOTA_NAME, TRUE)) {
        printf("SetPrivilege failed: Insufficient privileges\n");
        return -1;
    }

    ::SetTokenInformation(hPrimary, TokenSessionId, &session, sizeof(session));


    std::string commandLine;
    for (int i = 0; i < argc; i++)
        (commandLine += argv[i]) += " ";
    commandLine += "--run_as_system";

    if (!CreateProcessAsUser(hPrimary, nullptr, const_cast<wchar_t*>(string_to_wstring(commandLine).c_str()),
        nullptr, nullptr, FALSE, 0, nullptr, nullptr, &si, &pi)) {
        printf("Failed to create process (error=%ld)\n", ::GetLastError());
        return -1;
    }
    else {
        printf("Process created: %ld\n", pi.dwProcessId);
    }

    return 0;
}

int main(int argc, const char* argv[]) {
    bool runningAsSystem = argc > 1 && !strcmp(argv[argc - 1], "--run_as_system");
    if (!runningAsSystem) {
        if (runAsSystem(argc, argv) < 0)
            printf("Cannot run as SYSTEM\n");
        else
            return 0;
    }

    if (SUCCEEDED(OleInitialize(NULL))) {
        IDataObject* pdto = new MyDataObject();
        if (pdto) {
            OleSetClipboard(pdto);
            pdto->Release();
        }

        MSG msg;
        while (GetMessage(&msg, NULL, 0, 0)) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
        OleUninitialize();
    }

    return 0;
}