1 files changed, 79 insertions, 0 deletions
diff --git a/zen/privilege.cpp b/zen/privilege.cpp
new file mode 100644
index 00000000..495b1254
--- /dev/null
+++ b/zen/privilege.cpp
@@ -0,0 +1,79 @@
+#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());
+}