WHACK WHACK WHACK WHACK WHACK
- Fix many many bugs - Add GenericData use designation - Add invariant checking for hg - Add drag and drop for group merging - Better data path pruning based on ref/deref patterns
This commit is contained in:
parent
930b0ba0eb
commit
1b1122226b
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,)
|
||||
|
|
|
|||
Loading…
Reference in New Issue