#include "privilege.h"
#include "system_func.h"
#include "i18n.h"
#include "loki/ScopeGuard.h"

using namespace ffs3;


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
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::getLastErrorFormatted());
    }
    Loki::ScopeGuard dummy = Loki::MakeGuard(::CloseHandle, hToken);
    (void)dummy; //silence warning "unused variable"


    LUID luid = {};
    if (!::LookupPrivilegeValue(
            NULL,      //__in_opt  LPCTSTR lpSystemName,
            privilege, //__in      LPCTSTR lpName,
            &luid ))   //__out     PLUID lpLuid
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::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
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::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
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::getLastErrorFormatted());
    }
    Loki::ScopeGuard dummy = Loki::MakeGuard(::CloseHandle, hToken);
    (void)dummy; //silence warning "unused variable"

    LUID luid = {};
    if (!::LookupPrivilegeValue(
            NULL,      //__in_opt  LPCTSTR lpSystemName,
            privilege, //__in      LPCTSTR lpName,
            &luid ))   //__out     PLUID lpLuid
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::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
    {
        const wxString errorMessage = wxString(_("Error setting privilege:")) + wxT(" \"") + privilege +  wxT("\"") + wxT("\n\n");
        throw FileError(errorMessage + ffs3::getLastErrorFormatted());
    }

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