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