# -*- coding: utf-8 -*-

from ctypes import wintypes
import ctypes
import os

from win_structs import NTSTATUS, UNICODE_STRING

PVOID = ctypes.c_void_p
PULONG = ctypes.POINTER(wintypes.ULONG)
ULONG_PTR = wintypes.WPARAM
ACCESS_MASK = wintypes.DWORD
STATUS_INFO_LENGTH_MISMATCH = NTSTATUS(0xC0000004)


class SYSTEM_INFORMATION_CLASS(ctypes.c_ulong):

	def __repr__(self):
		return '%s(%s)' % (type(self).__name__, self.value)


SystemExtendedHandleInformation = SYSTEM_INFORMATION_CLASS(64)


class SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX(ctypes.Structure):

	_fields_ = [('Object', PVOID),
				('UniqueProcessId', wintypes.HANDLE),
				('HandleValue', wintypes.HANDLE),
				('GrantedAccess', ACCESS_MASK),
				('CreatorBackTraceIndex', wintypes.USHORT),
				('ObjectTypeIndex', wintypes.USHORT),
				('HandleAttributes', wintypes.ULONG),
				('Reserved', wintypes.ULONG)]


class SYSTEM_INFORMATION(ctypes.Structure):
	pass


PSYSTEM_INFORMATION = ctypes.POINTER(SYSTEM_INFORMATION)


class SYSTEM_HANDLE_INFORMATION_EX(SYSTEM_INFORMATION):

	_fields_ = [('NumberOfHandles', ULONG_PTR),
				('Reserved', ULONG_PTR),
				('_Handles', SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX * 1)]

	@property
	def Handles(self):
		arr_t = (SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX *
				 self.NumberOfHandles)
		return ctypes.POINTER(arr_t)(self._Handles)[0]


try:
	ntdll = ctypes.WinDLL('ntdll')
	ntdll.NtQuerySystemInformation.restype = NTSTATUS
	ntdll.NtQuerySystemInformation.argtypes = (SYSTEM_INFORMATION_CLASS,  # SystemInformationClass
											   PSYSTEM_INFORMATION,  # SystemInformation
											   wintypes.ULONG,  # SystemInformationLength
											   PULONG)  # ReturnLength
except WindowsError:
	# Just in case
	ntdll = None


ObjectBasicInformation = 0
ObjectNameInformation = 1
ObjectTypeInformation = 2


def _get_handle_info(handle, info_class):
	if hasattr(handle, "HandleValue"):
		handle = handle.HandleValue
	size_needed = wintypes.DWORD()
	buf = ctypes.c_buffer(0x1000)
	ntdll.NtQueryObject(handle, info_class, ctypes.byref(buf),
						ctypes.sizeof(buf), ctypes.byref(size_needed))
	return UNICODE_STRING.from_buffer_copy(buf[:size_needed.value]).Buffer


def get_handle_name(handle):
	return _get_handle_info(handle, ObjectNameInformation)


def get_handle_type(handle):
	return _get_handle_info(handle, ObjectTypeInformation)


def get_handles():
	info = SYSTEM_HANDLE_INFORMATION_EX()
	length = wintypes.ULONG()
	while True:
		status = ntdll.NtQuerySystemInformation(SystemExtendedHandleInformation,
												ctypes.byref(info),
												ctypes.sizeof(info),
												ctypes.byref(length))
		if status != STATUS_INFO_LENGTH_MISMATCH:
			break
		ctypes.resize(info, length.value)
	if status < 0:
		raise ctypes.WinError(ntdll.RtlNtStatusToDosError(status))
	return info.Handles


def get_process_handles(pid=None):
	"""
	Get handles of process <pid> (current process if not specified)
	
	"""
	if not pid:
		pid = os.getpid()
	handles = []
	for handle in get_handles():
		if handle.UniqueProcessId != pid:
			continue
		handles.append(handle)
	return handles


if __name__ == "__main__":
	import sys
	if len(sys.argv) > 1:
		pid = int(sys.argv[1])
	else:
		pid = None
	for handle in get_process_handles(pid):
		print("Handle = 0x%04x, Type = 0x%02x %r, Access = 0x%06x, Name = %r" %
			  (handle.HandleValue, handle.ObjectTypeIndex,
			   get_handle_type(handle), handle.GrantedAccess,
			   get_handle_name(handle)))
