1""" 2A module built to test if the current process has the privilege to 3create symlinks on Windows. 4""" 5# https://svn.python.org/projects/python/branches/pep-0384/Lib/test/symlink_support.py 6 7# allow script to run natively under python 2.6+ 8from __future__ import print_function 9 10import ctypes 11from ctypes import wintypes 12 13GetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess 14GetCurrentProcess.restype = wintypes.HANDLE 15OpenProcessToken = ctypes.windll.advapi32.OpenProcessToken 16OpenProcessToken.argtypes = (wintypes.HANDLE, wintypes.DWORD, 17 ctypes.POINTER(wintypes.HANDLE)) 18OpenProcessToken.restype = wintypes.BOOL 19 20class LUID(ctypes.Structure): 21 _fields_ = [ 22 ('low_part', wintypes.DWORD), 23 ('high_part', wintypes.LONG), 24 ] 25 26 def __eq__(self, other): 27 return ( 28 self.high_part == other.high_part and 29 self.low_part == other.low_part 30 ) 31 32 def __ne__(self, other): 33 return not (self==other) 34 35LookupPrivilegeValue = ctypes.windll.advapi32.LookupPrivilegeValueW 36LookupPrivilegeValue.argtypes = ( 37 wintypes.LPWSTR, # system name 38 wintypes.LPWSTR, # name 39 ctypes.POINTER(LUID), 40 ) 41LookupPrivilegeValue.restype = wintypes.BOOL 42 43class TOKEN_INFORMATION_CLASS: 44 TokenUser = 1 45 TokenGroups = 2 46 TokenPrivileges = 3 47 # ... see http://msdn.microsoft.com/en-us/library/aa379626%28VS.85%29.aspx 48 49SE_PRIVILEGE_ENABLED_BY_DEFAULT = (0x00000001) 50SE_PRIVILEGE_ENABLED = (0x00000002) 51SE_PRIVILEGE_REMOVED = (0x00000004) 52SE_PRIVILEGE_USED_FOR_ACCESS = (0x80000000) 53 54class LUID_AND_ATTRIBUTES(ctypes.Structure): 55 _fields_ = [ 56 ('LUID', LUID), 57 ('attributes', wintypes.DWORD), 58 ] 59 60 def is_enabled(self): 61 return bool(self.attributes & SE_PRIVILEGE_ENABLED) 62 63 def enable(self): 64 self.attributes |= SE_PRIVILEGE_ENABLED 65 66 def get_name(self): 67 size = wintypes.DWORD(10240) 68 buf = ctypes.create_unicode_buffer(size.value) 69 res = LookupPrivilegeName(None, self.LUID, buf, size) 70 if res == 0: 71 raise RuntimeError 72 return buf[:size.value] 73 74 def __str__(self): 75 name = self.name 76 fmt = ['{name}', '{name} (enabled)'][self.is_enabled()] 77 return fmt.format(**vars()) 78 79LookupPrivilegeName = ctypes.windll.advapi32.LookupPrivilegeNameW 80LookupPrivilegeName.argtypes = ( 81 wintypes.LPWSTR, # lpSystemName 82 ctypes.POINTER(LUID), # lpLuid 83 wintypes.LPWSTR, # lpName 84 ctypes.POINTER(wintypes.DWORD), #cchName 85 ) 86LookupPrivilegeName.restype = wintypes.BOOL 87 88class TOKEN_PRIVILEGES(ctypes.Structure): 89 _fields_ = [ 90 ('count', wintypes.DWORD), 91 ('privileges', LUID_AND_ATTRIBUTES*0), 92 ] 93 94 def get_array(self): 95 array_type = LUID_AND_ATTRIBUTES*self.count 96 privileges = ctypes.cast(self.privileges, 97 ctypes.POINTER(array_type)).contents 98 return privileges 99 100 def __iter__(self): 101 return iter(self.get_array()) 102 103PTOKEN_PRIVILEGES = ctypes.POINTER(TOKEN_PRIVILEGES) 104 105GetTokenInformation = ctypes.windll.advapi32.GetTokenInformation 106GetTokenInformation.argtypes = [ 107 wintypes.HANDLE, # TokenHandle 108 ctypes.c_uint, # TOKEN_INFORMATION_CLASS value 109 ctypes.c_void_p, # TokenInformation 110 wintypes.DWORD, # TokenInformationLength 111 ctypes.POINTER(wintypes.DWORD), # ReturnLength 112 ] 113GetTokenInformation.restype = wintypes.BOOL 114 115# http://msdn.microsoft.com/en-us/library/aa375202%28VS.85%29.aspx 116AdjustTokenPrivileges = ctypes.windll.advapi32.AdjustTokenPrivileges 117AdjustTokenPrivileges.restype = wintypes.BOOL 118AdjustTokenPrivileges.argtypes = [ 119 wintypes.HANDLE, # TokenHandle 120 wintypes.BOOL, # DisableAllPrivileges 121 PTOKEN_PRIVILEGES, # NewState (optional) 122 wintypes.DWORD, # BufferLength of PreviousState 123 PTOKEN_PRIVILEGES, # PreviousState (out, optional) 124 ctypes.POINTER(wintypes.DWORD), # ReturnLength 125 ] 126 127def get_process_token(): 128 "Get the current process token" 129 token = wintypes.HANDLE() 130 TOKEN_ALL_ACCESS = 0xf01ff 131 res = OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, token) 132 if not res > 0: 133 raise RuntimeError("Couldn't get process token") 134 return token 135 136def get_debug_luid(): 137 symlink_luid = LUID() 138 res = LookupPrivilegeValue(None, "SeDebugPrivilege", 139 symlink_luid) 140 if not res > 0: 141 raise RuntimeError("Couldn't lookup privilege value") 142 return symlink_luid 143 144def get_privilege_information(): 145 "Get all privileges associated with the current process." 146 # first call with zero length to determine what size buffer we need 147 148 return_length = wintypes.DWORD() 149 params = [ 150 get_process_token(), 151 TOKEN_INFORMATION_CLASS.TokenPrivileges, 152 None, 153 0, 154 return_length, 155 ] 156 157 res = GetTokenInformation(*params) 158 159 # assume we now have the necessary length in return_length 160 161 buffer = ctypes.create_string_buffer(return_length.value) 162 params[2] = buffer 163 params[3] = return_length.value 164 165 res = GetTokenInformation(*params) 166 assert res > 0, "Error in second GetTokenInformation (%d)" % res 167 168 privileges = ctypes.cast(buffer, ctypes.POINTER(TOKEN_PRIVILEGES)).contents 169 return privileges 170 171def report_privilege_information(): 172 "Report all privilege information assigned to the current process." 173 privileges = get_privilege_information() 174 print("found {0} privileges".format(privileges.count)) 175 tuple(map(print, privileges)) 176 177def enable_debug_privilege(): 178 """ 179 Try to assign the symlink privilege to the current process token. 180 Return True if the assignment is successful. 181 """ 182 # create a space in memory for a TOKEN_PRIVILEGES structure 183 # with one element 184 size = ctypes.sizeof(TOKEN_PRIVILEGES) 185 size += ctypes.sizeof(LUID_AND_ATTRIBUTES) 186 buffer = ctypes.create_string_buffer(size) 187 tp = ctypes.cast(buffer, ctypes.POINTER(TOKEN_PRIVILEGES)).contents 188 tp.count = 1 189 tp.get_array()[0].enable() 190 tp.get_array()[0].LUID = get_debug_luid() 191 token = get_process_token() 192 res = AdjustTokenPrivileges(token, False, tp, 0, None, None) 193 if res == 0: 194 raise RuntimeError("Error in AdjustTokenPrivileges") 195 196 ERROR_NOT_ALL_ASSIGNED = 1300 197 return ctypes.windll.kernel32.GetLastError() != ERROR_NOT_ALL_ASSIGNED 198