Python code coverage for Lib/test/symlink_support.py

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