def driver_type():

    implist = idaapi.get_import_module_qty()

    for i in range(0, implist):
        name = idaapi.get_import_module_name(i)
        idaapi.enum_import_names(i, cb)
    for name in names:
        if name == "FltRegisterFilter":
            return "Mini-Filter"
        elif name == "WdfVersionBind":
            return "WDF"
        elif name == "StreamClassRegisterMinidriver":
            return "Stream Minidriver"
        elif name == "KsCreateFilterFactory":
            return "AVStream"
        elif name == "PcRegisterSubdevice":
            return "PortCls"
    return "WDM" 

def get_iat_data(self):
        """
        Retrive data from IAT
        """
        imp_num = idaapi.get_import_module_qty()  # Number of imported modules

        for i in xrange(0,imp_num):
            name = idaapi.get_import_module_name(i).lower()
            if not name:
                #self.logger.error("Failed to get import module name for #%d", i)
                continue

            if not name in self.iat:
                self.iat[name]= []

            self.current_module = self.iat[name]
            idaapi.enum_import_names(i, self.imp_cb) 

def _build_imports(self):
        '''Build imports table. (Was taken from examples.)'''

        tree = {}
        nimps = idaapi.get_import_module_qty()

        for i in xrange(0, nimps):
            name = idaapi.get_import_module_name(i)
            if not name:
                continue
            # Create a list for imported names
            self.tmp_items = []

            # Enum imported entries in this module
            idaapi.enum_import_names(i, self._imports_names_cb)

            if name not in tree:
                tree[name] = []
            tree[name].extend(self.tmp_items)

        return tree 

def compute_imports():
        imports = {}
        current = ""

        def callback(ea, name, ordinal):
            imports[current].append((ea, name, ordinal))
            return True

        nimps = idaapi.get_import_module_qty()
        for i in xrange(0, nimps):
            current = idaapi.get_import_module_name(i)
            imports[current] = []
            idaapi.enum_import_names(i, callback)
        return imports 

def get_idata(self): #return tuple(start_addr, end_addr)
        ida_seg = idaapi.get_segm_by_name(".idata")
        if ida_seg is None:
            addr = None
            def cb(ea, name, i):
                addr = ea
            idaapi.enum_import_names(0, cb)
            ida_seg = idaapi.seg_by_addr(addr)
        return (ida_seg.start_ea, ida_seg.end_ea)
    
    #------------------------------------- 

def getImportTableData(self):
        """
        Update rt_import_table with current import table data.
        """

        def imp_cb(ea, name, ord):
            """
            Import enumeration callback function. used by idaapi.enum_import_names .
            """
            tmpImports.append([self.current_module_name, ea, name, ord])
            return True

        tmpImports = []  # Contains static import table data (w\o real function addresses)
        imp_num = idaapi.get_import_module_qty()  # Number of imported modules

        for i in xrange(0, imp_num):
            self.current_module_name = idaapi.get_import_module_name(i).lower()
            idaapi.enum_import_names(i, imp_cb)

        #  Get runtime function addresses and store in self.rt_import_table
        if not idaapi.is_debugger_on():
            raise RuntimeError("Debugger is not currently active.")

        for module_name, ea, name, ord in tmpImports:
            func_real_adrs = get_adrs_mem(ea)
            self.rt_import_table[func_real_adrs] = (module_name, ea, name, ord) 

def getApiMap(self):
        self._api_map = {}
        num_imports = ida_nalt.get_import_module_qty()
        for i in range(0, num_imports):
            self._import_module_name = ida_nalt.get_import_module_name(i)
            ida_nalt.enum_import_names(i, self._cbEnumImports)
        return self._api_map 

def getApiMap(self):
        self._api_map = {}
        num_imports = idaapi.get_import_module_qty()
        for i in range(0, num_imports):
            self._import_module_name = idaapi.get_import_module_name(i)
            idaapi.enum_import_names(i, self._cbEnumImports)
        return self._api_map 

def make_import_names_callback(library_calls, library_addr):
    """ Return a callback function used by idaapi.enum_import_names(). """
    def callback(ea, name, ordinal):
        """ Callback function to retrieve code references to library calls. """
        library_calls[name] = []
        library_addr[name] = ea
        for ref in idautils.CodeRefsTo(ea, 0):
            library_calls[name].append(ref)
        return True  # True -> Continue enumeration
    return callback 

def get_imports(library_calls, library_addr):
    """ Populate dictionaries with import information. """
    import_names_callback = make_import_names_callback(library_calls,
                                                       library_addr)
    for i in xrange(0, idaapi.get_import_module_qty()):
        idaapi.enum_import_names(i, import_names_callback) 

def find_pool_tags():
	""" Dirty hack around IDA's type information, find references to tag using functions then the comment marking the tag 
	then add the function caller/tag to output dictionary.
	"""
	
	funcs = [
		'ExAllocatePoolWithTag',
		'ExFreePoolWithTag',
		'ExAllocatePoolWithTagPriority'
	]

	tags = {}

	def imp_cb(ea, name, ord):
		if name in funcs:
			for xref in idautils.XrefsTo(ea):
				call_addr = xref.frm
				caller_name = idc.GetFunctionName(call_addr)
				prev = idc.PrevHead(call_addr)
				for _ in range(10):
					if idc.Comment(prev) == 'Tag' and idc.GetOpType(prev, 1) == 5:
						tag_raw = idc.GetOperandValue(prev, 1)
						tag = ''
						for i in range(3, -1, -1):
							tag += chr((tag_raw >> 8 * i) & 0xFF)
						if tag in tags.keys():
							tags[tag].add(caller_name)
						else:
							tags[tag] = set([caller_name])
						break
					prev = idc.PrevHead(prev)
		return True
	
	nimps = idaapi.get_import_module_qty()

	for i in xrange(0, nimps):
		name = idaapi.get_import_module_name(i)
		if not name:
			continue

		idaapi.enum_import_names(i, imp_cb)
	return tags