diff --git a/typetapper/analysis.py b/typetapper/analysis.py index beb7965..628aee7 100644 --- a/typetapper/analysis.py +++ b/typetapper/analysis.py @@ -49,8 +49,6 @@ class TypeTapperAnalysis(angr.Analysis): for pred, attrs in self._cfg.graph.pred[node].items(): if attrs['jumpkind'] == 'Ijk_FakeRet': continue - if pred.block is None: - continue pred_addr = pred.addr pred_blockinfo = self.manager.block_info[pred_addr] diff --git a/typetapper/hierarchy_graph.py b/typetapper/hierarchy_graph.py index 9ebe901..13efa13 100644 --- a/typetapper/hierarchy_graph.py +++ b/typetapper/hierarchy_graph.py @@ -1,5 +1,6 @@ -from typing import Set, Union, TYPE_CHECKING, List, Optional, Dict, Iterable, Tuple -from itertools import pairwise +from typing import Set, Union, TYPE_CHECKING, List, Optional, Dict, Iterable, Tuple, Callable, Any +from itertools import pairwise, chain +from collections import defaultdict import networkx @@ -48,6 +49,40 @@ class HierarchicalGraph(RelativeAtomGraph): self._current_group = group return super().expand(relatom) + def expand_while(self, seeds: Iterable[RelativeAtom], group: RelativeAtomGroup, predicate: Callable[[RelativeAtom], bool]): + for seed in seeds: + if not predicate(seed): + raise ValueError("Predicate does not hold for seed %s", seed) + + queue = set(seeds) + while queue: + seed = queue.pop() + 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): + queue.add(adj) + + def expand_by_key(self, seed: RelativeAtomOrGroup, group: RelativeAtomGroup, key: Callable[[RelativeAtom], Any]): + queues = defaultdict(set) + for atom in self.atoms(seed): + if atom in self.frontier: + queues[key(atom)].add(atom) + else: + for adj in chain(self.predecessors(atom), self.successors(atom)): + if adj in self.frontier: + queues[key(adj)].add(adj) + + for keyed, queue in queues.items(): + subgroup = self.create_group([], group) + self.expand_while(queue, subgroup, lambda atom: key(atom) == keyed) + + def atoms(self, node: RelativeAtomOrGroup) -> Iterable[RelativeAtom]: + if isinstance(node, RelativeAtomGroup): + for child in node.children: + yield from self.atoms(child) + else: + yield node + def _prop_propagate(self, node: RelativeAtomOrGroup, add: bool): prop = self.prop(node) for parent in self._ancestry(node): @@ -95,6 +130,7 @@ class HierarchicalGraph(RelativeAtomGraph): def _add_group(self, parent: RelativeAtomGroup) -> RelativeAtomGroup: group = RelativeAtomGroup(self, parent) parent.children.add(group) + self.__graph.add_node(group) return group def _paths_through(self, item: RelativeAtomOrGroup) -> Iterable[Tuple[Optional[MultiGraphEdge], Optional[MultiGraphEdge]]]: @@ -282,6 +318,14 @@ class HierarchicalGraph(RelativeAtomGraph): self.move_node_in(node, group) return group + def destroy_group(self, group: RelativeAtomGroup): + if group is self._root_group: + raise ValueError("Cannot break root group") + for child in list(group.children): + self.move_node_out(child) + group.parent.children.remove(group) + self.__graph.remove_node(group) + def _parent(self, item: RelativeAtomOrGroup) -> RelativeAtomGroup: result = item.parent if isinstance(item, RelativeAtomGroup) else self._atom_parents[item] if result is None: diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index 2c036d2..62e07bd 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -50,28 +50,36 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): self._layout() def _layout(self): + word_height = 40 + byte_height = word_height / self.hg.kp.kb._project.arch.bytes + ng = self.hg.local_graph(self.current_group) 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)) for node in lookup] + charts = [PropChart(None, self.hg.prop(node), byte_height=byte_height) 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['width'] = charts[idx].width / 72 + node.attr['height'] = charts[idx].height / 72 node.attr['shape'] = 'box' - for _, _, attr in ag.edge_attr.iteritems(): - attr['len'] = 100 + for edge in ag.edges_iter(): + edge.attr['len'] = 50 / 72 + 200 / 72 + ag.graph_attr['overlap'] = 'false' ag._layout() + ag.draw('/home/audrey/render.dot', prog='nop') 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))) + chart.setPos(QPointF(float(posx) - chart.width / 2, -float(posy) - chart.height / 2)) self.scene().addItem(chart) for u, v in ug.edges: @@ -82,7 +90,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): 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.): + def __init__(self, parent, prop: Prop, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5.): self.prop = prop self.max_width = max_width self.default_unit = default_unit @@ -143,7 +151,7 @@ class PropChart(QGraphicsItem): 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 + 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) @@ -163,7 +171,7 @@ class PropChart(QGraphicsItem): 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.): + def __init__(self, parent, start: PropChart, end: PropChart, toward_start: bool, toward_end: bool, stroke=3., arrow_size=8.): self.start = start self.end = end self.toward_start = toward_start @@ -301,12 +309,15 @@ def arrowhead_vectors(start: QPointF, end: QPointF, stem_width: float, flare_wid norm_out = rotate(norm, pi / 2) norm_back = -norm + stem_width /= 2 + flare_width /= 2 + 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_back * flare_width * 2 + norm_out * flare_width, + end + norm_back * flare_width * 2 + norm_out * stem_width, + end + norm_back * flare_width * 2 - norm_out * stem_width, + end + norm_back * flare_width * 2 - norm_out * flare_width, end + -norm_out * stem_width ) diff --git a/typetapper/knowledge.py b/typetapper/knowledge.py index 0633bc4..2e9a600 100644 --- a/typetapper/knowledge.py +++ b/typetapper/knowledge.py @@ -1,11 +1,11 @@ -from typing import Dict +from typing import Dict, Union from collections import defaultdict import angr import networkx -from angr.sim_variable import SimVariable, SimRegisterVariable, SimMemoryVariable, SimStackVariable, \ - SimTemporaryVariable +from angr.sim_variable import SimVariable, SimRegisterVariable, SimMemoryVariable, SimStackVariable from .data import BlockInfo, RegisterAtom, MemoryAtom, TmpAtom, Atom +from .relative_graph import RelativeAtom from .hierarchy_graph import HierarchicalGraph class TypeTapperManager(angr.knowledge_plugins.plugin.KnowledgeBasePlugin): @@ -74,3 +74,11 @@ class TypeTapperManager(angr.knowledge_plugins.plugin.KnowledgeBasePlugin): def session(self, atom: Atom) -> HierarchicalGraph: return HierarchicalGraph(self, [atom]) + + def atom_to_function(self, atom: Union[Atom, RelativeAtom]): + if isinstance(atom, Atom): + loc = atom.loc + else: + loc = atom.atom.loc + + return self.cfg.get_any_node(loc.bbl_addr).function_address diff --git a/typetapper/plugin.py b/typetapper/plugin.py index 14f3eed..8a86ee9 100644 --- a/typetapper/plugin.py +++ b/typetapper/plugin.py @@ -40,6 +40,11 @@ class TypeTapper(BasePlugin): def _start(self, node: Atom): hg = self.kp.session(node) + start_relatoms = list(hg.frontier) + func = self.kp.atom_to_function(node) + func_group = hg.create_group(start_relatoms, hg.root_group) + hg.expand_while(start_relatoms, func_group, lambda atom: self.kp.atom_to_function(atom) == func) + hg.expand_by_key(func_group, hg.root_group, self.kp.atom_to_function) view = HierarchicalGraphView(self.workspace.main_instance, "center", hg) self.workspace.add_view(view) diff --git a/typetapper/relative_graph.py b/typetapper/relative_graph.py index 3e953ff..af9201a 100644 --- a/typetapper/relative_graph.py +++ b/typetapper/relative_graph.py @@ -34,6 +34,10 @@ class RelativeAtomGraph: self.__graph = networkx.DiGraph() self.frontier = set() # nodes present in self.graph but haven't had all their edges analyzed + self.has_edge = self.__graph.has_edge + self.successors = self.__graph.successors + self.predecessors = self.__graph.predecessors + for atom in baseline: relative = RelativeAtom(atom=atom, callstack=()) self._add_node(relative, OpSequence()) @@ -89,6 +93,7 @@ class RelativeAtomGraph: if res is not None: result.add(res) self.frontier.update(result) + self.frontier.remove(relatom) return result def _expand_single(