def _make_wire_item(x1, y1, x2, y2):
path = QPainterPath()
path.moveTo(x1, y1)
path.lineTo(x2, y2)
path_item = QGraphicsPathItem(path)
stroker = QPainterPathStroker(QPen(Qt.black, 5, c=Qt.RoundCap))
stroke_path = stroker.createStroke(path)
stroke_item = QGraphicsPathItem(stroke_path)
path_item.setPen(QPen(Qt.white, 5, c=Qt.RoundCap))
stroke_item.setPen(QPen(Qt.black, 1.5, c=Qt.RoundCap))
group = QGraphicsItemGroup()
group.addToGroup(path_item)
group.addToGroup(stroke_item)
group.setFlag(QGraphicsItem.ItemIsSelectable)
return group
class QPOIViewer(QWidget):
"""
POI Viewer QWidget
"""
TAG_SPACING = 50
LEGEND_X = -50
LEGEND_Y = 0
LEGEND_WIDTH = 10
TRACE_FUNC_X = 0
TRACE_FUNC_Y = 0
TRACE_FUNC_WIDTH = 50
TRACE_FUNC_MINHEIGHT = 1000
TAB_HEADER_SIZE = 40
MAX_WINDOW_SIZE = 500
MARK_X = LEGEND_X
MARK_WIDTH = TRACE_FUNC_X - LEGEND_X + TRACE_FUNC_WIDTH
MARK_HEIGHT = 1
POIID_COLUMN = 0
CRASH_COLUMN = 1
CATEGORY_COLUMN = 2
DIAGNOSE_COLUMN = 3
COLUMN_FIELD = ['id', 'bbl', 'category', 'diagnose']
def __init__(self, workspace, parent=None, diagnose_handler=None):
super().__init__(parent=parent)
self.workspace = workspace
self._diagnose_handler = diagnose_handler
self.mark = None
self.legend = None
self.legend_height = 0
self.legend_img = None
self.trace_func_unit_height = 0
self.trace_func = None
self.trace_id = None
self.tabView = None
self.traceView = None
self.traceScene = None
self.POITraceTab = None
self.multiPOITab : QWidget = None
self.multiPOIList : QTableWidget = None
self.mark = None
self.curr_position = 0
self._use_precise_position = False
self._selected_traces = []
self._selected_poi = None
self._init_widgets()
self.selected_ins.am_subscribe(self._subscribe_select_ins)
self.poi_trace.am_subscribe(self._subscribe_set_trace)
self.multi_poi.am_subscribe(self._subscribe_add_poi)
self.multiPOIList.cellDoubleClicked.connect(self._on_cell_double_click)
self.multiPOIList.itemChanged.connect(self._on_diagnose_change)
#
# Forwarding properties
#
@property
def disasm_view(self):
"""
Get the current disassembly view (if there is one), or create a new one as needed.
"""
view = self.workspace.view_manager.current_view_in_category("disassembly")
if view is None:
view = self.workspace._get_or_create_disassembly_view()
return view
@property
def poi_trace(self):
return self.workspace.instance.poi_trace
@property
def multi_poi(self):
return self.workspace.instance.multi_poi
@property
def selected_ins(self):
return self.disasm_view.infodock.selected_insns
def _init_widgets(self):
_l.debug("QPOI Viewer Initiating")
self.tabView = QTabWidget() # QGraphicsView()
self.tabView.setContentsMargins(0, 0, 0, 0)
#
# POI trace Tab
#
self.POITraceTab = QWidget()
self.POITraceTab.setContentsMargins(0, 0, 0, 0)
singleLayout = QVBoxLayout()
singleLayout.setSpacing(0)
singleLayout.setContentsMargins(0, 0, 0, 0)
self.traceView = QGraphicsView()
self.traceScene = QGraphicsScene()
self.traceView.setScene(self.traceScene)
singleLayout.addWidget(self.traceView)
self.POITraceTab.setLayout(singleLayout)
#
# multiPOI Tab
#
self.multiPOITab = QMultiPOITab(self)
# self.multiPOITab = QWidget()
multiLayout = QVBoxLayout()
multiLayout.setSpacing(0)
multiLayout.setContentsMargins(0, 0, 0, 0)
self.multiPOIList = QTableWidget(0, 4) # row, col
self.multiPOIList.setHorizontalHeaderItem(0, QTableWidgetItem("ID"))
self.multiPOIList.setHorizontalHeaderItem(1, QTableWidgetItem("Crash Point"))
self.multiPOIList.setHorizontalHeaderItem(2, QTableWidgetItem("Tag"))
self.multiPOIList.setHorizontalHeaderItem(3, QTableWidgetItem("Diagnose"))
self.multiPOIList.horizontalHeader().setStretchLastSection(True)
self.multiPOIList.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeToContents)
self.multiPOIList.setSelectionBehavior(QAbstractItemView.SelectRows)
multiLayout.addWidget(self.multiPOIList)
self.multiPOITab.setLayout(multiLayout)
self.tabView.addTab(self.multiPOITab, "POI List")
self.tabView.addTab(self.POITraceTab, "POI Trace")
self.POI_TRACE = 1
self.MULTI_POI = 0
layout = QVBoxLayout()
layout.addWidget(self.tabView)
layout.setContentsMargins(0, 0, 0, 0)
self.setLayout(layout)
self.show()
def _reset(self):
self.traceScene.clear() #clear items
self.mark = None
self.legend = None
self.legend_height = 0
self.trace_func = QGraphicsItemGroup()
self.trace_id = QGraphicsItemGroup()
self.traceScene.addItem(self.trace_func)
self.hide()
#
# Event
#
def _on_cell_double_click(self, row, _):
_l.debug("row %d is double clicked", row)
first_cell = self.multiPOIList.item(row, 0)
if first_cell is None:
return
poi_id = first_cell.text()
poi = self.multi_poi.am_obj.get_poi_by_id(poi_id)
if poi is None:
return
# sanity checks
if not isinstance(poi, dict):
return
if 'output' not in poi or not isinstance(poi['output'], dict):
return
if 'bbl_history' not in poi['output']:
return
trace = poi['output']['bbl_history']
if self._selected_poi != poi_id and trace is not None:
# render the trace
self.poi_trace.am_obj = TraceStatistics(self.workspace, trace, trace_id=poi_id)
# show the trace statistic in POI trace
self.poi_trace.am_event()
# show covered basic blocks and functions
self.multi_poi.am_obj.reload_heatmap(poi_id)
# redraw function view
view = self.workspace.view_manager.first_view_in_category('functions')
if view is not None:
view.refresh()
# redraw disassembly view
view = self.workspace.view_manager.first_view_in_category('disassembly')
if view is not None:
view.redraw_current_graph()
if trace is not None:
# switch to POI trace tab
self.tabView.setCurrentIndex(self.POI_TRACE)
self._selected_poi = poi_id
second_cell = self.multiPOIList.item(row, 1)
crash_addr = None
if second_cell is not None:
try:
crash_addr = int(second_cell.text(), 16)
except ValueError:
pass
if crash_addr is not None:
# show the crashing address
view = self.workspace.view_manager.first_view_in_category('disassembly')
if view is not None:
crash_func = self._get_func_from_addr(crash_addr)
if crash_func is not None:
self.workspace.on_function_selected(crash_func)
self.selected_ins.clear()
self.selected_ins.update([crash_addr])
self.selected_ins.am_event()
view.current_graph.show_instruction(crash_addr)
def _on_diagnose_change(self, item: QTableWidgetItem):
column = item.column()
row = item.row()
poi_id = self.multiPOIList.item(row, self.POIID_COLUMN).text()
content = item.text()
original_content = self.multi_poi.am_obj.get_content_by_id_column(poi_id, column)
_l.debug('updaing %s, content: %s, original: %s', poi_id, content, original_content)
if not self._is_identical(content, original_content):
updated_poi = self.multi_poi.update_poi(poi_id, column, content)
self._diagnose_handler.submit_updated_poi(poi_id, updated_poi)
def _subscribe_add_poi(self):
_l.debug('add a poi to multi poi list')
if self.multi_poi.am_none:
self.multi_poi.am_obj = MultiPOI(self.workspace)
poi_ids = self.multi_poi.am_obj.get_all_poi_ids()
self.multiPOIList.clearContents()
self._populate_poi_table(self.multiPOIList, poi_ids)
self.show()
def _subscribe_set_trace(self):
_l.debug('on set trace in poi trace viewer')
self._reset()
if self.poi_trace.am_none:
return
_l.debug('minheight: %d, count: %d', self.TRACE_FUNC_MINHEIGHT,
self.poi_trace.count)
if self.poi_trace.count <= 0:
_l.warning("No valid addresses found in trace to show. Check base address offsets?")
self.poi_trace.am_obj = None
self.poi_trace.am_event()
return
if self.TRACE_FUNC_MINHEIGHT < self.poi_trace.count * 15:
self.trace_func_unit_height = 15
show_func_tag = True
else:
self.trace_func_unit_height = self.TRACE_FUNC_MINHEIGHT / self.poi_trace.count
show_func_tag = True
self.legend_height = int(self.poi_trace.count * self.trace_func_unit_height)
self._show_trace_func(show_func_tag)
self._show_legend()
self._set_mark_color()
self._refresh_multi_list()
# boundingSize = self.traceScene.itemsBoundingRect().width()
# windowSize = boundingSize
# if boundingSize > self.MAX_WINDOW_SIZE:
# windowSize = self.MAX_WINDOW_SIZE
# self.traceScene.setSceneRect(self.traceScene.itemsBoundingRect()) #resize
# if windowSize > self.width():
# self.setMinimumWidth(windowSize)
self.show()
def _subscribe_select_ins(self, **kwargs): # pylint: disable=unused-argument
if self.poi_trace.am_none:
return
if self.mark is not None:
for i in self.mark.childItems():
self.mark.removeFromGroup(i)
self.traceScene.removeItem(i)
self.traceScene.removeItem(self.mark)
self.mark = QGraphicsItemGroup()
self.traceScene.addItem(self.mark)
if self.selected_ins:
addr = next(iter(self.selected_ins))
positions = self.poi_trace.get_positions(addr)
if positions: #if addr is in list of positions
# handle case where insn was selected from disas view
if not self._use_precise_position:
self.curr_position = positions[0] - self.poi_trace.count
for p in positions:
color = self._get_mark_color(p, self.poi_trace.count)
y = self._get_mark_y(p)
if p == self.poi_trace.count + self.curr_position: #add thicker line for 'current' mark
self.mark.addToGroup(self.traceScene.addRect(self.MARK_X, y, self.MARK_WIDTH,
self.MARK_HEIGHT*4,
QPen(QColor('black')), QBrush(color)))
else:
self.mark.addToGroup(self.traceScene.addRect(self.MARK_X, y, self.MARK_WIDTH,
self.MARK_HEIGHT, QPen(color), QBrush(color)))
self.traceScene.update() #force redraw of the traceScene
self.scroll_to_position(self.curr_position)
def _get_func_from_addr(self, addr):
if self.workspace.instance.cfg.am_none:
return None
bbl = self.workspace.instance.cfg.get_any_node(addr, anyaddr=True)
function_addr = bbl.function_address
return self.workspace.instance.project.kb.functions.get(function_addr)
def _populate_poi_table(self, view, poi_ids):
view.clearContents()
view.setRowCount(len(poi_ids))
row = 0 #start after label row
for poi_id in poi_ids:
poi = self.multi_poi.am_obj.get_poi_by_id(poi_id)
_l.debug('populating poi: %s', poi)
category = poi['category']
output = poi['output']
crash_addr = output['bbl']
if crash_addr is not None:
crash = hex(crash_addr)
else:
crash = None
diagnose = output.get('diagnose')
_l.debug('poi_ids: %s', poi_ids)
_l.debug('current poi id: %s', poi_id)
self._set_item(view, row, self.POIID_COLUMN, poi_id, editable=False)
self._set_item(view, row, self.CRASH_COLUMN, crash, editable=True)
self._set_item(view, row, self.CATEGORY_COLUMN, category, editable=True)
self._set_item(view, row, self.DIAGNOSE_COLUMN, diagnose, editable=True)
row += 1
_l.debug('poi_ids: %s', poi_ids)
@staticmethod
def _set_item(view, row, column, text, editable=True):
if not text:
text = ""
item = QTableWidgetItem(text)
if not editable:
item.setFlags(item.flags() ^ Qt.ItemIsEditable)
view.setItem(row, column, item)
def _refresh_multi_list(self):
multiPOI = self.multi_poi.am_obj
trace_ids = multiPOI.get_all_poi_ids()
self.multiPOIList.clearContents()
self._populate_poi_table(self.multiPOIList, trace_ids)
if self._selected_traces and self.multiPOIList.rowCount() > 0:
self.multiPOIList.item(0, 0).setSelected(True)
self.multiPOIList.item(0, 1).setSelected(True)
else:
for row in range(self.multiPOIList.rowCount()):
item = self.multiPOIList.item(row, 0)
inputItem = self.multiPOIList.item(row, 1)
if item.text() in self._selected_traces:
item.setSelected(True)
inputItem.setSelected(True)
self.multi_poi.am_event()
def _on_tab_change(self):
multiPOI = self.multi_poi.am_obj
if self.tabView.currentIndex() == self.MULTI_POI:
multiPOI.is_active_tab = True
self._refresh_multi_list()
elif self.tabView.currentIndex() == self.POI_TRACE:
multiPOI = self.multi_poi.am_obj
multiPOI.is_active_tab = False
# self._show_trace_ids()
def scroll_to_position(self, position):
relative_pos = self.poi_trace.count + position
y_offset = self._get_mark_y(relative_pos)
scrollValue = 0
if y_offset > 0.5 * self.traceView.size().height():
scrollValue = y_offset - 0.5 * self.traceView.size().height()
scrollValue = min(scrollValue, self.traceView.verticalScrollBar().maximum())
self.traceView.verticalScrollBar().setValue(scrollValue)
self._use_precise_position = False
def jump_next_insn(self):
# for some reason indexing is done backwards
if self.curr_position + self.poi_trace.count < self.poi_trace.count - 1:
self.curr_position += 1
self._use_precise_position = True
bbl_addr = self.poi_trace.get_bbl_from_position(self.curr_position)
func = self.poi_trace.get_func_from_position(self.curr_position)
self._jump_bbl(func, bbl_addr)
def jump_prev_insn(self):
if self.curr_position + self.poi_trace.count > 0:
self.curr_position -= 1
self._use_precise_position = True
bbl_addr = self.poi_trace.get_bbl_from_position(self.curr_position)
func = self.poi_trace.get_func_from_position(self.curr_position)
self._jump_bbl(func, bbl_addr)
def mousePressEvent(self, event):
button = event.button()
pos = self._to_logical_pos(event.pos())
if button == Qt.LeftButton and self.tabView.currentIndex() == self.POI_TRACE and self._at_legend(pos):
func = self._get_func_from_y(pos.y())
bbl_addr = self._get_bbl_from_y(pos.y())
self._use_precise_position = True
self.curr_position = self._get_position(pos.y())
self._jump_bbl(func, bbl_addr)
def _jump_bbl(self, func, bbl_addr):
disasm_view = self.disasm_view
if disasm_view is not None:
all_insn_addrs = self.workspace.instance.project.factory.block(bbl_addr).instruction_addrs
# TODO: replace this with am_events perhaps?
self.workspace.on_function_selected(func)
self.selected_ins.clear()
self.selected_ins.update(all_insn_addrs)
self.selected_ins.am_event()
# TODO: this ought to happen automatically as a result of the am_event
disasm_view.current_graph.show_instruction(bbl_addr)
def _get_mark_color(self, i, total):
relative_gradient_pos = i * 1000 // total
return self.legend_img.pixelColor(self.LEGEND_WIDTH // 2,
relative_gradient_pos)
def _get_mark_y(self, i):
return self.TRACE_FUNC_Y + self.trace_func_unit_height * i
def _show_trace_func(self, show_func_tag=True):
x = self.TRACE_FUNC_X
y = self.TRACE_FUNC_Y
prev_name = None
for position in self.poi_trace.trace_func:
func_name = position.func_name
color = self.poi_trace.get_func_color(func_name)
self.trace_func.addToGroup(self.traceScene.addRect(x, y,
self.TRACE_FUNC_WIDTH, self.trace_func_unit_height,
QPen(color), QBrush(color)))
if show_func_tag is True and func_name != prev_name:
tag = self.traceScene.addText(func_name, QFont("Source Code Pro", 7))
tag.setPos(x + self.TRACE_FUNC_WIDTH +
self.TAG_SPACING, y -
tag.boundingRect().height() // 2)
self.trace_func.addToGroup(tag)
anchor = self.traceScene.addLine(
self.TRACE_FUNC_X + self.TRACE_FUNC_WIDTH, y,
x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING, y)
self.trace_func.addToGroup(anchor)
prev_name = func_name
y += self.trace_func_unit_height
@staticmethod
def _make_legend_gradient(x1, y1, x2, y2):
gradient = QLinearGradient(x1, y1, x2, y2)
gradient.setColorAt(0.0, Qt.red)
gradient.setColorAt(0.4, Qt.yellow)
gradient.setColorAt(0.6, Qt.green)
gradient.setColorAt(0.8, Qt.blue)
gradient.setColorAt(1.0, Qt.darkBlue)
return gradient
def _show_legend(self):
pen = QPen(Qt.transparent)
gradient = self._make_legend_gradient(self.LEGEND_X, self.LEGEND_Y,
self.LEGEND_X, self.LEGEND_Y + self.legend_height)
brush = QBrush(gradient)
self.legend = self.traceScene.addRect(self.LEGEND_X, self.LEGEND_Y,
self.LEGEND_WIDTH, self.legend_height, pen, brush)
reference_gradient = self._make_legend_gradient(0, 0, self.LEGEND_WIDTH, 1000)
base_img = QImage(self.LEGEND_WIDTH, 1000, QImage.Format.Format_ARGB32)
p = QPainter(base_img)
p.fillRect(base_img.rect(),reference_gradient)
self.legend_img = base_img #reference shade
def _set_mark_color(self):
_l.debug('trace count: %d', self.poi_trace.count)
for p in range(self.poi_trace.count):
color = self._get_mark_color(p, self.poi_trace.count)
self.poi_trace.set_mark_color(p, color)
def _at_legend(self, pos):
x = pos.x()
y = pos.y()
return self.TRACE_FUNC_X + self.LEGEND_X < x < self.traceView.width() and \
self.TRACE_FUNC_Y < y < self.TRACE_FUNC_Y + self.legend_height
def _to_logical_pos(self, pos):
x_offset = self.traceView.horizontalScrollBar().value()
y_offset = self.traceView.verticalScrollBar().value()
return QPoint(pos.x() + x_offset, pos.y() + y_offset)
def _get_position(self, y):
y_relative = y - self.legend_height - self.TAB_HEADER_SIZE
return int(y_relative // self.trace_func_unit_height)
def _get_bbl_from_y(self, y):
position = self._get_position(y)
return self.poi_trace.get_bbl_from_position(position)
def _get_func_from_y(self, y):
position = self._get_position(y)
func = self.poi_trace.get_func_from_position(position)
return func
#
# Context Menu
#
def menu_add_empty_poi(self):
_l.debug('adding a new empty poi item')
if self._diagnose_handler.get_image_id() is None:
QMessageBox.warning(self.workspace.main_window,
"No CHESS target available",
"No angr project is loaded, or you did not associate the current project with a CHESS "
"target. Please load a binary and associate it with a CHESS target before creating "
"POIs.")
return
poi_id = str(uuid4())
if self.multi_poi.am_none:
self.multi_poi.am_obj = MultiPOI(self.workspace)
empty_poi = deepcopy(EMPTY_POI)
self.multi_poi.add_poi(poi_id, empty_poi)
self.multi_poi.am_event()
self._diagnose_handler.submit_updated_poi(poi_id, empty_poi)
def menu_remove_poi(self):
items = self.multiPOIList.selectedItems()
row = items.pop().row()
poi_id = self.multiPOIList.item(row, 0).text()
_l.debug('removing ID %s', poi_id)
self.multi_poi.remove_poi(poi_id)
self.multi_poi.am_event()
@staticmethod
def _is_identical(content, original_content):
if content == original_content:
return True
if content == '' and original_content is None:
return True
try:
if int(content, 16) == int(original_content):
return True
except (TypeError, ValueError):
return False
return False
class QBlock(QCachedGraphicsItem):
TOP_PADDING = 5
BOTTOM_PADDING = 5
LEFT_PADDING = 10
RIGHT_PADDING = 10
SPACING = 0
def __init__(self,
workspace,
func_addr,
disasm_view,
disasm,
infodock,
addr,
cfg_nodes,
out_branches,
scene,
parent=None,
container=None):
super().__init__(parent=parent, container=container)
# initialization
self.workspace = workspace
self.func_addr = func_addr
self.disasm_view = disasm_view
self.disasm = disasm
self.infodock = infodock
self.variable_manager = infodock.variable_manager
self.addr = addr
self.cfg_nodes = cfg_nodes
self.out_branches = out_branches
self._scene = scene
self.margins = QMarginsF(self.LEFT_PADDING, self.TOP_PADDING,
self.RIGHT_PADDING, self.BOTTOM_PADDING)
self._config = Conf
self.objects = [] # instructions and labels
self._block_item = None # type: QPainterPath
self._block_item_obj = None # type: QGraphicsPathItem
self.qblock_annotations = None
self.addr_to_insns = {}
self.addr_to_labels = {}
self._init_widgets()
self._objects_are_hidden = False
self._create_block_item()
self.setAcceptHoverEvents(True)
#
# Properties
#
@property
def mode(self):
raise NotImplementedError
@property
def width(self):
return self.boundingRect().width()
@property
def height(self):
return self.boundingRect().height()
#
# Public methods
#
def clear_cache(self):
super().clear_cache()
for obj in self.objects:
obj.clear_cache()
def refresh(self):
for obj in self.objects:
obj.refresh()
self.layout_widgets()
self.recalculate_size()
self._create_block_item()
self.update()
def reload(self):
self._init_widgets()
self.refresh()
def size(self):
return self.width, self.height
def instruction_position(self, insn_addr):
if insn_addr in self.addr_to_insns:
insn = self.addr_to_insns[insn_addr]
pos = insn.pos()
return pos.x(), pos.y()
return None
#
# Initialization
#
def _create_block_item(self):
"""
Create the block background and border.
"""
if self._block_item_obj is not None and self._scene is not None:
self._scene.removeItem(self._block_item_obj)
self._block_item = None
self._block_item_obj = None
self._block_item = QPainterPath()
self._block_item.addRect(
self.block_object_group.childrenBoundingRect().marginsAdded(
self.margins))
def _init_widgets(self):
# TODO: hci: refactor: no need for self.objects anymore b/c of self.block_object_group. Using a
# QGraphicsItemGroup is a more natural way to group/work with multiple GraphicItems
if self._scene is not None:
for obj in self.objects:
self._scene.removeItem(obj)
self.objects.clear()
block_objects = get_block_objects(self.disasm, self.cfg_nodes,
self.func_addr)
self.block_object_group = QGraphicsItemGroup(parent=self)
self.block_object_group.setHandlesChildEvents(False)
for obj in block_objects:
if isinstance(obj, Instruction):
out_branch = get_out_branches_for_insn(self.out_branches,
obj.addr)
insn = QInstruction(self.workspace,
self.func_addr,
self.disasm_view,
self.disasm,
self.infodock,
obj,
out_branch,
self._config,
parent=self,
container=self._container)
self.objects.append(insn)
self.block_object_group.addToGroup(insn)
self.addr_to_insns[obj.addr] = insn
elif isinstance(obj, Label):
# label
label = QBlockLabel(obj.addr,
obj.text,
self._config,
self.disasm_view,
self.workspace,
self.infodock,
parent=self,
container=self._container)
self.objects.append(label)
self.block_object_group.addToGroup(label)
self.addr_to_labels[obj.addr] = label
elif isinstance(obj, PhiVariable):
if not isinstance(obj.variable, SimRegisterVariable):
phivariable = QPhiVariable(self.workspace,
self.disasm_view,
obj,
self._config,
parent=self,
container=self._container)
self.objects.append(phivariable)
self.block_object_group.addToGroup(phivariable)
elif isinstance(obj, Variables):
for var in obj.variables:
variable = QVariable(self.workspace,
self.disasm_view,
var,
self._config,
parent=self,
container=self._container)
self.objects.append(variable)
self.block_object_group.addToGroup(variable)
elif isinstance(obj, FunctionHeader):
header = QFunctionHeader(self.func_addr,
obj.name,
obj.prototype,
obj.args,
self._config,
self.disasm_view,
self.workspace,
self.infodock,
parent=self,
container=self._container)
self.objects.append(header)
self.block_object_group.addToGroup(header)
self.layout_widgets()
def layout_widgets(self):
raise NotImplementedError()
class QTraceViewer(QWidget):
TAG_SPACING = 50
LEGEND_X = -50
LEGEND_Y = 0
LEGEND_WIDTH = 10
TRACE_FUNC_X = 0
TRACE_FUNC_Y = 0
TRACE_FUNC_WIDTH = 50
TRACE_FUNC_MINHEIGHT = 1000
MARK_X = LEGEND_X
MARK_WIDTH = TRACE_FUNC_X - LEGEND_X + TRACE_FUNC_WIDTH
MARK_HEIGHT = 1
def __init__(self, workspace, disasm_view, parent=None):
super().__init__(parent=parent)
self.workspace = workspace
self.disasm_view = disasm_view
self.view = None
self.scene = None
self.mark = None
self.curr_position = 0
self._use_precise_position = False
self._init_widgets()
self.trace.am_subscribe(self._on_set_trace)
self.selected_ins.am_subscribe(self._on_select_ins)
self.view.installEventFilter(self)
#
# Forwarding properties
#
@property
def trace(self):
return self.workspace.instance.trace
@property
def selected_ins(self):
return self.disasm_view.infodock.selected_insns
def _init_widgets(self):
self.view = QGraphicsView()
self.scene = QGraphicsScene()
self.view.setScene(self.scene)
self._reset()
layout = QHBoxLayout()
layout.addWidget(self.view)
layout.setContentsMargins(0, 0, 0, 0)
layout.setAlignment(self.view, Qt.AlignLeft)
self.setLayout(layout)
def _reset(self):
self.scene.clear() #clear items
self.mark = None
self.legend = None
self.legend_height = 0
self.trace_func = QGraphicsItemGroup()
self.scene.addItem(self.trace_func)
self.hide()
def _on_set_trace(self, **kwargs):
self._reset()
if self.trace.am_obj is not None:
l.debug('minheight: %d, count: %d', self.TRACE_FUNC_MINHEIGHT,
self.trace.count)
if self.trace.count <= 0:
l.warning(
"No valid addresses found in trace to show. Check base address offsets?"
)
self.trace.am_obj = None
self.trace.am_event()
return
if self.TRACE_FUNC_MINHEIGHT < self.trace.count * 15:
self.trace_func_unit_height = 15
show_func_tag = True
else:
self.trace_func_unit_height = self.TRACE_FUNC_MINHEIGHT / self.trace.count
show_func_tag = True
self.legend_height = int(self.trace.count *
self.trace_func_unit_height)
self._show_trace_func(show_func_tag)
self._show_legend()
self._set_mark_color()
self.scene.setSceneRect(self.scene.itemsBoundingRect()) #resize
self.setFixedWidth(self.scene.itemsBoundingRect().width())
self.view.setFixedWidth(self.scene.itemsBoundingRect().width())
self.show()
def _on_select_ins(self, **kwargs):
if self.trace == None:
return
if self.mark is not None:
for i in self.mark.childItems():
self.mark.removeFromGroup(i)
self.scene.removeItem(i)
self.scene.removeItem(self.mark)
self.mark = QGraphicsItemGroup()
self.scene.addItem(self.mark)
if self.selected_ins:
addr = next(iter(self.selected_ins))
positions = self.trace.get_positions(addr)
if positions: #if addr is in list of positions
if not self._use_precise_position: #handle case where insn was selected from disas view
self.curr_position = positions[0] - self.trace.count
for p in positions:
color = self._get_mark_color(p, self.trace.count)
y = self._get_mark_y(p, self.trace.count)
if p == self.trace.count + self.curr_position: #add thicker line for 'current' mark
self.mark.addToGroup(
self.scene.addRect(self.MARK_X, y, self.MARK_WIDTH,
self.MARK_HEIGHT * 4,
QPen(QColor('black')),
QBrush(color)))
else:
self.mark.addToGroup(
self.scene.addRect(self.MARK_X, y, self.MARK_WIDTH,
self.MARK_HEIGHT, QPen(color),
QBrush(color)))
#y = self._get_mark_y(positions[0], self.trace.count)
#self.view.verticalScrollBar().setValue(y - 0.5 * self.view.size().height())
self.scene.update() #force redraw of the scene
self.scroll_to_position(self.curr_position)
def scroll_to_position(self, position):
relative_pos = self.trace.count + position
y_offset = self._get_mark_y(relative_pos, self.trace.count)
scrollValue = 0
if y_offset > 0.5 * self.view.size().height():
scrollValue = y_offset - 0.5 * self.view.size().height()
scrollValue = min(scrollValue, self.view.verticalScrollBar().maximum())
self.view.verticalScrollBar().setValue(scrollValue)
self._use_precise_position = False
def jump_next_insn(self):
if self.curr_position + self.trace.count < self.trace.count - 1: #for some reason indexing is done backwards
self.curr_position += 1
self._use_precise_position = True
func_name = self.trace.trace_func[self.curr_position].func_name
func = self._get_func_from_func_name(func_name)
bbl_addr = self.trace.trace_func[self.curr_position].bbl_addr
self._jump_bbl(func, bbl_addr)
def jump_prev_insn(self):
if self.curr_position + self.trace.count > 0:
self.curr_position -= 1
self._use_precise_position = True
func_name = self.trace.trace_func[self.curr_position].func_name
func = self._get_func_from_func_name(func_name)
bbl_addr = self.trace.trace_func[self.curr_position].bbl_addr
self._jump_bbl(func, bbl_addr)
def eventFilter(self, object, event): #specifically to catch arrow keys
# more elegant solution to link w/ self.view's scroll bar keypressevent?
if event.type() == QEvent.Type.KeyPress:
if not (event.modifiers() & Qt.ShiftModifier): #shift + arrowkeys
return False
key = event.key()
if key == Qt.Key_Up or key == Qt.Key_Left:
self.jump_prev_insn()
elif key == Qt.Key_Down or key == Qt.Key_Right:
self.jump_next_insn()
return True
return False # pass through all other events
def mousePressEvent(self, event):
button = event.button()
pos = self._to_logical_pos(event.pos())
if button == Qt.LeftButton and self._at_legend(pos):
func = self._get_func_from_y(pos.y())
bbl_addr = self._get_bbl_from_y(pos.y())
self._use_precise_position = True
self.curr_position = self._get_position(pos.y())
self._jump_bbl(func, bbl_addr)
def _jump_bbl(self, func, bbl_addr):
all_insn_addrs = self.workspace.instance.project.factory.block(
bbl_addr).instruction_addrs
# TODO: replace this with am_events perhaps?
self.workspace.on_function_selected(func)
self.selected_ins.clear()
self.selected_ins.update(all_insn_addrs)
self.selected_ins.am_event()
# TODO: this ought to happen automatically as a result of the am_event
self.disasm_view.current_graph.show_instruction(bbl_addr)
def _get_mark_color(self, i, total):
relative_gradient_pos = i * 1000 // total
return self.legend_img.pixelColor(self.LEGEND_WIDTH // 2,
relative_gradient_pos)
def _get_mark_y(self, i, total):
return self.TRACE_FUNC_Y + self.trace_func_unit_height * i
def _show_trace_func(self, show_func_tag):
x = self.TRACE_FUNC_X
y = self.TRACE_FUNC_Y
prev_name = None
for position in self.trace.trace_func:
bbl_addr = position.bbl_addr
func_name = position.func_name
l.debug('Draw function %x, %s', bbl_addr, func_name)
color = self.trace.get_func_color(func_name)
self.trace_func.addToGroup(
self.scene.addRect(x, y, self.TRACE_FUNC_WIDTH,
self.trace_func_unit_height, QPen(color),
QBrush(color)))
if show_func_tag is True and func_name != prev_name:
tag = self.scene.addText(func_name,
QFont("Source Code Pro", 7))
tag.setPos(x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING,
y - tag.boundingRect().height() // 2)
self.trace_func.addToGroup(tag)
anchor = self.scene.addLine(
self.TRACE_FUNC_X + self.TRACE_FUNC_WIDTH, y,
x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING, y)
self.trace_func.addToGroup(anchor)
prev_name = func_name
y += self.trace_func_unit_height
def _make_legend_gradient(self, x1, y1, x2, y2):
gradient = QLinearGradient(x1, y1, x2, y2)
gradient.setColorAt(0.0, Qt.red)
gradient.setColorAt(0.4, Qt.yellow)
gradient.setColorAt(0.6, Qt.green)
gradient.setColorAt(0.8, Qt.blue)
gradient.setColorAt(1.0, Qt.darkBlue)
return gradient
def _show_legend(self):
pen = QPen(Qt.transparent)
gradient = self._make_legend_gradient(
self.LEGEND_X, self.LEGEND_Y, self.LEGEND_X,
self.LEGEND_Y + self.legend_height)
brush = QBrush(gradient)
self.legend = self.scene.addRect(self.LEGEND_X, self.LEGEND_Y,
self.LEGEND_WIDTH, self.legend_height,
pen, brush)
reference_gradient = self._make_legend_gradient(
0, 0, self.LEGEND_WIDTH, 1000)
base_img = QImage(self.LEGEND_WIDTH, 1000, QImage.Format.Format_ARGB32)
p = QPainter(base_img)
p.fillRect(base_img.rect(), reference_gradient)
self.legend_img = base_img #reference shade
def _set_mark_color(self):
for p in range(self.trace.count):
color = self._get_mark_color(p, self.trace.count)
self.trace.set_mark_color(p, color)
def _at_legend(self, pos):
x = pos.x()
y = pos.y()
if self.TRACE_FUNC_X + self.LEGEND_X < x < self.view.width() and \
self.TRACE_FUNC_Y < y < self.TRACE_FUNC_Y + self.legend_height:
return True
else:
return False
def _to_logical_pos(self, pos):
x_offset = self.view.horizontalScrollBar().value()
y_offset = self.view.verticalScrollBar().value()
return QPoint(pos.x() + x_offset, pos.y() + y_offset)
def _get_position(self, y):
y_relative = y - self.legend_height
return int(y_relative // self.trace_func_unit_height)
def _get_bbl_from_y(self, y):
position = self._get_position(y)
return self.trace.get_bbl_from_position(position)
def _get_func_from_func_name(self, func_name):
return self.workspace.instance.kb.functions.function(name=func_name)
def _get_func_from_y(self, y):
position = self._get_position(y)
func_name = self.trace.get_func_name_from_position(position)
return self._get_func_from_func_name(func_name)
class QTraceViewer(QWidget):
"""
Load a basic block trace through json and visualize it in the disassembly
Ref: https://github.com/angr/angr-management/pull/122
"""
TAG_SPACING = 50
LEGEND_X = -50
LEGEND_Y = 0
LEGEND_WIDTH = 10
TRACE_FUNC_X = 0
TRACE_FUNC_Y = 0
TRACE_FUNC_WIDTH = 50
TRACE_FUNC_MINHEIGHT = 1000
TAB_HEADER_SIZE = 40
MAX_WINDOW_SIZE = 500
MARK_X = LEGEND_X
MARK_WIDTH = TRACE_FUNC_X - LEGEND_X + TRACE_FUNC_WIDTH
MARK_HEIGHT = 1
def __init__(self, workspace, disasm_view, parent=None):
super().__init__(parent=parent)
self.workspace = workspace
self.disasm_view = disasm_view
self.mark = None
self.legend = None
self.legend_height = 0
self.legend_img = None
self.trace_func_unit_height = 0
self.trace_func = None
self.trace_id = None
self.view = None
self.traceView = None
self.traceTab = None
self.traceScene = None
self.multiView = None
self.listView = None
self.mark = None
self.curr_position = 0
self._use_precise_position = False
self._selected_traces = []
self._init_widgets()
self.trace.am_subscribe(self._on_set_trace)
self.selected_ins.am_subscribe(self._on_select_ins)
self.traceTab.installEventFilter(self)
#
# Forwarding properties
#
@property
def trace(self):
return self.workspace.instance.trace
@property
def multi_trace(self):
return self.workspace.instance.multi_trace
@property
def selected_ins(self):
return self.disasm_view.infodock.selected_insns
def _init_widgets(self):
self.view = QTabWidget() # QGraphicsView()
self.traceTab = QWidget()
tracelayout = QVBoxLayout()
self.traceView = QGraphicsView()
self.traceScene = QGraphicsScene()
self.traceView.setScene(self.traceScene)
self.listView = QTableWidget(0, 2) # row, col
self.listView.setHorizontalHeaderItem(0, QTableWidgetItem("Trace ID"))
self.listView.setHorizontalHeaderItem(1, QTableWidgetItem("Input ID"))
self.listView.setSelectionMode(QAbstractItemView.SingleSelection)
self.listView.setSelectionBehavior(QAbstractItemView.SelectRows)
# self.listView.horizontalHeader().setStretchLastSection(True)
# self.listView.horizontalHeader().setSectionResizeModel(0, QHeaderView.Stretch)
self.listView.cellClicked.connect(self._switch_current_trace)
self.traceSeedButton = QPushButton("View Input Seed")
self.traceSeedButton.clicked.connect(self._view_input_seed)
tracelayout.addWidget(self.traceView)
tracelayout.addWidget(self.listView)
tracelayout.addWidget(self.traceSeedButton)
self.traceTab.setLayout(tracelayout)
self.multiView = QWidget()
multiLayout = QVBoxLayout()
self.multiTraceList = QTableWidget(0, 2) # row, col
self.multiTraceList.setSelectionMode(QAbstractItemView.MultiSelection)
self.multiTraceList.setSelectionBehavior(QAbstractItemView.SelectRows)
self.multiTraceList.setHorizontalScrollMode(
self.multiTraceList.ScrollPerPixel)
self.multiTraceList.setHorizontalHeaderItem(
0, QTableWidgetItem("Trace ID"))
self.multiTraceList.setHorizontalHeaderItem(
1, QTableWidgetItem("Input ID"))
self.selectMultiTrace = QPushButton("Refresh Heatmap")
self.selectMultiTrace.clicked.connect(self._refresh_heatmap)
multiLayout.addWidget(self.multiTraceList)
multiLayout.addWidget(self.selectMultiTrace)
self.multiView.setLayout(multiLayout)
self.view.addTab(self.traceTab, "SingleTrace")
self.view.addTab(self.multiView, "MultiTrace HeatMap")
self.SINGLE_TRACE = 0
self.MULTI_TRACE = 1
self.view.currentChanged.connect(self._on_tab_change)
self._reset()
layout = QVBoxLayout()
layout.addWidget(self.view)
layout.setContentsMargins(0, 0, 0, 0)
layout.setAlignment(self.view, Qt.AlignLeft)
self.setLayout(layout)
def _reset(self):
self.traceScene.clear() #clear items
self.listView.clearContents()
self.multiTraceList.clearContents()
self.mark = None
self.legend = None
self.legend_height = 0
self.trace_func = QGraphicsItemGroup()
self.trace_id = QGraphicsItemGroup()
self.traceScene.addItem(self.trace_func)
self.hide()
def _view_input_seed(self):
current_trace_stats = self.trace.am_obj
input_id = current_trace_stats.input_id
inputSeed = self.multi_trace.am_obj.get_input_seed_for_id(input_id)
msgText = "%s" % inputSeed
msgDetails = "Input for [%s]" % current_trace_stats.id
msgbox = QMessageBox()
msgbox.setWindowTitle("Seed Input")
msgbox.setDetailedText(msgDetails)
msgbox.setText(msgText)
msgbox.setStandardButtons(QMessageBox.Ok)
msgbox.exec()
def _switch_current_trace(self, row):
if self.listView.rowCount() <= 0:
return
current_trace = self.trace.am_obj.id
new_trace = self.multiTraceList.item(row, 0).text()
if current_trace == new_trace:
return
trace_stats = self.multi_trace.am_obj.get_trace_with_id(new_trace)
if trace_stats:
self.trace.am_obj = trace_stats
self._on_set_trace()
def _on_set_trace(self):
self._reset()
if self.trace.am_none or self.trace.count is None:
return
l.debug('minheight: %d, count: %d', self.TRACE_FUNC_MINHEIGHT,
self.trace.count)
if self.trace.count <= 0:
l.warning(
"No valid addresses found in trace to show. Check base address offsets?"
)
self.trace.am_obj = None
self.trace.am_event()
return
if self.TRACE_FUNC_MINHEIGHT < self.trace.count * 15:
self.trace_func_unit_height = 15
show_func_tag = True
else:
self.trace_func_unit_height = self.TRACE_FUNC_MINHEIGHT / self.trace.count
show_func_tag = True
self.legend_height = int(self.trace.count *
self.trace_func_unit_height)
self._show_trace_func(show_func_tag)
self._show_legend()
self._show_trace_ids()
self._set_mark_color()
self._refresh_multi_list()
boundingSize = self.traceScene.itemsBoundingRect().width()
windowSize = boundingSize
if boundingSize > self.MAX_WINDOW_SIZE:
windowSize = self.MAX_WINDOW_SIZE
self.traceScene.setSceneRect(
self.traceScene.itemsBoundingRect()) #resize
self.setFixedWidth(windowSize)
# self.listScene.setSceneRect(self.listScene.itemsBoundingRect()) #resize
self.multiView.setFixedWidth(windowSize)
cellWidth = windowSize // 2
self.listView.setColumnWidth(0, cellWidth)
self.listView.setColumnWidth(1, cellWidth)
self.listView.setFixedHeight(self.multiView.height() // 4)
self.multiTraceList.setColumnWidth(0, cellWidth)
self.multiTraceList.setColumnWidth(1, cellWidth)
self.view.setFixedWidth(windowSize)
self.show()
def _populate_trace_table(self, view, trace_ids):
numIDs = len(trace_ids)
view.clearContents()
view.setRowCount(numIDs)
row = 0 #start after label row
for traceID in trace_ids:
inputID = self.multi_trace.am_obj.get_input_id_for_trace_id(
traceID)
if inputID is None:
self.workspace.log("No inputID found for trace %s" % traceID)
view.setItem(row, 0, QTableWidgetItem(traceID))
view.setItem(row, 1, QTableWidgetItem(inputID))
row += 1
def _refresh_heatmap(self):
multiTrace = self.multi_trace.am_obj
multiTrace.clear_heatmap()
multiTrace.is_active_tab = True
selected_items = self.multiTraceList.selectedItems()
self._selected_traces.clear()
for row in range(self.multiTraceList.rowCount()):
item = self.multiTraceList.item(row, 0)
if item in selected_items:
self._selected_traces.append(item.text())
multiTrace.reload_heatmap(self._selected_traces)
self.multi_trace.am_event()
def _refresh_multi_list(self):
multiTrace = self.multi_trace.am_obj
trace_ids = multiTrace.get_all_trace_ids()
self.multiTraceList.clearContents()
self._populate_trace_table(self.multiTraceList, trace_ids)
if self._selected_traces and self.multiTraceList.rowCount() > 0:
self.multiTraceList.item(0, 0).setSelected(True)
self.multiTraceList.item(0, 1).setSelected(True)
else:
for row in range(self.multiTraceList.rowCount()):
item = self.multiTraceList.item(row, 0)
inputItem = self.multiTraceList.item(row, 1)
if item.text() in self._selected_traces:
item.setSelected(True)
inputItem.setSelected(True)
self.multi_trace.am_event()
def _on_tab_change(self):
# self._reset()
multiTrace = self.multi_trace.am_obj
if self.view.currentIndex() == self.MULTI_TRACE:
multiTrace.is_active_tab = True
self._refresh_multi_list()
elif self.view.currentIndex() == self.SINGLE_TRACE:
multiTrace = self.multi_trace.am_obj
multiTrace.is_active_tab = False
self._show_trace_ids()
def _on_select_ins(self, **kwargs): # pylint: disable=unused-argument
if self.trace.am_none:
return
if self.mark is not None:
for i in self.mark.childItems():
self.mark.removeFromGroup(i)
self.traceScene.removeItem(i)
self.traceScene.removeItem(self.mark)
self.mark = QGraphicsItemGroup()
self.traceScene.addItem(self.mark)
if self.selected_ins:
addr = next(iter(self.selected_ins))
positions = self.trace.get_positions(addr)
if positions: #if addr is in list of positions
if not self._use_precise_position: #handle case where insn was selected from disas view
self.curr_position = positions[0] - self.trace.count
for p in positions:
color = self._get_mark_color(p, self.trace.count)
y = self._get_mark_y(p)
if p == self.trace.count + self.curr_position: #add thicker line for 'current' mark
self.mark.addToGroup(
self.traceScene.addRect(self.MARK_X, y,
self.MARK_WIDTH,
self.MARK_HEIGHT * 4,
QPen(QColor('black')),
QBrush(color)))
else:
self.mark.addToGroup(
self.traceScene.addRect(self.MARK_X, y,
self.MARK_WIDTH,
self.MARK_HEIGHT,
QPen(color),
QBrush(color)))
self.traceScene.update() #force redraw of the traceScene
self.scroll_to_position(self.curr_position)
def scroll_to_position(self, position):
relative_pos = self.trace.count + position
y_offset = self._get_mark_y(relative_pos)
scrollValue = 0
if y_offset > 0.5 * self.traceView.size().height():
scrollValue = y_offset - 0.5 * self.traceView.size().height()
scrollValue = min(scrollValue,
self.traceView.verticalScrollBar().maximum())
self.traceView.verticalScrollBar().setValue(scrollValue)
self._use_precise_position = False
def jump_next_insn(self):
if self.curr_position + self.trace.count < self.trace.count - 1: #for some reason indexing is done backwards
self.curr_position += 1
self._use_precise_position = True
bbl_addr = self.trace.get_bbl_from_position(self.curr_position)
func = self.trace.get_func_from_position(self.curr_position)
self._jump_bbl(func, bbl_addr)
def jump_prev_insn(self):
if self.curr_position + self.trace.count > 0:
self.curr_position -= 1
self._use_precise_position = True
bbl_addr = self.trace.get_bbl_from_position(self.curr_position)
func = self.trace.get_func_from_position(self.curr_position)
self._jump_bbl(func, bbl_addr)
def eventFilter(self, obj, event): #specifically to catch arrow keys #pylint: disable=unused-argument
# more elegant solution to link w/ self.view's scroll bar keypressevent?
if event.type() == QEvent.Type.KeyPress:
if not event.modifiers() & Qt.ShiftModifier: #shift + arrowkeys
return False
key = event.key()
if key in [Qt.Key_Up, Qt.Key_Left]:
self.jump_prev_insn()
elif key in [Qt.Key_Down, Qt.Key_Right]:
self.jump_next_insn()
return True
return False # pass through all other events
def mousePressEvent(self, event):
button = event.button()
pos = self._to_logical_pos(event.pos())
if button == Qt.LeftButton and self.view.currentIndex(
) == self.SINGLE_TRACE and self._at_legend(pos):
func = self._get_func_from_y(pos.y())
bbl_addr = self._get_bbl_from_y(pos.y())
self._use_precise_position = True
self.curr_position = self._get_position(pos.y())
self._jump_bbl(func, bbl_addr)
def _jump_bbl(self, func, bbl_addr):
all_insn_addrs = self.workspace.instance.project.factory.block(
bbl_addr).instruction_addrs
# TODO: replace this with am_events perhaps?
if func is None:
return
self.workspace.on_function_selected(func)
self.selected_ins.clear()
self.selected_ins.update(all_insn_addrs)
self.selected_ins.am_event()
# TODO: this ought to happen automatically as a result of the am_event
self.disasm_view.current_graph.show_instruction(bbl_addr)
def _get_mark_color(self, i, total):
relative_gradient_pos = i * 1000 // total
return self.legend_img.pixelColor(self.LEGEND_WIDTH // 2,
relative_gradient_pos)
def _get_mark_y(self, i):
return self.TRACE_FUNC_Y + self.trace_func_unit_height * i
def _show_trace_ids(self):
trace_ids = self.multi_trace.get_all_trace_ids()
# traceID = self.listScene.addText(id_txt, QFont("Source Code Pro", 7))
# traceID.setPos(5,5)
self.listView.clearContents()
self._populate_trace_table(self.listView, trace_ids)
if len(self.listView.selectedItems()) <= 0 and not self.trace.am_none:
for row in range(self.listView.rowCount()):
item = self.listView.item(row, 0)
inputItem = self.listView.item(row, 1)
if self.trace.id in item.text():
item.setSelected(True)
inputItem.setSelected(True)
break
def _show_trace_func(self, show_func_tag):
x = self.TRACE_FUNC_X
y = self.TRACE_FUNC_Y
prev_name = None
for position in self.trace.trace_func:
bbl_addr = position.bbl_addr
func_name = position.func_name
l.debug('Draw function %x, %s', bbl_addr, func_name)
color = self.trace.get_func_color(func_name)
self.trace_func.addToGroup(
self.traceScene.addRect(x, y, self.TRACE_FUNC_WIDTH,
self.trace_func_unit_height,
QPen(color), QBrush(color)))
if show_func_tag is True and func_name != prev_name:
tag = self.traceScene.addText(func_name,
QFont("Source Code Pro", 7))
tag.setPos(x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING,
y - tag.boundingRect().height() // 2)
self.trace_func.addToGroup(tag)
anchor = self.traceScene.addLine(
self.TRACE_FUNC_X + self.TRACE_FUNC_WIDTH, y,
x + self.TRACE_FUNC_WIDTH + self.TAG_SPACING, y)
self.trace_func.addToGroup(anchor)
prev_name = func_name
y += self.trace_func_unit_height
@staticmethod
def _make_legend_gradient(x1, y1, x2, y2):
gradient = QLinearGradient(x1, y1, x2, y2)
gradient.setColorAt(0.0, Qt.red)
gradient.setColorAt(0.4, Qt.yellow)
gradient.setColorAt(0.6, Qt.green)
gradient.setColorAt(0.8, Qt.blue)
gradient.setColorAt(1.0, Qt.darkBlue)
return gradient
def _show_legend(self):
pen = QPen(Qt.transparent)
gradient = self._make_legend_gradient(
self.LEGEND_X, self.LEGEND_Y, self.LEGEND_X,
self.LEGEND_Y + self.legend_height)
brush = QBrush(gradient)
self.legend = self.traceScene.addRect(self.LEGEND_X, self.LEGEND_Y,
self.LEGEND_WIDTH,
self.legend_height, pen, brush)
reference_gradient = self._make_legend_gradient(
0, 0, self.LEGEND_WIDTH, 1000)
base_img = QImage(self.LEGEND_WIDTH, 1000, QImage.Format.Format_ARGB32)
p = QPainter(base_img)
p.fillRect(base_img.rect(), reference_gradient)
self.legend_img = base_img #reference shade
def _set_mark_color(self):
for p in range(self.trace.count):
color = self._get_mark_color(p, self.trace.count)
self.trace.set_mark_color(p, color)
def _at_legend(self, pos):
x = pos.x()
y = pos.y()
return self.TRACE_FUNC_X + self.LEGEND_X < x < self.traceView.width() and \
self.TRACE_FUNC_Y < y < self.TRACE_FUNC_Y + self.legend_height
def _to_logical_pos(self, pos):
x_offset = self.traceView.horizontalScrollBar().value()
y_offset = self.traceView.verticalScrollBar().value()
return QPoint(pos.x() + x_offset, pos.y() + y_offset)
def _get_position(self, y):
y_relative = y - self.legend_height - self.TAB_HEADER_SIZE
return int(y_relative // self.trace_func_unit_height)
def _get_bbl_from_y(self, y):
position = self._get_position(y)
return self.trace.get_bbl_from_position(position)
def _get_func_from_y(self, y):
position = self._get_position(y)
func = self.trace.get_func_from_position(position)
return func
class SectionScene(QGraphicsScene):
def __init__(self, section, model, commodity_types, process_cores):
super().__init__()
self._grid_size = QSize(GRID_WIDTH, GRID_HEIGHT)
self._drop_indicator = self.init_drop_indicator()
self._process_items = QGraphicsItemGroup()
self._commodity_items = List([])
self._bounding_rect = BoundingRect(section, model.process_list,
commodity_types)
self._clicked_item = None
self._item_mouse_offset = None
self._connect_line = None
self._items_border = QRect()
self._grid_border = QRect()
self._section = section
self._model = model
self._cores = process_cores
self._edit_mode = SelectConnect.SELECT
self._draft_mode = False
self.init_scene()
@property
def edit_mode(self):
return self._edit_mode
@edit_mode.setter
def edit_mode(self, value):
if value in SelectConnect:
self._edit_mode = value
if value == SelectConnect.SELECT:
self._process_items.setFlag(QGraphicsItem.ItemIsMovable)
else:
self._process_items.setFlag(QGraphicsItem.ItemIsMovable, False)
else:
raise TypeError
@property
def draft_mode(self):
return self._draft_mode
@draft_mode.setter
def draft_mode(self, value):
if value in [True, False]:
self._draft_mode = value
self._process_items.setFlag(QGraphicsItem.ItemIsMovable, value)
else:
raise TypeError
def init_scene(self):
"""initialize content of scene based on list of elements"""
for process in self._model.process_list:
if process.core.section == self._section:
process_item = self.draw_process(process)
self.draw_connection(process_item)
self.draw_connection(process_item, False)
self.update_commodities()
def draw_commodity(self, commodity):
"""create commodity item"""
commodity_item = CommodityItem(commodity.locations[self._section],
self._bounding_rect)
commodity_item.setData(0, commodity)
commodity_item.setData(1, 0)
self.addItem(commodity_item)
return commodity_item
def draw_connection(self, process_item, input_com=True):
commodity_types = []
process = process_item.data(0)
connections = process_item.data(1)
# define necessary connections to commodity item
commodities = process.inputs if input_com else process.outputs
for commodity in commodities:
if commodity.commodity_type not in commodity_types:
commodity_types.append(commodity.commodity_type)
commodity_difference = PROCESS_SIZE / (len(commodity_types) + 1)
for commodity_type in commodity_types:
commodity_position = commodity_type.locations[self._section]
item_position = process.coordinate.x() + PROCESS_SIZE / 2
# set item_position to left/right border of process depending on commodity_position
if commodity_position < item_position:
item_position -= PROCESS_SIZE
start_position = commodity_position if input_com else item_position
end_position = item_position if input_com else commodity_position
commodity_item = [
item for item in self._commodity_items
if item.data(0) is commodity_type
]
connection_item = [
connection for connection in connections
if connection.data(1) is commodity_item
]
if commodity_item:
commodity_item = commodity_item[0]
else:
# create new commodity item
commodity_item = self.draw_commodity(commodity_type)
self._commodity_items.add(commodity_item)
if connection_item:
connection_item = connection_item[0]
else:
# ignore second double arrow for storage processes if start_position is left of end_position
if (process.core.category is ProcessCategory.STORAGE) & (
end_position - start_position < 0):
return
# create new connection item
connection_item = ConnectionItem(
end_position - start_position,
process.core.category is ProcessCategory.STORAGE)
commodity_item.setData(1, commodity_item.data(1) + 1)
connection_item.setData(1, commodity_item)
connection_list = process_item.data(1)
connection_list.append(connection_item)
process_item.setData(1, connection_list)
connection_item.setX(start_position)
index = commodity_types.index(commodity_type) + 1
connection_item.setY(process.coordinate.y() - PROCESS_SIZE / 2 +
commodity_difference * index)
self.addItem(connection_item)
def draw_process(self, process):
"""create process item and add to scene & process items list"""
process_item = ProcessItem(process.core.icon)
process_item.setData(0, process)
process_item.setData(1, [])
process_item.setPos(process.coordinate)
process_item.setFlag(QGraphicsItem.ItemIsMovable)
self._process_items.prepareGeometryChange()
self._process_items.addToGroup(process_item)
self.addItem(process_item)
return process_item
def delete_process(self, item):
process = item.data(0)
self._model.process_list.remove(process)
# remove all connection items and the commodity item if applicable
for connection_item in item.data(1):
commodity_item = connection_item.data(1)
commodity_references = commodity_item.data(1) - 1
commodity_item.setData(1, commodity_references)
if commodity_references == 0:
# remove commodity of section count and commodity item
commodity_item.data(0).locations.remove(self._section)
self._commodity_items.remove(commodity_item)
self.removeItem(commodity_item)
self.removeItem(connection_item)
# reduce connection count of all commodities connected to process
for input_com in process.inputs:
input_com.connection_count[self._section] -= 1
if input_com.connection_count[self._section] == 0:
input_com.connection_count.remove(self._section)
if not input_com.connection_count:
self._model.commodity_list.remove(input_com)
for output in process.outputs:
output.connection_count[self._section] -= 1
if output.connection_count[self._section] == 0:
output.connection_count.remove(self._section)
if not output.connection_count:
self._model.commodity_list.remove(output)
self._process_items.removeFromGroup(item)
self.removeItem(item)
self.update_commodities()
def init_drop_indicator(self):
rect = QRect(-DROP_INDICATOR_SIZE / 2, -DROP_INDICATOR_SIZE / 2,
DROP_INDICATOR_SIZE, DROP_INDICATOR_SIZE)
brush = QBrush(Qt.darkBlue)
pen = QPen(Qt.darkBlue, 1)
ellipse = self.addEllipse(rect, pen, brush)
ellipse.setVisible(False)
return ellipse
def disable_drop_indicator(self):
self._drop_indicator.setX(-DROP_INDICATOR_SIZE / 2)
self._drop_indicator.setY(-DROP_INDICATOR_SIZE / 2)
self._drop_indicator.setVisible(False)
self.update()
def align_drop_indicator(self, point):
"""align the drop indicator to grid while mouse movement"""
# move drop indicator only within grid boundaries
if not self.get_grid_border(-MIN_GRID_MARGIN).contains(
point.toPoint()):
self._drop_indicator.setVisible(False)
return
# calculate the nearby position of grid interceptions
grid_x = (round(point.x() / self._grid_size.width() / 2 - 1 / 2) * 2 +
1) * self._grid_size.width()
grid_y = round(
point.y() / self._grid_size.height()) * self._grid_size.height()
# move drop indicator to grid interception
self._drop_indicator.setX(grid_x)
self._drop_indicator.setY(grid_y)
self._drop_indicator.setVisible(True)
self.update()
def get_grid_border(self, margin=0):
def get_raster_length(length, raster):
return (int((length + MIN_GRID_MARGIN + raster / 2) / raster) -
1 / 2) * raster
left = get_raster_length(self.sceneRect().left(),
self._grid_size.width())
right = get_raster_length(self.sceneRect().right(),
self._grid_size.width())
top = get_raster_length(self.sceneRect().top(),
self._grid_size.height())
bottom = get_raster_length(self.sceneRect().bottom(),
self._grid_size.height())
return QRect(left, top, right - left, bottom - top).marginsAdded(
QMargins(margin, margin, margin, margin))
def execute_connection(self, item, position):
if not self._connect_line:
if self.show_connection_menu(item, position):
return
else:
selected_commodity = self._connect_line.data(1)
if isinstance(item, ProcessItem):
# create incoming connection
if selected_commodity in item.data(0).core.inputs:
if self._connect_line.data(0) is None:
# establish connection from commodity
if selected_commodity not in item.data(0).inputs:
self.connect_process(item.data(0),
selected_commodity, False)
self.draw_connection(item)
else:
# establish connection between processes
self.connect_processes(selected_commodity,
self._connect_line.data(0),
item)
elif item is None:
# create outgoing connection
menu = QMenu()
menu.addAction("> Commodity")
if menu.exec_(QCursor.pos()):
self.connect_process(
self._connect_line.data(0).data(0), selected_commodity)
self.draw_connection(self._connect_line.data(0), False)
# end connection
self.views()[1].setMouseTracking(False)
self.removeItem(self._connect_line)
self._connect_line = None
def show_connection_menu(self, item, position):
menu = QMenu()
# add all output commodities of process to menu
for commodity in item.data(0).core.outputs:
action = menu.addAction(str(commodity))
action.setData(commodity)
# execute menu
action = menu.exec_(QCursor.pos())
if action:
# start connection
line_pen = QPen(Qt.lightGray, 1, Qt.DashLine)
line = QLineF(position, position)
self._connect_line = self.addLine(line, line_pen)
self._connect_line.setData(0, item)
self._connect_line.setData(1, action.data())
self.views()[1].setMouseTracking(True)
return True
return False
def connect_processes(self, selected_commodity, start_process_item,
end_process_item):
# establish connection
connect_commodity = None
commodity_in_input = False
commodity_in_output = False
start_process = start_process_item.data(0)
end_process = end_process_item.data(0)
# check if selected_commodity is in outputs or inputs of related processes
outputs = start_process.outputs
if selected_commodity in outputs:
commodity_in_output = True
connect_commodity = outputs[outputs.index(selected_commodity)]
inputs = end_process.inputs
if selected_commodity in inputs:
commodity_in_input = True
connect_commodity = inputs[inputs.index(selected_commodity)]
# determine necessary connection steps
if commodity_in_input & commodity_in_output:
# selected_commodity already exists in input and output -> connection exists
QMessageBox.warning(self.views()[1], "Connection exists",
"Connection already established",
QMessageBox.Ok)
return
elif commodity_in_input:
# add commodity to output of start process
self.connect_process(start_process, connect_commodity)
self.draw_connection(start_process_item, False)
elif commodity_in_output:
# add commodity to input of end process
self.connect_process(end_process, connect_commodity, False)
self.draw_connection(end_process_item)
else:
# establish new commodity in both processes
connect_commodity = selected_commodity.copy()
self.connect_process(start_process, connect_commodity)
self.connect_process(end_process, connect_commodity, False)
self.draw_connection(start_process_item, False)
self.draw_connection(end_process_item)
def connect_process(self, process, commodity, output=True):
if self._section not in commodity.connection_count.keys():
commodity.connection_count[self._section] = 0
if commodity not in self._model.commodity_list:
self._model.commodity_list.add(commodity)
commodity.connection_count[self._section] += 1
commodity_list = process.outputs if output else process.inputs
commodity_list.add(commodity)
# add commodity in input and output list
if process.core.category is ProcessCategory.STORAGE:
commodity_list = process.outputs if not output else process.inputs
commodity_list.add(commodity)
self.set_commodity_position(commodity.commodity_type, process, output)
def set_commodity_position(self, commodity_type, process, right=True):
"""changes the commodity position according to newly connected process"""
if self._section not in commodity_type.locations.keys():
# set initial position to right of process
location = process.coordinate.x() + self._grid_size.width() * (
1 if right else -1)
commodity_type.locations[self._section] = location
else:
# todo change position based on newly connected process
pass
def update_commodities(self):
"""update length of commodity line based on bounding rect"""
self._bounding_rect.update()
top_border = self._bounding_rect.top()
height_border = self._bounding_rect.height()
for commodity_item in self._commodity_items:
commodity_item.update_line(top_border, height_border)
def drawBackground(self, painter, rect):
if self.draft_mode:
border_rect = self.get_grid_border()
grid_rect = self.get_grid_border(-self._grid_size.width() / 2)
line_list = []
# horizontal grid lines
for line_coordinate in range(grid_rect.top(), border_rect.bottom(),
self._grid_size.height()):
line_list.append(
QLineF(border_rect.left(), line_coordinate,
border_rect.right(), line_coordinate))
# vertical process lines
left_border = int((grid_rect.left() + self._grid_size.width()) / (2 * self._grid_size.width())) * \
self._grid_size.width() * 2 - self._grid_size.width()
for line_coordinate in range(left_border, border_rect.right(),
self._grid_size.width() * 2):
line_list.append(
QLineF(line_coordinate, border_rect.top(), line_coordinate,
border_rect.bottom()))
# vertical commodity lines
left_border = int(
grid_rect.left() /
(2 * self._grid_size.width())) * self._grid_size.width() * 2
for line_coordinate in range(left_border, border_rect.right(),
self._grid_size.width() * 2):
line_list.append(
QLineF(line_coordinate - 1, border_rect.top(),
line_coordinate - 1, border_rect.bottom()))
line_list.append(
QLineF(line_coordinate + 1, border_rect.top(),
line_coordinate + 1, border_rect.bottom()))
painter.setPen(QPen(Qt.lightGray, 1))
painter.drawLines(line_list)
super().drawBackground(painter, rect)
def mousePressEvent(self, event):
# ignore right button click
if event.button() == Qt.RightButton:
return
self._clicked_item = self.itemAt(event.scenePos(), QTransform())
if self.draft_mode & (self.edit_mode == SelectConnect.SELECT):
if isinstance(self._clicked_item, ProcessItem):
self._clicked_item.setOpacity(0.5)
self._drop_indicator.setPos(self._clicked_item.pos())
self._drop_indicator.setVisible(True)
# set offset between mouse and center of item for mouseMoveEvent
self._item_mouse_offset = self._clicked_item.mapFromScene(
event.scenePos())
else:
self._clicked_item = None
def mouseMoveEvent(self, event):
if self.draft_mode:
if self.edit_mode == SelectConnect.SELECT:
if self._clicked_item:
# move drop indicator along grid and process item according to mouse position
self.align_drop_indicator(event.scenePos() -
self._item_mouse_offset)
self._clicked_item.setPos(event.scenePos() -
self._item_mouse_offset)
elif self._connect_line:
# update connect line in CONNECT mode
line = self._connect_line.line()
# add offset to position to avoid itemAt function to return lineItem
line.setP2(event.scenePos() - QPoint(2, 2))
self._connect_line.setLine(line)
super().mouseMoveEvent(event)
def mouseReleaseEvent(self, event):
# normal mode
if not self.draft_mode:
if event.button() == Qt.LeftButton:
# toggle sidebar with properties
clicked_item = self.itemAt(event.scenePos(), QTransform())
if self._clicked_item:
if isinstance(clicked_item, ProcessItem):
self.views()[0].sidebar_toggled.emit(
clicked_item.data(0))
self.views()[0].commodity_clicked.emit(None)
elif isinstance(clicked_item, CommodityItem):
self.views()[0].sidebar_toggled.emit(None)
self.views()[0].commodity_clicked.emit(
clicked_item.data(0))
else:
self.views()[0].sidebar_toggled.emit(None)
self.views()[0].commodity_clicked.emit(None)
else:
self.views()[0].sidebar_toggled.emit(None)
self.views()[0].commodity_clicked.emit(None)
return
# draft mode in selection
if self.edit_mode == SelectConnect.SELECT:
# release process item
if self._clicked_item:
self._clicked_item.setOpacity(1.0)
# remove clicked item from collision list with drop indicator
collision_items = self.collidingItems(self._drop_indicator)
if collision_items:
collision_items.remove(self._clicked_item)
# avoid placement if colliding with other items
if not collision_items:
self._clicked_item.setPos(self._drop_indicator.scenePos())
self._clicked_item.data(
0).coordinate = self._drop_indicator.pos()
self._bounding_rect.update()
self.disable_drop_indicator()
self._clicked_item = None
else:
# connecting processes
if event.button() == Qt.LeftButton:
clicked_item = self.itemAt(event.scenePos(), QTransform())
self.execute_connection(clicked_item, event.scenePos())
super().mouseReleaseEvent(event)
def dragEnterEvent(self, event):
pass
def dragLeaveEvent(self, event):
self.disable_drop_indicator()
def dragMoveEvent(self, event):
if self.draft_mode & (self.edit_mode == SelectConnect.SELECT):
self.align_drop_indicator(event.scenePos())
def dropEvent(self, event):
"""add new process to list and process item to scene"""
# prevent process placement if item exists there
if len(self.collidingItems(self._drop_indicator)) > 0:
return
core_name = event.mimeData().data(
MimeType.PROCESS_CORE.value).data().decode('UTF-8')
process_core = list(
filter(lambda element: element.name == core_name, self._cores))[0]
# define name based on process core and existing names
remainder_name = [
process.name.split(process_core.name)[1].strip()
for process in self._model.process_list
if process_core.name in process.name
]
if not remainder_name:
process_name = process_core.name
elif max(remainder_name).isdigit():
process_name = process_core.name + " " + str(
int(max(remainder_name)) + 1)
else:
process_name = process_core.name + " 1"
# create new process based on process core and add to element list
process = Process(process_name, self._drop_indicator.pos(),
process_core)
self._model.add_process(process)
# create process item and connections
self.draw_process(process)
self.update_commodities()
self.disable_drop_indicator()
def contextMenuEvent(self, event):
"""context menu to interact with process items - delete item"""
# prevent context menu of scene not in draft mode
if not self.draft_mode:
return
if self._edit_mode == SelectConnect.SELECT:
# open context menu only for process items
self._clicked_item = self.itemAt(event.scenePos(), QTransform())
if self._clicked_item:
if isinstance(self._clicked_item, ProcessItem):
menu = QMenu()
delete_action = QAction("Delete", None)
delete_action.triggered.connect(
lambda: self.delete_process(self._clicked_item))
menu.addAction(delete_action)
menu.exec_(event.screenPos())
self._clicked_item = None
else:
# open context menu for commodities of other sections
menu = QMenu()
section_list = [
item for item in OverviewSelection if item is not self._section
and item is not OverviewSelection.OVERVIEW
]
for section in section_list:
submenu = QMenu(section.name)
section_coms = [
commodity for commodity in self._model.commodity_list
if section in commodity.connection_count.keys()
]
for commodity in section_coms:
action = submenu.addAction(str(commodity))
action.setData(commodity)
# only add sub menu if commodities are available
if section_coms:
menu.addMenu(submenu)
# execute menu
action = menu.exec_(QCursor.pos())
if action:
# start connection
line_pen = QPen(Qt.lightGray, 1, Qt.DashLine)
line = QLineF(event.scenePos(), event.scenePos())
self._connect_line = self.addLine(line, line_pen)
self._connect_line.setData(0, None)
self._connect_line.setData(1, action.data())
self.views()[1].setMouseTracking(True)
def setSceneRect(self, rect):
"""set sceneRect to view boundaries if necessary space is less"""
super().setSceneRect(
self._bounding_rect.scene_rect(rect, 2 * MIN_GRID_MARGIN))