#include "privilege.h"
#include <map>
#include "thread.h" //includes <boost/thread.hpp>
#include "zstring.h"
#include "scope_guard.h"

using namespace zen;


namespace
{
bool privilegeIsActive(LPCTSTR privilege) //throw FileError
{
    HANDLE hToken = nullptr;
    if (!::OpenProcessToken(::GetCurrentProcess(), //__in   HANDLE ProcessHandle,
                            TOKEN_QUERY,           //__in   DWORD DesiredAccess,
                            &hToken))              //__out  PHANDLE TokenHandle
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"OpenProcessToken", getLastError()));
    ZEN_ON_SCOPE_EXIT(::CloseHandle(hToken));

    LUID luid = {};
    if (!::LookupPrivilegeValue(nullptr,   //__in_opt  LPCTSTR lpSystemName,
                                privilege, //__in      LPCTSTR lpName,
                                &luid ))   //__out     PLUID lpLuid
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"LookupPrivilegeValue", getLastError()));

    PRIVILEGE_SET priv  = {};
    priv.PrivilegeCount = 1;
    priv.Control        = PRIVILEGE_SET_ALL_NECESSARY;
    priv.Privilege[0].Luid = luid;
    priv.Privilege[0].Attributes = SE_PRIVILEGE_ENABLED;

    BOOL alreadyGranted = FALSE;
    if (!::PrivilegeCheck(hToken,           //__in     HANDLE ClientToken,
                          &priv,            //__inout  PPRIVILEGE_SET RequiredPrivileges,
                          &alreadyGranted)) //__out    LPBOOL pfResult
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"PrivilegeCheck", getLastError()));

    return alreadyGranted != FALSE;
}


void setPrivilege(LPCTSTR privilege, bool enable) //throw FileError
{
    HANDLE hToken = nullptr;
    if (!::OpenProcessToken(::GetCurrentProcess(),   //__in   HANDLE ProcessHandle,
                            TOKEN_ADJUST_PRIVILEGES, //__in   DWORD DesiredAccess,
                            &hToken))                //__out  PHANDLE TokenHandle
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"OpenProcessToken", getLastError()));
    ZEN_ON_SCOPE_EXIT(::CloseHandle(hToken));

    LUID luid = {};
    if (!::LookupPrivilegeValue(nullptr,   //__in_opt  LPCTSTR lpSystemName,
                                privilege, //__in      LPCTSTR lpName,
                                &luid ))   //__out     PLUID lpLuid
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"LookupPrivilegeValue", getLastError()));

    TOKEN_PRIVILEGES tp = {};
    tp.PrivilegeCount   = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = enable ? SE_PRIVILEGE_ENABLED : 0;

    if (!::AdjustTokenPrivileges(hToken,   //__in       HANDLE TokenHandle,
                                 false,    //__in       BOOL DisableAllPrivileges,
                                 &tp,      //__in_opt   PTOKEN_PRIVILEGES NewState,
                                 0,        //__in       DWORD BufferLength,
                                 nullptr,  //__out_opt  PTOKEN_PRIVILEGES PreviousState,
                                 nullptr)) //__out_opt  PDWORD ReturnLength
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"AdjustTokenPrivileges", getLastError()));

    const ErrorCode lastError = getLastError();
    if (lastError == ERROR_NOT_ALL_ASSIGNED) //check although previous function returned with success!
        throw FileError(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\""), formatSystemError(L"AdjustTokenPrivileges", lastError));
}


class Privileges
{
public:
    static Privileges& getInstance()
    {
        static Privileges inst;
        return inst;
    }

    void ensureActive(LPCTSTR privilege) //throw FileError
    {
        if (activePrivileges.find(privilege) != activePrivileges.end())
            return; //privilege already active

        if (privilegeIsActive(privilege)) //privilege was already active before starting this tool
            activePrivileges.insert(std::make_pair(privilege, false));
        else
        {
            setPrivilege(privilege, true);
            activePrivileges.insert(std::make_pair(privilege, true));
        }
    }

private:
    Privileges() {}
    Privileges(Privileges&);
    Privileges& operator=(const Privileges&);

    ~Privileges() //clean up: deactivate all privileges that have been activated by this application
    {
        for (auto it = activePrivileges.begin(); it != activePrivileges.end(); ++it)
            if (it->second)
                try
                {
                    setPrivilege(it->first.c_str(), false); //throw FileError
                }
                catch (...) {}
    }

    std::map<Zstring, bool> activePrivileges; //bool: enabled by this application
};

boost::mutex lockPrivileges;
}


void zen::activatePrivilege(LPCTSTR privilege) //throw FileError
{
    boost::lock_guard<boost::mutex> dummy(lockPrivileges);
    Privileges::getInstance().ensureActive(privilege);
}