From 6be9ba072749b27cee4144a91e43f5125bae9d47 Mon Sep 17 00:00:00 2001 From: Audrey Dutcher Date: Sat, 29 Oct 2022 14:58:34 -0700 Subject: [PATCH] we have rendered... something --- typetapper/__init__.py | 1 + typetapper/analysis.py | 14 +- typetapper/data.py | 4 +- typetapper/hierarchy_graph_view.py | 325 +++++++++++++++++++++++++++++ typetapper/knowledge.py | 10 + typetapper/plugin.py | 46 ++++ 6 files changed, 391 insertions(+), 9 deletions(-) create mode 100644 typetapper/hierarchy_graph_view.py create mode 100644 typetapper/plugin.py diff --git a/typetapper/__init__.py b/typetapper/__init__.py index e69de29..dd2a003 100644 --- a/typetapper/__init__.py +++ b/typetapper/__init__.py @@ -0,0 +1 @@ +from .plugin import TypeTapper diff --git a/typetapper/analysis.py b/typetapper/analysis.py index ec2daff..beb7965 100644 --- a/typetapper/analysis.py +++ b/typetapper/analysis.py @@ -1,21 +1,21 @@ +from typing import Optional import logging from collections import OrderedDict import angr from angr import Block -from angr.analyses.cfg import CFGBase -from angr.knowledge_plugins.cfg import CFGNode +from angr.knowledge_plugins.cfg import CFGNode, CFGModel from .engine import TypeTapperEngine from .knowledge import TypeTapperManager l = logging.getLogger(__name__) -class TypeTapper(angr.Analysis): - def __init__(self, cfg: CFGBase, tmp_atoms: bool=False): - self._cfg = cfg +class TypeTapperAnalysis(angr.Analysis): + def __init__(self, cfg: Optional[CFGModel]=None, tmp_atoms: bool=False): + self._cfg = cfg if cfg is not None else self.kb.cfgs.get_most_accurate() self.manager = self.kb.request_knowledge(TypeTapperManager) - self.manager.cfg = cfg.model + self.manager.cfg = cfg self._engine = TypeTapperEngine(self.project, self.manager, tmp_atoms=tmp_atoms) if not self._cfg.normalized: @@ -44,7 +44,7 @@ class TypeTapper(angr.Analysis): block_addr = next(iter(reversed(queue.keys()))) queue.pop(block_addr) node_blockinfo = self.manager.block_info[block_addr] - node = self._cfg.model.get_any_node(block_addr) + node = self._cfg.get_any_node(block_addr) fakeret_addr = next((pred.addr for pred, attrs in self._cfg.graph.pred[node].items() if attrs['jumpkind'] == 'Ijk_FakeRet'), None) for pred, attrs in self._cfg.graph.pred[node].items(): if attrs['jumpkind'] == 'Ijk_FakeRet': diff --git a/typetapper/data.py b/typetapper/data.py index e6ae9db..cb909dc 100644 --- a/typetapper/data.py +++ b/typetapper/data.py @@ -1,6 +1,6 @@ from typing import Tuple, Any, List, Set, Optional, Dict from collections import defaultdict, Counter -from enum import Enum, auto +from enum import IntEnum, auto from dataclasses import dataclass, field import copy @@ -127,7 +127,7 @@ def simplify_op_sequence(seq: List[Op]): # noinspection PyArgumentList -class DataKind(Enum): +class DataKind(IntEnum): Int = auto() Float = auto() Pointer = auto() diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py new file mode 100644 index 0000000..2c036d2 --- /dev/null +++ b/typetapper/hierarchy_graph_view.py @@ -0,0 +1,325 @@ +from typing import Tuple, Optional +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.QtWidgets import QGraphicsItem, QGraphicsRectItem, QGraphicsSimpleTextItem, QHBoxLayout, QGraphicsScene + +import networkx + +from angrmanagement.ui.views import BaseView +from angrmanagement.ui.widgets.qgraph import QZoomableDraggableGraphicsView + +from .hierarchy_graph import HierarchicalGraph, RelativeAtomGroup +from .data import Prop, DataKind + +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) + +def kind_to_color(kind: DataKind) -> QColor: + if kind == DataKind.Pointer: + return QColor(0xe0, 0x10, 0x10) + elif kind == DataKind.Int: + return QColor(0x20, 0x20, 0xf0) + elif kind == DataKind.Float: + return QColor(0x10, 0xe0, 0x10) + else: + raise ValueError(kind) + +class HierarchicalGraphView(BaseView): + def __init__(self, instance, default_docking_position, hg: HierarchicalGraph, *args, current_group: Optional[RelativeAtomGroup]=None, **kwargs): + super().__init__("typetapper", instance, default_docking_position, *args, **kwargs) + self.base_caption = "TypeTapper" + + layout = QHBoxLayout(self) + self.graph = HierarchicalGraphWidget(self, hg, current_group) + layout.addWidget(self.graph) + self.setLayout(layout) + +class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): + def __init__(self, parent, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None): + self.hg = hg + self.current_group = hg.root_group if current_group is None else current_group + super().__init__(parent) + self.setScene(QGraphicsScene()) + self._layout() + + def _layout(self): + ng = self.hg.local_graph(self.current_group) + ug = networkx.Graph(ng) + lookup = list(ug.nodes()) + reverse = {x: i for i, x in enumerate(lookup)} + charts = [PropChart(None, self.hg.prop(node)) for node in lookup] + networkx.relabel_nodes(ug, reverse, copy=False) + + ag = networkx.drawing.nx_agraph.to_agraph(ug) + for node in ag.nodes_iter(): + idx = int(node) + node.attr['width'] = charts[idx].width + node.attr['height'] = charts[idx].height + node.attr['shape'] = 'box' + for _, _, attr in ag.edge_attr.iteritems(): + attr['len'] = 100 + ag._layout() + + for node in ag.nodes_iter(): + idx = int(node) + chart = charts[idx] + posx, posy = node.attr['pos'].split(',') + chart.setPos(QPointF(float(posx), float(posy))) + self.scene().addItem(chart) + + for u, v in ug.edges: + towards_end = lookup[v] in ng.succ[lookup[u]] + towards_start = lookup[u] in ng.succ[lookup[v]] + arrow = PropArrow(None, charts[u], charts[v], towards_start, towards_end) + self.scene().addItem(arrow) + + +class PropChart(QGraphicsItem): + def __init__(self, parent, prop: Prop, max_width=300., default_unit=10., byte_height=20., margin_x=15., margin_y=5., padding_left=5.): + self.prop = prop + self.max_width = max_width + self.default_unit = default_unit + self.unit = default_unit + self.byte_height = byte_height + self.margin_x = margin_x + self.margin_y = margin_y + self.padding_left = padding_left + self.width = 0. + self.height = 0. + super().__init__(parent) + + self.objects = [] + self._layout_marks() + + def boundingRect(self): + return QRectF(0, 0, self.width, self.height) + + def paint(self, painter, option, widget=...): + pass + + @property + def center(self) -> QPointF: + pos = self.pos() + pos.setX(pos.x() + self.width / 2) + pos.setY(pos.y() + self.height / 2) + return pos + + def _layout_marks(self): + data = [(offset, kind, size, count) for offset, v1 in self.prop.struct_data.items() for size, v2 in v1.items() for kind, count in v2.items()] + data.sort() + cols_allocated = defaultdict(int) + row_allocated = 0 + offset_allocated = 0 + marks = [] + ticks = [] + for offset, kind, size, count in data: + xpos = max(cols_allocated[suboffset] for suboffset in range(offset, offset + size)) + for suboffset in range(offset, offset + size): + cols_allocated[suboffset] = xpos + count + + if offset >= offset_allocated: + ypos = row_allocated + row_allocated += size + offset_allocated = offset + size + ticks.append((offset, ypos)) + else: + ypos = row_allocated - (offset_allocated - offset) + stretch = max(0, offset + size - offset_allocated) + row_allocated += stretch + offset_allocated += stretch + + marks.append((offset, kind, size, count, xpos, ypos)) + + total_width = max(cols_allocated.values()) if cols_allocated else 0 + + if total_width * self.unit > self.max_width: + self.unit = self.max_width / total_width + + self.width = self.max_width + 2 * self.margin_x + self.padding_left + self.height = row_allocated + 2 * self.margin_y + box = QGraphicsRectItem(QRectF(0, 0, self.width, self.height), self) + box.setBrush(BOX_COLOR) + box.setPen(BOX_BORDER_COLOR) + self.objects.append(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) + rect.setBrush(kind_to_color(kind)) + rect.setPen(QColor(0x10, 0x10, 0x10)) + self.objects.append(rect) + + for offset, ypos in ticks: + tick = QGraphicsSimpleTextItem(str(offset), self) + tick.setBrush(TEXT_COLOR) + tick.setFont(TEXT_FONT) + tick.setPos(QPointF(self.margin_x - tick.boundingRect().width(), self.margin_y + ypos * self.byte_height)) + self.objects.append(tick) + +class PropArrow(QGraphicsItem): + def __init__(self, parent, start: PropChart, end: PropChart, toward_start: bool, toward_end: bool, stroke=2., arrow_size=4.): + self.start = start + self.end = end + self.toward_start = toward_start + self.toward_end = toward_end + self.stroke = stroke + self.arrow_size = arrow_size + self._percentage = 1. + + super().__init__(parent) + self.layout() + + @property + def percentage(self): + return self._percentage + + @percentage.setter + def percentage(self, value): + self._percentage = value + self.layout() + + def layout(self): + center_start = self.start.center + center_end = self.end.center + delta_x = center_end.x() - center_start.x() + delta_y = center_end.y() - center_start.y() + slope = delta_y / delta_x if delta_x else float('inf') + + start_corner_slope = self.start.height / self.start.width + end_corner_slope = self.end.height / self.end.width + + start_horiz = start_corner_slope > slope > -start_corner_slope + end_horiz = end_corner_slope > slope > -end_corner_slope + + if start_horiz: + start_right = center_end.x() > center_start.x() + if start_right: + start_ex = self.start.width / 2 + else: + start_ex = -self.start.width / 2 + start_ey = start_ex * slope + else: + start_bottom = center_end.y() > center_start.y() + if start_bottom: + start_ey = self.start.height / 2 + else: + start_ey = -self.start.height / 2 + start_ex = start_ey / slope + + if end_horiz: + end_right = center_start.x() > center_end.x() + if end_right: + end_ex = self.end.width / 2 + else: + end_ex = -self.end.width / 2 + end_ey = end_ex * slope + else: + end_bottom = center_start.y() > center_end.y() + if end_bottom: + end_ey = self.end.height / 2 + else: + end_ey = -self.end.height / 2 + end_ex = end_ey / slope + + start_pt = QPointF(center_start.x() + start_ex, center_start.y() + start_ey) + end_pt = QPointF(center_end.x() + end_ex, center_end.y() + end_ey) + + if self.percentage != 1.: + end_pt = start_pt + (end_pt - start_pt) * self.percentage + + self.prepareGeometryChange() + self.setPos(start_pt) + self.rel_end = end_pt - start_pt + + + def orient(self, start: PropChart): + if start is self.start: + pass + elif start is self.end: + self.start, self.end = self.end, self.start + self.toward_start, self.toward_end = self.toward_end, self.toward_start + self.layout() + else: + raise ValueError("Orienting from node which is not associated with edge") + + def boundingRect(self): + x = 0 + y = 0 + w = self.rel_end.x() + h = self.rel_end.y() + + if w < 0: + x, w = w, -w + if h < 0: + y, h = h, -h + + 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) + path = self.shape() + painter.fillPath(path, brush) + + def shape(self): + path = QPainterPath(QPointF()) + s_tip1, s_flare1, s_base1, s_base2, s_flare2, s_tip2 = arrowhead_vectors(self.rel_end, QPointF(), self.stroke, self.arrow_size) + e_tip1, e_flare1, e_base1, e_base2, e_flare2, e_tip2 = arrowhead_vectors(QPointF(), self.rel_end, self.stroke, self.arrow_size) + if self.toward_start: + path.lineTo(s_flare1) + path.lineTo(s_base1) + else: + path.lineTo(s_tip1) + + if self.toward_end: + path.lineTo(e_base2) + path.lineTo(e_flare2) + path.lineTo(self.rel_end) + path.lineTo(e_flare1) + path.lineTo(e_base1) + else: + path.lineTo(e_tip2) + path.lineTo(e_tip1) + + if self.toward_start: + path.lineTo(s_base2) + path.lineTo(s_flare2) + else: + path.lineTo(s_tip2) + path.lineTo(QPointF()) + path.closeSubpath() + return path + +def arrowhead_vectors(start: QPointF, end: QPointF, stem_width: float, flare_width: float) -> Tuple[QPointF, QPointF, QPointF, QPointF, QPointF, QPointF]: + direction = end - start + norm = normalized(direction) + norm_out = rotate(norm, pi / 2) + norm_back = -norm + + return ( + end + norm_out * stem_width, + end + norm_back * flare_width + norm_out * flare_width, + end + norm_back * flare_width + norm_out * stem_width, + end + norm_back * flare_width - norm_out * stem_width, + end + norm_back * flare_width - norm_out * flare_width, + end + -norm_out * stem_width + ) + +def normalized(pt: QPointF) -> QPointF: + length = sqrt(pt.x() ** 2 + pt.y() ** 2) + if length: + return pt / length + else: + return QPointF() + +def rotate(pt: QPointF, angle: float) -> QPointF: + cost = cos(angle) + sint = sin(angle) + x = pt.x() + y = pt.y() + return QPointF(x * cost - y * sint, x * sint + y * cost) diff --git a/typetapper/knowledge.py b/typetapper/knowledge.py index b4a2fb9..0633bc4 100644 --- a/typetapper/knowledge.py +++ b/typetapper/knowledge.py @@ -3,6 +3,8 @@ from collections import defaultdict import angr import networkx +from angr.sim_variable import SimVariable, SimRegisterVariable, SimMemoryVariable, SimStackVariable, \ + SimTemporaryVariable from .data import BlockInfo, RegisterAtom, MemoryAtom, TmpAtom, Atom from .hierarchy_graph import HierarchicalGraph @@ -62,5 +64,13 @@ class TypeTapperManager(angr.knowledge_plugins.plugin.KnowledgeBasePlugin): return atom raise LookupError("Cannot find tmp %d in instruction %#x. Are your temp numbers based on the .normalized_block?" % (tmp, addr)) + def lookup_variable(self, addr: int, var: SimVariable): + if isinstance(var, SimRegisterVariable): + return self.lookup_reg(addr, self.kb._project.arch.translate_register_name(var.reg, var.size)) + elif isinstance(var, (SimMemoryVariable, SimStackVariable)): + return self.lookup_mem(addr) + else: + raise LookupError("Cannot handle var of type %s" % type(var)) + def session(self, atom: Atom) -> HierarchicalGraph: return HierarchicalGraph(self, [atom]) diff --git a/typetapper/plugin.py b/typetapper/plugin.py new file mode 100644 index 0000000..14f3eed --- /dev/null +++ b/typetapper/plugin.py @@ -0,0 +1,46 @@ +from typing import Optional + +from angr.analyses.decompiler.structured_codegen.c import CVariable +from angrmanagement.plugins import BasePlugin +from angrmanagement.ui.views import CodeView + +from .analysis import TypeTapperAnalysis +from .data import Atom +from .knowledge import TypeTapperManager +from .hierarchy_graph_view import HierarchicalGraphView + + +class TypeTapper(BasePlugin): + def __init__(self, workspace): + super().__init__(workspace) + self.kp: Optional[TypeTapperManager] = None + + MENU_BUTTONS = ["TypeTapper initial analysis"] + + def handle_click_menu(self, idx): + assert idx == 0 + cfg = self.workspace.main_instance.cfg + tt = self.workspace.main_instance.project.analyses[TypeTapperAnalysis](cfg) + self.kp = tt.manager + + def build_context_menu_node(self, node): + # this is bad lol + code_view = self.workspace.view_manager.current_tab + if not isinstance(code_view, CodeView): + return + addr = code_view._textedit.get_closest_insaddr(node) + + if self.kp is not None and isinstance(node, CVariable): + try: + atom = self.kp.lookup_variable(addr, node.variable) + except: + pass + else: + yield 'Start TypeTapper', lambda: self._start(atom) + + def _start(self, node: Atom): + hg = self.kp.session(node) + + view = HierarchicalGraphView(self.workspace.main_instance, "center", hg) + self.workspace.add_view(view) + self.workspace.raise_view(view)