import ctypes.wintypes


class LUID(ctypes.Structure):
	_fields_ = [
		('low_part', ctypes.wintypes.DWORD),
		('high_part', ctypes.wintypes.LONG),
	]

	def __eq__(self, other):
		return (
			self.high_part == other.high_part and
			self.low_part == other.low_part
		)

	def __ne__(self, other):
		return not (self == other)


LookupPrivilegeValue = ctypes.windll.advapi32.LookupPrivilegeValueW
LookupPrivilegeValue.argtypes = (
	ctypes.wintypes.LPWSTR,  # system name
	ctypes.wintypes.LPWSTR,  # name
	ctypes.POINTER(LUID),
)
LookupPrivilegeValue.restype = ctypes.wintypes.BOOL


class TOKEN_INFORMATION_CLASS:
	TokenUser = 1
	TokenGroups = 2
	TokenPrivileges = 3
	# ... see http://msdn.microsoft.com/en-us/library/aa379626%28VS.85%29.aspx


SE_PRIVILEGE_ENABLED_BY_DEFAULT = 0x00000001
SE_PRIVILEGE_ENABLED = 0x00000002
SE_PRIVILEGE_REMOVED = 0x00000004
SE_PRIVILEGE_USED_FOR_ACCESS = 0x80000000


class LUID_AND_ATTRIBUTES(ctypes.Structure):
	_fields_ = [
		('LUID', LUID),
		('attributes', ctypes.wintypes.DWORD),
	]

	def is_enabled(self):
		return bool(self.attributes & SE_PRIVILEGE_ENABLED)

	def enable(self):
		self.attributes |= SE_PRIVILEGE_ENABLED

	def get_name(self):
		size = ctypes.wintypes.DWORD(10240)
		buf = ctypes.create_unicode_buffer(size.value)
		res = LookupPrivilegeName(None, self.LUID, buf, size)
		if res == 0:
			raise RuntimeError
		return buf[:size.value]

	def __str__(self):
		res = self.get_name()
		if self.is_enabled():
			res += ' (enabled)'
		return res


LookupPrivilegeName = ctypes.windll.advapi32.LookupPrivilegeNameW
LookupPrivilegeName.argtypes = (
	ctypes.wintypes.LPWSTR,  # lpSystemName
	ctypes.POINTER(LUID),  # lpLuid
	ctypes.wintypes.LPWSTR,  # lpName
	ctypes.POINTER(ctypes.wintypes.DWORD),  # cchName
)
LookupPrivilegeName.restype = ctypes.wintypes.BOOL


class TOKEN_PRIVILEGES(ctypes.Structure):
	_fields_ = [
		('count', ctypes.wintypes.DWORD),
		('privileges', LUID_AND_ATTRIBUTES * 0),
	]

	def get_array(self):
		array_type = LUID_AND_ATTRIBUTES * self.count
		privileges = ctypes.cast(
			self.privileges, ctypes.POINTER(array_type)).contents
		return privileges

	def __iter__(self):
		return iter(self.get_array())


PTOKEN_PRIVILEGES = ctypes.POINTER(TOKEN_PRIVILEGES)

GetTokenInformation = ctypes.windll.advapi32.GetTokenInformation
GetTokenInformation.argtypes = [
	ctypes.wintypes.HANDLE,  # TokenHandle
	ctypes.c_uint,  # TOKEN_INFORMATION_CLASS value
	ctypes.c_void_p,  # TokenInformation
	ctypes.wintypes.DWORD,  # TokenInformationLength
	ctypes.POINTER(ctypes.wintypes.DWORD),  # ReturnLength
]
GetTokenInformation.restype = ctypes.wintypes.BOOL

# http://msdn.microsoft.com/en-us/library/aa375202%28VS.85%29.aspx
AdjustTokenPrivileges = ctypes.windll.advapi32.AdjustTokenPrivileges
AdjustTokenPrivileges.restype = ctypes.wintypes.BOOL
AdjustTokenPrivileges.argtypes = [
	ctypes.wintypes.HANDLE,                # TokenHandle
	ctypes.wintypes.BOOL,                  # DisableAllPrivileges
	PTOKEN_PRIVILEGES,              # NewState (optional)
	ctypes.wintypes.DWORD,                 # BufferLength of PreviousState
	PTOKEN_PRIVILEGES,              # PreviousState (out, optional)
	ctypes.POINTER(ctypes.wintypes.DWORD),  # ReturnLength
]