Esempio n. 1
0
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]
Esempio n. 2
0
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])
    [2.0 None->None] '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])
    [1745645548.0 None->None] '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__()
Esempio n. 3
0
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])
    [2.0 None->None] '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])
    [1745645548.0 None->None] '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__()