class QLinearViewer(QWidget): def __init__(self, workspace, disasm_view, parent=None): super(QLinearViewer, self).__init__(parent) self.workspace = workspace self.disasm_view = disasm_view self.objects = [] # Objects that will be painted self.cfg = None self.cfb = None self._offset_to_addr = AVLTree() self._addr_to_offset = AVLTree() self._offset_to_object = AVLTree() self._offset = 0 self._paint_start_offset = 0 self._linear_view = None # type: QLinearGraphicsView self._disasms = {} self._init_widgets() # # Properties # @property def offset(self): return self._offset @offset.setter def offset(self, v): self._offset = v @property def paint_start_offset(self): return self._paint_start_offset @property def max_offset(self): # TODO: Cache it try: max_off, obj = self._offset_to_object.max_item() except (KeyError, ValueError): return 0 return max_off + obj.height # # Proxy properties # @property def selected_operands(self): return self._linear_view.selected_operands @property def selected_insns(self): return self._linear_view.selected_insns # # Public methods # def initialize(self): self._make_objects() def navigate_to_addr(self, addr): if not self._addr_to_offset: return try: _, floor_offset = self._addr_to_offset.floor_item(addr) except KeyError: _, floor_offset = floor_offset = self._addr_to_offset.min_item() self.navigate_to(floor_offset) def refresh(self): self._linear_view.refresh() def navigate_to(self, offset): self._linear_view.verticalScrollBar().setValue(int(offset)) self.prepare_objects(offset) self._linear_view.refresh() def prepare_objects(self, offset): if offset == self._offset: return try: start_offset = self._offset_to_object.floor_key(offset) except (KeyError, ValueError): try: start_offset = self._offset_to_object.min_key() except ValueError: # Tree is empty return # Update offset self._offset = offset self._paint_start_offset = start_offset self.objects = [] max_height = self.height() for off, obj in self._offset_to_object.iter_items( start_key=start_offset): self.objects.append(obj) if off - offset > max_height: break # # Private methods # def _init_widgets(self): self._linear_view = QLinearGraphicsView(self, self.disasm_view) layout = QHBoxLayout() layout.addWidget(self._linear_view) layout.setContentsMargins(0, 0, 0, 0) self.setLayout(layout) # Setup proxy methods self.update_label = self._linear_view.update_label self.select_instruction = self._linear_view.select_instruction self.unselect_instruction = self._linear_view.unselect_instruction self.unselect_all_instructions = self._linear_view.unselect_all_instructions self.select_operand = self._linear_view.select_operand self.unselect_operand = self._linear_view.unselect_operand self.unselect_all_operands = self._linear_view.unselect_all_operands self.show_selected = self._linear_view.show_selected self.show_instruction = self._linear_view.show_instruction def _make_objects(self): self._addr_to_offset.clear() self._offset_to_addr.clear() self._offset_to_object.clear() y = 0 self._linear_view._clear_insn_addr_block_mapping() for obj_addr, obj in self.cfb.floor_items(): if isinstance(obj, Block): func_addr = self.cfg.get_any_node( obj.addr).function_address # FIXME: Resiliency func = self.cfg.kb.functions[func_addr] # FIXME: Resiliency disasm = self._get_disasm(func) qobject = QBlock( self.workspace, func_addr, self.disasm_view, disasm, self.disasm_view.infodock, obj.addr, [obj], {}, mode='linear', ) for insn_addr in qobject.addr_to_insns.keys(): self._linear_view._add_insn_addr_block_mapping( insn_addr, qobject) elif isinstance(obj, Unknown): qobject = QUnknownBlock(self.workspace, obj_addr, obj.bytes) else: continue self._offset_to_object[y] = qobject if obj_addr not in self._addr_to_offset: self._addr_to_offset[obj_addr] = y y += qobject.height def _get_disasm(self, func): """ :param func: :return: """ if func.addr not in self._disasms: self._disasms[ func. addr] = self.workspace.instance.project.analyses.Disassembly( function=func) return self._disasms[func.addr]
class SortedTypedList(object): """This data structure allows to sort and maintain sorted a list of objects inheriting from :class:`netzob.Common.Utils.SortableObject.SortableObject`. It uses an AVLTree proposed by :mod:`bintrees` to represent elements in the list. :TODO: The inner object __tree stores the given element using an AVLTree. Thus exposing new methods only requires the creation of a wrapper method in this class. >>> from netzob.all import * >>> from netzob.Common.Utils.SortedTypedList import SortedTypedList >>> msg1 = RawMessage(b"msg1", date=25.0) >>> msg2 = RawMessage(b"msg2", date=2.0) >>> msg3 = RawMessage(b"msg3", date=1456487548.0) >>> l = SortedTypedList(RawMessage, [msg2, msg3, msg1]) >>> print(list(l.values())[0]) [0;32m[2.0 [0;m[1;32mNone[1;m[0;32m->[0;m[1;32mNone[1;m[0;32m][0;m 'msg2' >>> msg4 = RawMessage(b"msg4", date=145548.0) >>> l.add(msg4) >>> msg5 = RawMessage(b"msg5", date=14.0) >>> msg6 = RawMessage(b"msg6", date=1745645548.0) >>> l.addAll([msg5, msg6]) >>> print(list(l.values())[5]) [0;32m[1745645548.0 [0;m[1;32mNone[1;m[0;32m->[0;m[1;32mNone[1;m[0;32m][0;m 'msg6' >>> print(len(l)) 6 """ def __init__(self, membersTypes, elements=None): self.membersTypes = membersTypes self.__treePriorities = AVLTree() self.__mapMessages = dict() if elements is not None and len(elements) > 0: self._extend(elements) def add(self, element): """Insert in the proper place the specified element. :type: any object that comply with the typed of the current list and inherits from :class:`netzob.Common.Utils.SortableObject.SortableObject`. :raises: a TypeError if element is None or if its type doesn't comply with the definition of the list. """ if element is None: raise TypeError("Element cannot be None") self._extend([element]) def addAll(self, elements): """Insert in their proper place all the specified element. :type: a list of any object that comply with the typed of the current list and inherits from :class:`netzob.Common.Utils.SortableObject.SortableObject`. :raises: a TypeError if element is None or if its type doesn't comply with the definition of the list. """ if elements is None: raise TypeError("Element cannot be None") self._extend(elements) def values(self): """Return a list sorted with the values of the current SortedTypedList. :warning: modifying this list has no effect on the SortedTypedList. :rtype: :mod:list """ l = [] for x in list(self.__treePriorities.keys()): l.extend(self.__mapMessages[x]) return l def clear(self): """remove all items from the list. It's a O(n) operation""" self.__treePriorities.clear() def _extend(self, elements): """Add all the elements in the current list. :parameter elements: a list of :class:`netzob.Common.Utils.SortableObject.SortableObject` to insert. :raises: TypeError if something is wrong with the given elements """ for e in elements: self._check(e) d = dict() for e in elements: d[e.priority()] = None if e.priority() in self.__mapMessages: self.__mapMessages[e.priority()].append(e) else: self.__mapMessages[e.priority()] = [e] self.__treePriorities.update(d) def _check(self, v): if not isinstance(v, self.membersTypes): raise TypeError( "Invalid type for argument, expecting: {0}, received : {1}". format(self.membersTypes, v.__class__.__name__)) if not isinstance(v, SortableObject): raise TypeError( "Objects inserted in a SortedTypedList must inherits from SortableObject class" ) def __len__(self): """Returns the number of elements in the sorted list which takes O(1) operation :)""" return len(self.__treePriorities) def __str__(self): return ', \n'.join([str(v) for v in list(self.values())]) def __repr__(self): return repr(str(self)) def __iter__(self): """SortedTypedList is an iterable over its values (and not its keys).""" return list(self.__treePriorities.values()).__iter__()