From 0583d70c059be5c44deaffa77ab42062f732b99a Mon Sep 17 00:00:00 2001 From: Audrey Dutcher Date: Wed, 2 Nov 2022 20:07:43 -0700 Subject: [PATCH] some algorithm fixes. the girls are interactinggggg --- typetapper/analysis.py | 32 ++++--- typetapper/data.py | 14 ++- typetapper/engine.py | 4 +- typetapper/hierarchy_graph.py | 1 + typetapper/hierarchy_graph_view.py | 147 ++++++++++++++++++++++++++--- typetapper/relative_graph.py | 19 ++-- 6 files changed, 183 insertions(+), 34 deletions(-) diff --git a/typetapper/analysis.py b/typetapper/analysis.py index 628aee7..83b7cf6 100644 --- a/typetapper/analysis.py +++ b/typetapper/analysis.py @@ -8,6 +8,7 @@ from angr.knowledge_plugins.cfg import CFGNode, CFGModel from .engine import TypeTapperEngine from .knowledge import TypeTapperManager +from .data import Atom, RegisterAtom, CodeLoc, OpSequence, ControlFlowActionPop, ControlFlowActionPush l = logging.getLogger(__name__) @@ -51,24 +52,33 @@ class TypeTapperAnalysis(angr.Analysis): continue pred_addr = pred.addr pred_blockinfo = self.manager.block_info[pred_addr] + callsite_addr = fakeret_addr if attrs['jumpkind'] == 'Ijk_Ret' else pred_addr if attrs['jumpkind'] == 'Ijk_Call' else None # TAKE IT BACK NOW Y'ALL for name in node_blockinfo.ready_inputs: - input_info = node_blockinfo.inputs[name] - input_info_new = input_info.step(pred_addr, block_addr, attrs['jumpkind'], fakeret_addr) - if input_info_new is None: - continue - output_atom = pred_blockinfo.outputs.get(input_info.atom.slot_name, None) + input_atom = node_blockinfo.inputs[name] + output_atom = pred_blockinfo.outputs.get(input_atom.slot_name, None) if output_atom is not None: if output_atom.name == name: - input_info_new.commit(self.manager.graph, output_atom) + self._passive_link(input_atom, output_atom, attrs['jumpkind'], callsite_addr) else: pass # alias mismatch - elif name not in pred_blockinfo.inputs: # sketchy... this means that we can't account for multiple paths to the same atom - pred_blockinfo.inputs[name] = input_info_new - pred_blockinfo.ready_inputs.add(name) - queue[pred_addr] = None - queue.move_to_end(pred_addr, last=False) + else: + if name not in pred_blockinfo.inputs: + pred_blockinfo.inputs[name] = RegisterAtom(CodeLoc(pred_addr, 0, pred_addr), input_atom.size, input_atom.name, input_atom.slot_name) + pred_blockinfo.ready_inputs.add(name) + queue[pred_addr] = None + queue.move_to_end(pred_addr, last=False) + self._passive_link(input_atom, pred_blockinfo.inputs[name], attrs['jumpkind'], callsite_addr) # this is safe because if there is a loop edge, the atom will already be in the inputs so it won't be re-upped node_blockinfo.ready_inputs.clear() + + def _passive_link(self, input_atom: Atom, output_atom: Atom, jumpkind: str, callsite: Optional[int]): + if jumpkind == 'Ijk_Ret': + cf = [ControlFlowActionPop(callsite)] + elif jumpkind == 'Ijk_Call': + cf = [ControlFlowActionPush(callsite)] + else: + cf = [] + self.manager.graph.add_edge(output_atom, input_atom, ops=OpSequence(), cf=cf) diff --git a/typetapper/data.py b/typetapper/data.py index cb909dc..84d4291 100644 --- a/typetapper/data.py +++ b/typetapper/data.py @@ -23,18 +23,30 @@ class RegisterAtom(Atom): name: str slot_name: str + def __repr__(self): + return f'{self.name} @ {self.loc.ins_addr:#x}' + @dataclass(frozen=True) class MemoryAtom(Atom): endness: str + def __repr__(self): + return f'MEM @ {self.loc.ins_addr:#x}' + @dataclass(frozen=True) class TmpAtom(Atom): tmp: int + def __repr__(self): + return f'TMP#{self.tmp} @ {self.loc.ins_addr:#x}' + @dataclass(frozen=True) class ConstAtom(Atom): value: int + def __repr__(self): + return f'CONST#{self.value:#x} @ {self.loc.ins_addr:#x}' + @dataclass(frozen=True) class Op: @@ -277,6 +289,6 @@ class ControlFlowActionPop(ControlFlowAction): @dataclass class BlockInfo: outputs: Dict[str, RegisterAtom] = field(default_factory=dict) # slot names - inputs: Dict[str, RegisterInputInfo] = field(default_factory=dict) # alias names + inputs: Dict[str, RegisterAtom] = field(default_factory=dict) # alias names atoms: List[Atom] = field(default_factory=list) ready_inputs: Set[str] = field(default_factory=set) # alias names diff --git a/typetapper/engine.py b/typetapper/engine.py index c6ababa..c9081b2 100644 --- a/typetapper/engine.py +++ b/typetapper/engine.py @@ -64,9 +64,9 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): else: pass # alias mismatch elif name in self.blockinfo.inputs: - return LiveData.new_atom(self.blockinfo.inputs[name].atom) + return LiveData.new_atom(self.blockinfo.inputs[name]) else: - self.blockinfo.inputs[name] = RegisterInputInfo(atom=reg_atom, callsites=(), reverse_callsites=()) + self.blockinfo.inputs[name] = reg_atom self.blockinfo.ready_inputs.add(name) return LiveData.new_atom(reg_atom) diff --git a/typetapper/hierarchy_graph.py b/typetapper/hierarchy_graph.py index 13efa13..4cf508d 100644 --- a/typetapper/hierarchy_graph.py +++ b/typetapper/hierarchy_graph.py @@ -57,6 +57,7 @@ class HierarchicalGraph(RelativeAtomGraph): queue = set(seeds) while queue: seed = queue.pop() + self.move_node(seed, group) self.expand(seed, group=group) # no-op if not in frontier for adj in chain(self.predecessors(seed), self.successors(seed)): if adj in self.frontier and adj not in queue and predicate(adj): diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index 62e07bd..98a8b12 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -1,25 +1,31 @@ -from typing import Tuple, Optional +from typing import Tuple, Optional, Union, Dict from collections import defaultdict from math import sqrt, sin, cos, pi from PySide6.QtCore import QRectF, QPointF -from PySide6.QtGui import QColor, QFont, QPen, QBrush, QPainterPath +from PySide6.QtGui import QColor, QFont, QBrush, QPainterPath, QPen from PySide6.QtWidgets import QGraphicsItem, QGraphicsRectItem, QGraphicsSimpleTextItem, QHBoxLayout, QGraphicsScene import networkx +from angrmanagement.config import Conf +from angrmanagement.data.instance import Instance +from angrmanagement.data.object_container import ObjectContainer from angrmanagement.ui.views import BaseView from angrmanagement.ui.widgets.qgraph import QZoomableDraggableGraphicsView -from .hierarchy_graph import HierarchicalGraph, RelativeAtomGroup +from .hierarchy_graph import HierarchicalGraph, RelativeAtomGroup, RelativeAtomOrGroup from .data import Prop, DataKind +from .relative_graph import RelativeAtom TEXT_COLOR = QColor(0, 0, 0) TEXT_FONT = QFont("default", 10) BACKGROUND_COLOR = QColor(255, 255, 255) BOX_COLOR = QColor(0xff, 0xc0, 0x80) -BOX_BORDER_COLOR = QColor(0x10, 0x10, 0x10) -ARROW_COLOR = QColor(0, 0, 0) +BOX_BORDER_COLOR = QColor(0x60, 0x60, 0x60) +BOX_BORDER_SELECTED_COLOR = QColor(0x60, 0x60, 0xff) +BOX_BORDER_WIDTH = 1 +BOX_BORDER_SELECTED_WIDTH = 2 def kind_to_color(kind: DataKind) -> QColor: if kind == DataKind.Pointer: @@ -37,30 +43,67 @@ class HierarchicalGraphView(BaseView): self.base_caption = "TypeTapper" layout = QHBoxLayout(self) - self.graph = HierarchicalGraphWidget(self, hg, current_group) + self.graph = HierarchicalGraphWidget(self, instance, hg, current_group) layout.addWidget(self.graph) self.setLayout(layout) class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): - def __init__(self, parent, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None): + def __init__(self, parent, instance: Instance, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None): self.hg = hg - self.current_group = hg.root_group if current_group is None else current_group + self.current_group: Union[ObjectContainer, RelativeAtomGroup] = ObjectContainer(hg.root_group if current_group is None else current_group) + self.current_node: Union[ObjectContainer, RelativeAtomOrGroup, None] = ObjectContainer(None) + self.instance = instance + self._labels = {} + self.nodes: Dict[RelativeAtomOrGroup, PropChartHG] = {} + super().__init__(parent) + self.setScene(QGraphicsScene()) self._layout() + self.current_group.am_subscribe(self.on_new_group) + + def on_new_group(self, **kwargs): + self._layout() + + def navigate(self, node: RelativeAtomOrGroup): + if isinstance(node, RelativeAtomGroup): + self.current_group.am_obj = node + self.current_group.am_event() + elif isinstance(node, RelativeAtom): + self.instance.workspace.jump_to(node.atom.loc.ins_addr) + else: + assert False + + def select(self, node: Optional[RelativeAtomOrGroup]): + if not self.current_node.am_none: + self.nodes[self.current_node.am_obj].update() + self.current_node.am_obj = node + self.current_node.am_event() + if not self.current_node.am_none: + self.nodes[self.current_node.am_obj].update() def _layout(self): + for item in list(self.scene().items()): + self.scene().removeItem(item) + self.nodes.clear() + word_height = 40 byte_height = word_height / self.hg.kp.kb._project.arch.bytes - ng = self.hg.local_graph(self.current_group) + ng = self.hg.local_graph(self.current_group.am_obj) ug = networkx.Graph(ng) for edge in ug.edges: del ug.edges[edge]['prev'] del ug.edges[edge]['next'] lookup = list(ug.nodes()) reverse = {x: i for i, x in enumerate(lookup)} - charts = [PropChart(None, self.hg.prop(node), byte_height=byte_height) for node in lookup] + charts = [PropChartHG( + None, + self.hg.prop(node), + byte_height=byte_height, + nav=node, + hgw=self, + label=self.label(node)) for node in lookup] networkx.relabel_nodes(ug, reverse, copy=False) ag = networkx.drawing.nx_agraph.to_agraph(ug) @@ -73,7 +116,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): edge.attr['len'] = 50 / 72 + 200 / 72 ag.graph_attr['overlap'] = 'false' ag._layout() - ag.draw('/home/audrey/render.dot', prog='nop') + #ag.draw('/home/audrey/render.dot', prog='nop') for node in ag.nodes_iter(): idx = int(node) @@ -88,6 +131,36 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): arrow = PropArrow(None, charts[u], charts[v], towards_start, towards_end) self.scene().addItem(arrow) + self.nodes = {chart.nav: chart for chart in charts} + + def label(self, node: RelativeAtomOrGroup) -> str: + if node not in self._labels: + if isinstance(node, RelativeAtom): + self._labels[node] = str(node.atom) + else: + funcs = {self.hg.kp.atom_to_function(atom) for atom in self.hg.atoms(node) if atom not in self.hg.frontier} + if len(funcs) == 1: + func = funcs.pop() + self._labels[node] = self.instance.kb.functions[func].name + else: + self._labels[node] = 'Unknown group' + return self._labels[node] + + def mouseDoubleClickEvent(self, event): + super().mouseDoubleClickEvent(event) + if not event.isAccepted(): + event.accept() + if self.current_group.am_obj is not self.hg.root_group: + self.current_group.am_obj = self.current_group.parent + self.current_group.am_event() + + def mousePressEvent(self, event): + super().mousePressEvent(event) + item = self.itemAt(event.pos()) + if item is None: + event.accept() + self.select(None) + class PropChart(QGraphicsItem): def __init__(self, parent, prop: Prop, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5.): @@ -101,9 +174,11 @@ class PropChart(QGraphicsItem): self.padding_left = padding_left self.width = 0. self.height = 0. + self.objects = [] + self.object_bg: Optional[QGraphicsRectItem] = None + super().__init__(parent) - self.objects = [] self._layout_marks() def boundingRect(self): @@ -154,8 +229,9 @@ class PropChart(QGraphicsItem): self.height = row_allocated * self.byte_height + 2 * self.margin_y box = QGraphicsRectItem(QRectF(0, 0, self.width, self.height), self) box.setBrush(BOX_COLOR) - box.setPen(BOX_BORDER_COLOR) + box.setPen(QPen(BOX_BORDER_COLOR, BOX_BORDER_WIDTH)) self.objects.append(box) + self.object_bg = box for offset, kind, size, count, xpos, ypos in marks: rect = QGraphicsRectItem(QRectF(self.margin_x + self.padding_left + xpos * self.unit, self.margin_y + ypos * self.byte_height, count * self.unit, size * self.byte_height), self) @@ -170,6 +246,49 @@ class PropChart(QGraphicsItem): tick.setPos(QPointF(self.margin_x - tick.boundingRect().width(), self.margin_y + ypos * self.byte_height)) self.objects.append(tick) +class PropChartHG(PropChart): + def __init__(self, *args, nav: RelativeAtomOrGroup, hgw: HierarchicalGraphWidget, label: str, **kwargs): + self.nav = nav + self.hgw = hgw + self.label = label + self.object_label = None + super().__init__(*args, **kwargs) + + def mouseDoubleClickEvent(self, event): + event.accept() + self.hgw.navigate(self.nav) + + def mousePressEvent(self, event): + event.accept() + self.hgw.select(self.nav) + + def _layout_marks(self): + lbl_height = 20 + super()._layout_marks() + self.height += lbl_height + for obj in self.objects: + if obj is self.object_bg: + obj: QGraphicsRectItem + rect = obj.rect() + rect.setHeight(obj.rect().height() + lbl_height) + obj.setRect(rect) + else: + obj.setY(obj.y() + lbl_height) + + self.object_label = QGraphicsSimpleTextItem(self.label, self) + self.object_label.setPos(QPointF(self.margin_x, self.margin_y)) + self.object_label.setFont(TEXT_FONT) + self.objects.append(self.object_label) + + def paint(self, painter, option, widget=...): + if self.nav is self.hgw.current_node.am_obj: + self.object_bg.setPen(QPen(BOX_BORDER_SELECTED_COLOR, BOX_BORDER_SELECTED_WIDTH)) + else: + self.object_bg.setPen(QPen(BOX_BORDER_COLOR, BOX_BORDER_WIDTH)) + + super().paint(painter, option, widget) + + class PropArrow(QGraphicsItem): def __init__(self, parent, start: PropChart, end: PropChart, toward_start: bool, toward_end: bool, stroke=3., arrow_size=8.): self.start = start @@ -270,7 +389,7 @@ class PropArrow(QGraphicsItem): return QRectF(x - self.arrow_size, y - self.arrow_size, w + self.arrow_size * 2, h + self.arrow_size * 2) def paint(self, painter, option, widget=...): - brush = QBrush(ARROW_COLOR) + brush = QBrush(Conf.palette_windowtext) path = self.shape() painter.fillPath(path, brush) diff --git a/typetapper/relative_graph.py b/typetapper/relative_graph.py index af9201a..50f9ac1 100644 --- a/typetapper/relative_graph.py +++ b/typetapper/relative_graph.py @@ -15,6 +15,7 @@ l = logging.getLogger(__name__) class RelativeAtom: atom: Atom callstack: Tuple[int, ...] + rcallstack: Tuple[int, ...] @dataclass class RelativeAtomAttrs: @@ -39,7 +40,7 @@ class RelativeAtomGraph: self.predecessors = self.__graph.predecessors for atom in baseline: - relative = RelativeAtom(atom=atom, callstack=()) + relative = RelativeAtom(atom=atom, callstack=(), rcallstack=()) self._add_node(relative, OpSequence()) self.frontier.add(relative) # TODO ??? @@ -105,7 +106,7 @@ class RelativeAtomGraph: edge_ops: OpSequence, is_pred: bool, ) -> Optional[RelativeAtom]: - callstack = self._update_callstack(relatom.callstack, edge_cf, is_pred) + callstack, rcallstack = self._update_callstack(relatom.callstack, relatom.rcallstack, edge_cf, is_pred) if callstack is None: return None @@ -114,7 +115,7 @@ class RelativeAtomGraph: else: path = attrs.path + edge_ops - relsucc = RelativeAtom(atom=succ, callstack=callstack) + relsucc = RelativeAtom(atom=succ, callstack=callstack, rcallstack=rcallstack) res = self._add_node(relsucc, path) if is_pred: if not self.__graph.has_edge(relsucc, relatom): @@ -127,9 +128,10 @@ class RelativeAtomGraph: @staticmethod def _update_callstack( callstack: Tuple[int, ...], + rcallstack: Tuple[int, ...], cf: List[ControlFlowAction], reverse: bool - ) -> Optional[Tuple[int, ...]]: + ) -> Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]]: for directive in reversed(cf) if reverse else cf: if isinstance(directive, ControlFlowActionPop): pop = True @@ -147,7 +149,12 @@ class RelativeAtomGraph: if callstack[-1] != callsite: return None callstack = callstack[:-1] + else: + rcallstack = rcallstack + (callsite,) else: - callstack = callstack + (callsite,) + if rcallstack and rcallstack[-1] == callsite: + rcallstack = rcallstack[:-1] + else: + callstack = callstack + (callsite,) - return callstack + return callstack, rcallstack