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