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
130
131
132
133
134
135
136
137
138
|
#include "privilege.h"
#include <map>
#include "thread.h" //includes <boost/thread.hpp>
#include "zstring.h"
#include "scope_guard.h"
#include "win_ver.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()));
ErrorCode lastError = getLastError();
if (lastError == ERROR_NOT_ALL_ASSIGNED) //check although previous function returned with success!
{
#ifdef __MINGW32__ //Shobjidl.h
#define ERROR_ELEVATION_REQUIRED 740L
#endif
if (vistaOrLater()) //replace this useless error code with what it *really* means!
lastError = ERROR_ELEVATION_REQUIRED;
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);
}
|
