#include "privilege.h"
#include "scope_guard.h"

using namespace zen;


Privileges& Privileges::getInstance()
{
    static Privileges instance;
    return instance;
}


bool Privileges::privilegeIsActive(LPCTSTR privilege) //throw FileError
{
    HANDLE hToken = NULL;
    if (!::OpenProcessToken(::GetCurrentProcess(), //__in   HANDLE ProcessHandle,
                            TOKEN_QUERY,           //__in   DWORD DesiredAccess,
                            &hToken))              //__out  PHANDLE TokenHandle
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());
    ZEN_ON_BLOCK_EXIT(::CloseHandle(hToken));

    LUID luid = {};
    if (!::LookupPrivilegeValue(
            NULL,      //__in_opt  LPCTSTR lpSystemName,
            privilege, //__in      LPCTSTR lpName,
            &luid ))   //__out     PLUID lpLuid
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());

    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(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());

    return alreadyGranted == TRUE;
}


void Privileges::setPrivilege(LPCTSTR privilege, bool enable) //throw FileError
{
    HANDLE hToken = NULL;
    if (!::OpenProcessToken(::GetCurrentProcess(),   //__in   HANDLE ProcessHandle,
                            TOKEN_ADJUST_PRIVILEGES, //__in   DWORD DesiredAccess,
                            &hToken))                //__out  PHANDLE TokenHandle
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());
    ZEN_ON_BLOCK_EXIT(::CloseHandle(hToken));

    LUID luid = {};
    if (!::LookupPrivilegeValue(
            NULL,      //__in_opt  LPCTSTR lpSystemName,
            privilege, //__in      LPCTSTR lpName,
            &luid ))   //__out     PLUID lpLuid
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());

    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,
            NULL,   //__out_opt  PTOKEN_PRIVILEGES PreviousState,
            NULL))  //__out_opt  PDWORD ReturnLength
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());

    if (::GetLastError() == ERROR_NOT_ALL_ASSIGNED) //check although previous function returned with success!
        throw FileError(_("Error setting privilege:") + " \"" + privilege +  "\"" + "\n\n" + getLastErrorFormatted());
}