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
127
128
129
|
#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 (const auto& priv : activePrivileges)
if (priv.second)
{
try
{
setPrivilege(priv.first.c_str(), false); //throw FileError
}
catch (FileError&) {}
}
}
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);
}
|
