summaryrefslogtreecommitdiff
path: root/zen/privilege.cpp
blob: 69c820f3c6ad6ad7a759bf9d3efc5ce45ca53508 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#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"\"") + L"\n\n" + getLastErrorFormatted());
    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"\"") + L"\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(replaceCpy(_("Cannot set privilege %x."), L"%x", std::wstring(L"\"") + privilege +  L"\"") + L"\n\n" + getLastErrorFormatted());

    return alreadyGranted == TRUE;
}


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"\"") + L"\n\n" + getLastErrorFormatted());
    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"\"") + L"\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,
                                 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"\"") + L"\n\n" + getLastErrorFormatted());

    if (::GetLastError() == 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"\"") + L"\n\n" + getLastErrorFormatted());
}


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&);
    void operator=(Privileges&);

    ~Privileges() //clean up: deactivate all privileges that have been activated by this application
    {
        for (auto iter = activePrivileges.begin(); iter != activePrivileges.end(); ++iter)
            if (iter->second)
                try
                {
                    setPrivilege(iter->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);
}
bgstack15