diff --git a/typetapper/data.py b/typetapper/data.py index e054045..c8e56e9 100644 --- a/typetapper/data.py +++ b/typetapper/data.py @@ -60,6 +60,19 @@ class ConstOffsetOp(Op): def invert(self): return ConstOffsetOp(-self.const) + def __init__(self, const: int): + # ummmmmm missing size + while const < -2**63: + const += 2**64 + while const > 2**63: + const -= 2**64 + object.__setattr__(self, 'const', const) + +@dataclass(frozen=True) +class NegOp(Op): + def invert(self): + return self + @dataclass(frozen=True) class VarOffsetOp(Op): var: Any @@ -120,19 +133,33 @@ def simplify_op_sequence(seq: List[Op]): cur = seq[i] if isinstance(cur, ConstOffsetOp) and cur.const == 0: seq.pop(i) + if i > 0: + i -= 1 continue nex = seq[i + 1] if i + 1 < len(seq) else None if isinstance(cur, ConstOffsetOp) and isinstance(nex, ConstOffsetOp): seq[i] = ConstOffsetOp(cur.const + nex.const) seq.pop(i + 1) + if i > 0: + i -= 1 continue if isinstance(cur, RefOp) and isinstance(nex, DerefOp) and cur.size == nex.size: seq.pop(i) seq.pop(i) + if i > 0: + i -= 1 continue if isinstance(cur, DerefOp) and isinstance(nex, RefOp) and cur.size == nex.size: seq.pop(i) seq.pop(i) + if i > 0: + i -= 1 + continue + if isinstance(cur, NegOp) and isinstance(nex, NegOp): + seq.pop(i) + seq.pop(i) + if i > 0: + i -= 1 continue i += 1 @@ -140,6 +167,7 @@ def simplify_op_sequence(seq: List[Op]): # noinspection PyArgumentList class DataKind(IntEnum): + GenericData = auto() Int = auto() Float = auto() Pointer = auto() @@ -232,8 +260,12 @@ class LiveData: def new_const(cls, value: int, size: int, codeloc: CodeLoc) -> 'LiveData': return cls([(ConstAtom(codeloc, size, value), OpSequence())], value, size) - def appended(self, op: Op, size: int) -> 'LiveData': - return LiveData([(atom, seq.appended(op)) for atom, seq in self.sources], self.const, size) + def appended(self, op: Op, size: int, const: Optional[int]=None) -> 'LiveData': + return LiveData( + [(atom, seq.appended(op)) for atom, seq in self.sources], + self.const if const is None else const, + size + ) def unioned(self, other: 'LiveData', size: int, const: Optional[int]=None) -> 'LiveData': return LiveData(self.sources + other.sources, const, size) diff --git a/typetapper/engine.py b/typetapper/engine.py index f1407f7..6ba82fc 100644 --- a/typetapper/engine.py +++ b/typetapper/engine.py @@ -97,7 +97,10 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): return self.load(addr, size, endness) def load(self, addr, size, endness): - addr.prop_self(DataKind.Pointer, self.graph) + prop = Prop() + prop.self_data[DataKind.Pointer] += 1 + prop.struct_data[0][size][DataKind.GenericData] += 1 + addr.prop(prop, self.graph) mem_atom = MemoryAtom(self.codeloc, size, endness) self.blockinfo.atoms.append(mem_atom) @@ -137,8 +140,11 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): const = addend0 + addend1 * sign else: const = None - input0 = args[0].appended(ConstOffsetOp(addend1 * sign) if addend1 is not None else VarOffsetOp(args[1]), size) - input1 = args[1].appended(ConstOffsetOp(addend0) if addend0 is not None else VarOffsetOp(args[0]), size) + neg1 = args[1] + if sign == -1: + neg1 = neg1.appended(NegOp(), neg1.size, -addend1 if addend1 is not None else None) + input0 = args[0].appended(ConstOffsetOp(addend1 * sign) if addend1 is not None else VarOffsetOp(neg1), size) + input1 = neg1.appended(ConstOffsetOp(addend0) if addend0 is not None else VarOffsetOp(args[0]), size) result = input0.unioned(input1, size, const) else: result = LiveData.new_null(size) @@ -193,6 +199,7 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin): def store(self, addr, data, endness): prop = Prop() prop.self_data[DataKind.Pointer] += 1 + prop.struct_data[0][data.size][DataKind.GenericData] += 1 addr.prop(prop, self.graph) mem_atom = MemoryAtom(self.codeloc, data.size, endness) diff --git a/typetapper/hierarchy_graph.py b/typetapper/hierarchy_graph.py index c23ae27..df3dfea 100644 --- a/typetapper/hierarchy_graph.py +++ b/typetapper/hierarchy_graph.py @@ -31,6 +31,44 @@ class HierarchicalGraph(RelativeAtomGraph): super().__init__(kp, baseline) + def _check_invariants(self): + for node in self.__graph.nodes: + if isinstance(node, RelativeAtom): + assert node in self._atom_parents + for succ in self.successors(node): + for edge in pairwise(self._hierarchy_path(node, succ)): + assert edge[1] in self.__graph.succ[edge[0]] + else: + assert isinstance(node, RelativeAtomGroup) + assert node is self._root_group or node.parent is self._root_group or node.parent in self.__graph.nodes + + pool = set(self.__graph.edges) + while pool: + edge = pool.pop() + + next_edge = self.__graph.edges[edge]['next'] + while next_edge is not None: + assert next_edge in pool + pool.remove(next_edge) + next_edge = self.__graph.edges[next_edge]['next'] + + prev_edge = self.__graph.edges[edge]['prev'] + while prev_edge is not None: + assert prev_edge in pool + pool.remove(prev_edge) + prev_edge = self.__graph.edges[prev_edge]['prev'] + + def check_invariants(self): + return + try: + self._check_invariants() + except AssertionError as e: + import ipdb + while True: + ipdb.post_mortem(e.__traceback__) + ipdb.set_trace() + + @property def root_group(self) -> RelativeAtomGroup: return self._root_group @@ -63,7 +101,7 @@ class HierarchicalGraph(RelativeAtomGraph): 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]) -> Dict[Any, RelativeAtomGroup]: + def expand_by_key(self, seed: RelativeAtomOrGroup, group: RelativeAtomGroup, key: Callable[[RelativeAtom], Any], existing_groups: Optional[Dict[Any, RelativeAtomGroup]]=None) -> Dict[Any, RelativeAtomGroup]: queues = defaultdict(set) for atom in self.atoms(seed): if atom in self.frontier: @@ -73,10 +111,13 @@ class HierarchicalGraph(RelativeAtomGraph): if adj in self.frontier: queues[key(adj)].add(adj) - result = {} + result = existing_groups if existing_groups is not None else {} for keyed, queue in queues.items(): - subgroup = self.create_group([], group) - result[keyed] = subgroup + if keyed in result: + subgroup = result[keyed] + else: + subgroup = self.create_group([], group) + result[keyed] = subgroup self.expand_while(queue, subgroup, lambda atom: key(atom) == keyed) return result @@ -103,6 +144,7 @@ class HierarchicalGraph(RelativeAtomGraph): self._current_group.children.add(relatom) self.__graph.add_node(relatom) self._prop_propagate(relatom, True) + self.check_invariants() return res def _add_edge(self, relatom1: RelativeAtom, relatom2: RelativeAtom): @@ -118,6 +160,7 @@ class HierarchicalGraph(RelativeAtomGraph): self.__graph.edges[edge]['prev'] = None prev_edge = edge self.__graph.edges[prev_edge]['next'] = None + self.check_invariants() def _remove_node(self, relatom: RelativeAtom): super()._remove_node(relatom) @@ -127,15 +170,18 @@ class HierarchicalGraph(RelativeAtomGraph): self._atom_parents[relatom].children.remove(relatom) del self._atom_parents[relatom] self.__graph.remove_node(relatom) + self.check_invariants() def _remove_edge(self, relatom1: RelativeAtom, relatom2: RelativeAtom): super()._remove_edge(relatom1, relatom2) self.__graph.remove_edges_from(pairwise(self._hierarchy_path(relatom1, relatom2))) + self.check_invariants() def _add_group(self, parent: RelativeAtomGroup) -> RelativeAtomGroup: group = RelativeAtomGroup(self, parent) parent.children.add(group) self.__graph.add_node(group) + self.check_invariants() return group def _paths_through(self, item: RelativeAtomOrGroup) -> Iterable[Tuple[Optional[MultiGraphEdge], Optional[MultiGraphEdge]]]: @@ -145,7 +191,7 @@ class HierarchicalGraph(RelativeAtomGraph): yield from ((None, (item, succ, key)) for succ, keys in self.__graph.succ[item].items() for key in keys) yield from (((pred, item, key), None) for pred, keys in self.__graph.pred[item].items() for key in keys) else: - yield from ((self.__graph.edges[(item, succ, key)]['prev'], (item, succ, key)) for succ in self.__graph.succ[item] for key in succ) + yield from ((self.__graph.edges[(item, succ, key)]['prev'], (item, succ, key)) for succ, keys in self.__graph.succ[item].items() for key in keys) def move_node_out(self, item: RelativeAtomOrGroup): # item will now be a neighbor of its parent @@ -197,7 +243,7 @@ class HierarchicalGraph(RelativeAtomGraph): self.__graph.edges[parent_item_edge][prev_next[outward]] = parent_neighbor_edge self.__graph.edges[parent_neighbor_edge][prev_next[not outward]] = parent_item_edge self.__graph.edges[parent_neighbor_edge][prev_next[outward]] = further_edge - if further_edge: self.__graph.edges[further_edge][prev_next[not outward]] = further_edge + if further_edge: self.__graph.edges[further_edge][prev_next[not outward]] = parent_neighbor_edge else: # contract assert outer is parent @@ -222,6 +268,7 @@ class HierarchicalGraph(RelativeAtomGraph): else: self._atom_parents[item] = new_parent self._prop_propagate(item, True) + self.check_invariants() def move_node_in(self, item: RelativeAtomOrGroup, new_parent: RelativeAtomGroup): if item is self._root_group: @@ -274,7 +321,7 @@ class HierarchicalGraph(RelativeAtomGraph): self.__graph.edges[parent_item_edge][prev_next[outward]] = parent_outer_edge self.__graph.edges[parent_outer_edge][prev_next[not outward]] = parent_item_edge self.__graph.edges[parent_outer_edge][prev_next[outward]] = further_edge - if further_edge: self.__graph.edges[further_edge][prev_next[not outward]] = further_edge + if further_edge: self.__graph.edges[further_edge][prev_next[not outward]] = parent_outer_edge else: # contract assert further_edge @@ -298,6 +345,7 @@ class HierarchicalGraph(RelativeAtomGraph): else: self._atom_parents[item] = new_parent self._prop_propagate(item, True) + self.check_invariants() def move_node(self, node: RelativeAtomOrGroup, new_parent: RelativeAtomGroup): new_parent_ancestry = set(self._ancestry(new_parent)) diff --git a/typetapper/hierarchy_graph_view.py b/typetapper/hierarchy_graph_view.py index 66f2237..7825e59 100644 --- a/typetapper/hierarchy_graph_view.py +++ b/typetapper/hierarchy_graph_view.py @@ -36,7 +36,7 @@ BOX_BORDER_SELECTED_COLOR = QColor(0x60, 0x60, 0xff) BOX_BORDER_WIDTH = 1 BOX_BORDER_SELECTED_WIDTH = 2 PLUS_COLOR = QColor(0x40, 0xc0, 0xb0) -PLUS_ICON_COLOR = QColor(0x50, 0xd0, 0xf0) +PLUS_ICON_COLOR = QColor(0xff, 0xff, 0xff) PLUS_HOVERED_COLOR = QColor(0x40, 0xc0, 0xc0) PLUS_BORDER_COLOR = QColor(0x30, 0xc0, 0x50) PLUS_ICON_FONT = QFont("default", 12) @@ -51,16 +51,18 @@ def kind_to_color(kind: DataKind) -> QColor: return QColor(0x20, 0x20, 0xf0) elif kind == DataKind.Float: return QColor(0x10, 0xe0, 0x10) + elif kind == DataKind.GenericData: + return QColor(0xc0, 0xc0, 0xc0) else: raise ValueError(kind) class HierarchicalGraphView(BaseView): - def __init__(self, instance, default_docking_position, hg: HierarchicalGraph, faux_frontier: Set[RelativeAtomOrGroup], *args, current_group: Optional[RelativeAtomGroup]=None, **kwargs): + def __init__(self, instance, default_docking_position, hg: HierarchicalGraph, faux_frontier: Set[RelativeAtomOrGroup], ignore: Set[RelativeAtomOrGroup], *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, instance, hg, current_group, faux_frontier) + self.graph = HierarchicalGraphWidget(self, instance, hg, current_group, faux_frontier, ignore) layout.addWidget(self.graph) self.setLayout(layout) @@ -71,7 +73,7 @@ class ExternNode: RelativeAtomOrGroupOrExtern = Union[RelativeAtomOrGroup, ExternNode] class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): - def __init__(self, parent, instance: Instance, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None, faux_frontier: Optional[Set[RelativeAtomOrGroup]]=None): + def __init__(self, parent, instance: Instance, hg: HierarchicalGraph, current_group: Optional[RelativeAtomGroup]=None, faux_frontier: Optional[Set[RelativeAtomOrGroup]]=None, ignore: Optional[Set[RelativeAtomOrGroup]]=None): self.hg = hg self.current_group: Union[ObjectContainer, RelativeAtomGroup] = ObjectContainer(hg.root_group if current_group is None else current_group) self.selected_nodes: Union[ObjectContainer, Set[RelativeAtomOrGroupOrExtern]] = ObjectContainer(set()) @@ -82,6 +84,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): self.expansion_qnodes: List[HGNode] = [] self.expansion_qedges: List[HGArrow] = [] self.expansion_models: List[RelativeAtomOrGroup] = [] + self.drop_target: Optional['PropChartHG'] = None self.layout_prog = 'neato' self.layout_social_distancing = 200 / DPI @@ -102,6 +105,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): self._selection_event_occurring = False self._layout_event_occurring = False self.faux_frontier = faux_frontier or set() + self.ignore = ignore or set() super().__init__(parent) @@ -121,7 +125,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): @property def selected_node(self) -> Optional[RelativeAtomOrGroupOrExtern]: - return self.selected_nodes[0] if len(self.selected_nodes) == 1 else None + return next(iter(self.selected_nodes)) if len(self.selected_nodes) == 1 else None def navigate(self, node: RelativeAtomOrGroupOrExtern): if isinstance(node, RelativeAtomGroup): @@ -149,6 +153,11 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): return any(adj in self.faux_frontier for adj in chain(self.ng.succ[model], self.ng.pred[model])) def begin_expand(self, model: RelativeAtomOrGroup): + self.expansion_models = [adj for adj in set(chain(self.ng.succ[model], self.ng.pred[model])) if adj in self.faux_frontier] + if not self.expansion_models: + l.error("How'd you do that") + return + qnode = self.qnodes[self.ug_reverse[model]] self.scene().clearSelection() qnode.setSelected(True) @@ -163,14 +172,11 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): edge.setOpacity(0.4) byte_height = self._byte_height - self.expansion_models = [adj for adj in chain(self.ng.succ[model], self.ng.pred[model]) if adj in self.faux_frontier] self.expansion_qnodes = [PropChartHG( None, - prop=self.hg.prop(other_model), byte_height=byte_height, model=other_model, hgw=self, - label=self.label(other_model) ) for other_model in self.expansion_models] self.expansion_qedges = [HGArrow( None, @@ -236,12 +242,118 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): if accepted_qnode: self.faux_frontier.remove(accepted_qnode.model) - newgroups = self.hg.expand_by_key(accepted_qnode.model, self.current_group.am_obj, self.hg.kp.atom_to_function) - self.faux_frontier.update(newgroups.values()) + existing_groups = [(g, self.label(g)) for g in self.current_group.children if isinstance(g, RelativeAtomGroup)] + existing_groups = { + self.instance.kb.functions[lbl].addr: g + for (g, lbl) in existing_groups + if lbl in self.instance.kb.functions + } + newgroups = self.hg.expand_by_key(accepted_qnode.model, self.current_group.am_obj, self.hg.kp.atom_to_function, dict(existing_groups)) + for k, g in newgroups.items(): + if k in existing_groups: + continue + if k in self.hg.kp.kb.functions and self.hg.kp.kb.functions[k].alignment: + self.ignore.add(g) + else: + self.faux_frontier.add(g) self._layout(0) + def set_drop_target(self, item: Optional['PropChartHG']): + if item is self.drop_target: + return + if self.drop_target is not None: + self.drop_target.deemphasize() + self.drop_target = item + if item is not None: + item.emphasize() + for selected in self.scene().selectedItems(): + if isinstance(selected, PropChartHG): + selected.setOpacity(0.5 if item is not None else 1) + selected.setZValue(10 if item is not None else 0) + + def confirm_drop(self): + if self.drop_target is None: + return + + model = self.drop_target.model + if isinstance(model, RelativeAtom): + # new group time! + idx = self.ug_reverse[model] + old_model = model + model = self.hg.create_group([old_model], self.current_group.am_obj) + self.ug_lookup[idx] = model + self.ug_reverse[model] = idx + self.drop_target.model = model + self._labels[model] = 'Manual group' + edge_replacement: Optional[Tuple[RelativeAtomOrGroupOrExtern, RelativeAtomOrGroupOrExtern]] = None + for key, edge in self.drop_target.edges.items(): + if isinstance(key, ExternNode): + extern_qnode = edge.start if edge.end is self.drop_target else edge.end + assert isinstance(extern_qnode, HGExtern) + edge = extern_qnode.edges.pop(old_model) + new_extern_model = ExternNode(model) + extern_qnode.model = new_extern_model + idx = self.ug_reverse.pop(key) + self.ug_reverse[new_extern_model] = idx + self.ug_lookup[idx] = new_extern_model + extern_qnode.edges[self.drop_target.model] = edge + edge_replacement = key, new_extern_model + else: + if edge.start is self.drop_target: + edge.end.edges.pop(old_model) + edge.end.edges[model] = edge + else: + edge.start.edges.pop(old_model) + edge.start.edges[model] = edge + if edge_replacement is not None: + self.drop_target.edges[edge_replacement[1]] = self.drop_target.edges.pop(edge_replacement[0]) + + + assert isinstance(model, RelativeAtomGroup) + for item in self.scene().selectedItems(): + if isinstance(item, PropChartHG): + self.hg.move_node_in(item.model, model) + self._rekey_externs(item, self.drop_target) + + self.drop_target.deemphasize() + self.drop_target._layout_marks() + self.drop_target = None + self._layout(0) + + # private interfaces + def _rekey_externs(self, old_qnode: 'PropChartHG', new_qnode: 'PropChartHG'): + old_extern_model = ExternNode(old_qnode.model) + new_extern_model = ExternNode(new_qnode.model) + if old_extern_model in old_qnode.edges: + extern_qnode = self.qnodes[self.ug_reverse[old_extern_model]] + if new_extern_model not in new_qnode.edges: + edge = extern_qnode.edges.pop(old_qnode.model) + assert edge is old_qnode.edges.pop(old_extern_model) + extern_qnode.model = new_extern_model + idx = self.ug_reverse.pop(old_extern_model) + self.ug_reverse[new_extern_model] = idx + self.ug_lookup[idx] = new_extern_model + extern_qnode.edges[new_qnode.model] = edge + new_qnode.edges[new_extern_model] = edge + if edge.start is old_qnode: + edge.start = new_qnode + else: + edge.end = new_qnode + edge.layout() + else: + old_edge = old_qnode.edges[old_extern_model] + new_edge = new_qnode.edges[new_extern_model] + swap = (old_edge.start is old_qnode) ^ (new_edge.start is new_qnode) + if swap: + to_start, to_end = old_edge.toward_end, old_edge.toward_start + else: + to_end, to_start = old_edge.toward_end, old_edge.toward_start + + new_edge.toward_start |= to_start + new_edge.toward_end |= to_end + @property def _byte_height(self): word_height = 40 @@ -260,7 +372,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): self.ng = self.hg.local_graph(current_group) self.ug = networkx.Graph(self.ng) for model in list(self.ug.nodes): - if model in self.faux_frontier: + if model in self.faux_frontier or model in self.ignore: self.ug.remove_node(model) for edge in self.ug.edges: del self.ug.edges[edge]['prev'] @@ -284,7 +396,7 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): ibr = QRectF(0, -1000, 1000, 1000) for i, node in enumerate(self.ug_lookup): if isinstance(node, ExternNode): - qnode = old_nodes.pop(node.adj_to, None) + qnode = old_nodes.pop(node, None) if qnode is None: qnode = HGExtern(None, model=node, hgw=self) entering.add(qnode) @@ -295,11 +407,9 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): if qnode is None: qnode = PropChartHG( None, - prop=self.hg.prop(node), byte_height=byte_height, model=node, hgw=self, - label=self.label(node) ) entering.add(qnode) qnode.center = rand_point(ibr, i) @@ -309,26 +419,30 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): qnode.set_expandable(self.can_expand(model)) for u, v in self.ug.edges: - lu = self.ug_lookup[u] - lv = self.ug_lookup[v] + model_u = self.ug_lookup[u] + model_v = self.ug_lookup[v] qu = self.qnodes[u] qv = self.qnodes[v] - if lu in qv.edges: - assert lv in qu.edges + if model_u in qv.edges: + assert model_v in qu.edges + edge = qv.edges[model_u] + edge.prepareGeometryChange() + edge.toward_start = edge.start.model in self.ng.succ[edge.end.model] + edge.toward_end = edge.end.model in self.ng.succ[edge.start.model] else: - assert lv not in qu.edges - if isinstance(lu, ExternNode): - towards_start = current_group in self.ng.succ[lv] - towards_end = lv in self.ng.succ[current_group] - elif isinstance(lv, ExternNode): - towards_start = lu in self.ng.succ[current_group] - towards_end = current_group in self.ng.succ[lu] + assert model_v not in qu.edges + if isinstance(model_u, ExternNode): + towards_start = current_group in self.ng.succ[model_v] + towards_end = model_v in self.ng.succ[current_group] + elif isinstance(model_v, ExternNode): + towards_start = model_u in self.ng.succ[current_group] + towards_end = current_group in self.ng.succ[model_u] else: - towards_end = lv in self.ng.succ[lu] - towards_start = lu in self.ng.succ[lv] + towards_end = model_v in self.ng.succ[model_u] + towards_start = model_u in self.ng.succ[model_v] arrow = HGArrow(None, self.qnodes[u], self.qnodes[v], towards_start, towards_end) - qu.edges[lv] = arrow - qv.edges[lu] = arrow + qu.edges[model_v] = arrow + qv.edges[model_u] = arrow self.scene().addItem(arrow) arrow.enter() self.qedges.append(arrow) @@ -420,6 +534,8 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): def mouseReleaseEvent(self, event): self._mouse_is_down = False + if self.drop_target is not None: + self.confirm_drop() super().mouseReleaseEvent(event) # objectcontainer event handlers @@ -454,6 +570,8 @@ class HierarchicalGraphWidget(QZoomableDraggableGraphicsView): elif not self._selection_event_occurring: self.selected_nodes.am_obj = {model for model, qnode in zip(self.ug_lookup, self.qnodes) if qnode.isSelected()} self.selected_nodes.am_event(src='qt') + for qedge in self.qedges: + qedge.update() class AnimatableItem(QGraphicsItem): @@ -482,6 +600,7 @@ class HGNode(AnimatableItem): self.setFlag(QGraphicsItem.ItemIsSelectable) self.setFlag(QGraphicsItem.ItemIsMovable) self.setFlag(QGraphicsItem.ItemSendsGeometryChanges) + self.setZValue(0) # public interfaces @@ -548,8 +667,7 @@ class HGNode(AnimatableItem): class PropChart(QGraphicsItem): - def __init__(self, *args, prop: Prop, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5., **kwargs): - self.prop = prop + def __init__(self, parent, max_width=200., default_unit=10., byte_height=10., margin_x=15., margin_y=5., padding_left=5.): self.max_width = max_width self.default_unit = default_unit self.unit = default_unit @@ -562,10 +680,14 @@ class PropChart(QGraphicsItem): self.objects = [] self.object_bg: Optional[QGraphicsRectItem] = None - super().__init__(*args, **kwargs) + super().__init__(parent) self._layout_marks() + @property + def prop(self) -> Prop: + raise NotImplementedError() + def boundingRect(self): return QRectF(0, 0, self.width, self.height) @@ -573,6 +695,10 @@ class PropChart(QGraphicsItem): pass def _layout_marks(self): + self.objects.clear() + for item in list(self.childItems()): + self.scene().removeItem(item) + 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) @@ -617,19 +743,40 @@ class PropChart(QGraphicsItem): rect.setPen(QColor(0x10, 0x10, 0x10)) self.objects.append(rect) + tick_width = 0 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)) + width = tick.boundingRect().width() + tick_width = max(width, tick_width) + tick.setPos(QPointF(self.margin_x - width, self.margin_y + ypos * self.byte_height)) self.objects.append(tick) + if tick_width > self.margin_x: + shift = tick_width - self.margin_x + for obj in self.objects: + if obj is self.object_bg: + rect = obj.rect() + rect.setWidth(rect.width() + shift) + obj.setRect(rect) + else: + obj.setX(obj.x() + shift) + self.width += shift + class PropChartHG(HGNode, PropChart): - def __init__(self, *args, label: str, **kwargs): - self.label = label + def __init__(self, parent, **kwargs): self.object_label = None self.object_btn: Optional[PlusButton] = None - super().__init__(*args, **kwargs) + super().__init__(parent, **kwargs) + + @property + def label(self): + return self.hgw.label(self.model) + + @property + def prop(self): + return self.hgw.hg.prop(self.model) @property def center(self) -> QPointF: @@ -663,7 +810,7 @@ class PropChartHG(HGNode, PropChart): self.object_btn = PlusButton(self, self.begin_expand, 16) self.object_btn.setPos(self.width - 18, 2) - self.setTransformOriginPoint(self.center) + self.setTransformOriginPoint(self.center - self.pos()) def set_expandable(self, expandable: bool): self.object_btn.setVisible(expandable) @@ -714,6 +861,29 @@ class PropChartHG(HGNode, PropChart): node.attr['height'] = self.height / DPI node.attr['shape'] = 'box' + def mouseMoveEvent(self, event): + if self.isSelected() and bool(event.buttons() & Qt.LeftButton): + for item in self.scene().items(event.scenePos()): + if isinstance(item, PropChartHG) and not item.isSelected(): + self.hgw.set_drop_target(item) + break + else: + self.hgw.set_drop_target(None) + super().mouseMoveEvent(event) + + def emphasize(self, duration=100): + self._start_animation(self._animate_emphasize, duration) + + def _animate_emphasize(self, value): + self.setScale(1 + 0.3 * value) + + def deemphasize(self, duration=100): + self._start_animation(self._animate_deemphasize, duration) + + def _animate_deemphasize(self, value): + self.setScale(1 + 0.3 * (1 - value)) + + class HGExtern(HGNode): def __init__(self, *args, radius=28., **kwargs): self.radius = radius @@ -778,6 +948,8 @@ class HGArrow(AnimatableItem): super().__init__(parent) + self.setZValue(-10) + def enter(self, duration=250): self._start_animation(self._animate_enter, duration) @@ -843,7 +1015,10 @@ class HGArrow(AnimatableItem): 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(Conf.palette_windowtext) + color = QColor(Conf.palette_windowtext) + if self.start.isSelected() or self.end.isSelected(): + color.setRedF(1 - color.redF()) + brush = QBrush(color) path = self.shape() painter.fillPath(path, brush) diff --git a/typetapper/knowledge.py b/typetapper/knowledge.py index b9c0aea..764296c 100644 --- a/typetapper/knowledge.py +++ b/typetapper/knowledge.py @@ -27,7 +27,7 @@ class TypeTapperManager(angr.knowledge_plugins.plugin.KnowledgeBasePlugin): if block is None: raise LookupError("No such block %#x" % addr) - return self.block_info[addr] + return self.block_info[block.addr] def lookup_reg(self, addr: int, register: str) -> RegisterAtom: blockinfo = self._block_info(addr) @@ -42,12 +42,12 @@ class TypeTapperManager(angr.knowledge_plugins.plugin.KnowledgeBasePlugin): continue if atom.slot_name == register and atom.loc.ins_addr == addr: return atom - for info in blockinfo.inputs.values(): - if info.atom.name == register: - return info.atom - for info in blockinfo.inputs.values(): - if info.atom.slot_name == register: - return info.atom + for atom in blockinfo.inputs.values(): + if atom.name == register: + return atom + for atom in blockinfo.inputs.values(): + if atom.slot_name == register: + return atom raise LookupError("Cannot find register %s in instruction %#x" % (register, addr)) def lookup_mem(self, addr: int) -> MemoryAtom: diff --git a/typetapper/plugin.py b/typetapper/plugin.py index 340a461..5010807 100644 --- a/typetapper/plugin.py +++ b/typetapper/plugin.py @@ -22,6 +22,7 @@ class TypeTapper(BasePlugin): cfg = self.workspace.main_instance.cfg tt = self.workspace.main_instance.project.analyses[TypeTapperAnalysis](cfg) self.kp = tt.manager + self._start(self.kp.lookup_reg(0x4144a5, "rdi")) def build_context_menu_node(self, node): # this is bad lol @@ -46,6 +47,12 @@ class TypeTapper(BasePlugin): hg.expand_while(start_relatoms, func_group, lambda atom: self.kp.atom_to_function(atom) == func) newgroups = hg.expand_by_key(func_group, hg.root_group, self.kp.atom_to_function) - view = HierarchicalGraphView(self.workspace.main_instance, "center", hg, set(newgroups.values())) + view = HierarchicalGraphView( + self.workspace.main_instance, + "center", + hg, + set(g for k, g in newgroups.items() if not self.kp.kb.functions[k].alignment), + set(g for k, g in newgroups.items() if self.kp.kb.functions[k].alignment), + ) self.workspace.add_view(view) self.workspace.raise_view(view) diff --git a/typetapper/relative_graph.py b/typetapper/relative_graph.py index 73aed09..cb57ad7 100644 --- a/typetapper/relative_graph.py +++ b/typetapper/relative_graph.py @@ -4,7 +4,8 @@ import logging import networkx -from .data import Atom, Prop, OpSequence, ControlFlowActionPop, ControlFlowActionPush, ControlFlowAction, RefOp, DerefOp +from .data import Atom, Prop, OpSequence, ControlFlowActionPop, ControlFlowActionPush, ControlFlowAction, RefOp, \ + DerefOp, ConstOffsetOp if TYPE_CHECKING: from .knowledge import TypeTapperManager @@ -28,7 +29,12 @@ class RelativeAtomAttrs: # TODO has_gone_down if self.path != other.path: # TODO unifications - pass + # compute a stupidity score + score_0 = sum(0 if isinstance(op, (ConstOffsetOp, RefOp, DerefOp)) else 1 for op in self.path.ops) + score_1 = sum(0 if isinstance(op, (ConstOffsetOp, RefOp, DerefOp)) else 1 for op in other.path.ops) + if score_1 < score_0: + self.path = other.path + return True return False class RelativeAtomGraph: @@ -51,6 +57,7 @@ class RelativeAtomGraph: If relatom is not present in the graph, add it. If it is present in the graph, merge the new information into its attrs """ + from .data import RegisterAtom newattrs = RelativeAtomAttrs( prop=self.kp.graph.nodes[relatom.atom].get('prop', Prop()).transform(path.invert()), path=path, @@ -109,12 +116,8 @@ class RelativeAtomGraph: edge_ops: OpSequence, is_pred: bool, ) -> Optional[RelativeAtom]: - goes_down = any(isinstance(op, DerefOp) for op in edge_ops.ops) - goes_up = any(isinstance(op, DerefOp) for op in edge_ops.ops) - if is_pred: - goes_down, goes_up = goes_up, goes_down - if attrs.has_gone_down and goes_up: - return None + #if str(relatom.atom) == 'rsp @ 0x4144a8' and str(succ) == 'MEM @ 0x4144d9': + # import ipdb; ipdb.set_trace() weh = self._update_callstack(relatom.callstack, relatom.rcallstack, edge_cf, is_pred) if weh is None: @@ -124,8 +127,23 @@ class RelativeAtomGraph: path = attrs.path path += edge_ops.invert() if is_pred else edge_ops + derefs = sum(isinstance(op, DerefOp) for op in path.ops) + refs = sum(isinstance(op, RefOp) for op in path.ops) + goes_down = any(isinstance(op, DerefOp) for op in edge_ops.ops) + goes_up = any(isinstance(op, RefOp) for op in edge_ops.ops) + if is_pred: + goes_down, goes_up = goes_up, goes_down + reset_down = len(path.ops) == 0 + + if derefs > 0 and refs > 0: + return None + if derefs > 1: + return None + if goes_up and attrs.has_gone_down: + return None + relsucc = RelativeAtom(atom=succ, callstack=callstack, rcallstack=rcallstack) - res = self._add_node(relsucc, path, attrs.has_gone_down or goes_down) + res = self._add_node(relsucc, path, (attrs.has_gone_down or goes_down) and not reset_down) if is_pred: if not self.__graph.has_edge(relsucc, relatom): self._add_edge(relsucc, relatom) @@ -161,7 +179,9 @@ class RelativeAtomGraph: else: rcallstack = rcallstack + (callsite,) else: - if rcallstack and rcallstack[-1] == callsite: + if rcallstack: + if rcallstack[-1] != callsite: + return None rcallstack = rcallstack[:-1] else: callstack = callstack + (callsite,)