def draw_shortest_path(self, scene: QGraphicsScene, navmesh: NavMesh, destination: Point, player: bool) -> None: for line in self.shortest_path_segments: try: scene.removeItem(line) except RuntimeError: pass if player: origin = self.game.theater.player_points()[0] else: origin = self.game.theater.enemy_points()[0] prev_pos = self._transform_point(origin.position) try: path = navmesh.shortest_path(origin.position, destination) except ValueError: return for waypoint in path[1:]: new_pos = self._transform_point(waypoint) flight_path_pen = self.flight_path_pen(player, selected=True) # Draw the line to the *middle* of the waypoint. offset = self.WAYPOINT_SIZE // 2 self.shortest_path_segments.append( scene.addLine(prev_pos[0] + offset, prev_pos[1] + offset, new_pos[0] + offset, new_pos[1] + offset, flight_path_pen)) self.shortest_path_segments.append( scene.addEllipse(new_pos[0], new_pos[1], self.WAYPOINT_SIZE, self.WAYPOINT_SIZE, flight_path_pen, flight_path_pen)) prev_pos = new_pos
def draw_test_flight_plan(self, scene: QGraphicsScene, task: FlightType, point_near_target: Point, player: bool) -> None: for line in self.shortest_path_segments: try: scene.removeItem(line) except RuntimeError: pass self.clear_flight_paths(scene) target = self.game.theater.closest_target(point_near_target) if player: origin = self.game.theater.player_points()[0] else: origin = self.game.theater.enemy_points()[0] package = Package(target) flight = Flight(package, F_16C_50, 2, task, start_type="Warm", departure=origin, arrival=origin, divert=None) package.add_flight(flight) planner = FlightPlanBuilder(self.game, package, is_player=player) try: planner.populate_flight_plan(flight) except InvalidObjectiveLocation: return package.time_over_target = TotEstimator(package).earliest_tot() self.draw_flight_plan(scene, flight, selected=True)
def clear_flight_paths(self, scene: QGraphicsScene) -> None: for item in self.flight_path_items: try: scene.removeItem(item) except RuntimeError: # Something may have caused those items to already be removed. pass self.flight_path_items.clear()
def testIt(self): scene = QGraphicsScene() i1 = QGraphicsRectItem() scene.addItem(i1) i2 = QGraphicsRectItem(i1) i3 = QGraphicsRectItem() i4 = QGraphicsRectItem() group = scene.createItemGroup((i2, i3, i4)) scene.removeItem(i1) del i1 # this shouldn't delete i2 self.assertEqual(i2.scene(), scene) scene.destroyItemGroup(group) self.assertRaises(RuntimeError, group.type)
def testIt(self): scene = QGraphicsScene() i1 = QGraphicsRectItem() scene.addItem(i1) i2 = QGraphicsRectItem(i1) i3 = QGraphicsRectItem() i4 = QGraphicsRectItem() group = scene.createItemGroup((i2, i3, i4)) scene.removeItem(i1) del i1 # this shouldn't delete i2 self.assertEqual(i2.scene(), scene) scene.destroyItemGroup(group) self.assertRaises(RuntimeError, group.type)
def highlight_mouse_navmesh(self, scene: QGraphicsScene, navmesh: NavMesh, mouse_position: Point) -> None: if self.navmesh_highlight is not None: try: scene.removeItem(self.navmesh_highlight) except RuntimeError: pass navpoly = navmesh.localize(mouse_position) if navpoly is None: return self.navmesh_highlight = self.draw_shapely_poly( scene, navpoly.poly, CONST.COLORS["transparent"], CONST.COLORS["light_green_transparent"])
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 DesktopWideOverlay(QMainWindow): _instances: Dict[int, QGraphicsItem] view: QGraphicsView scene: QGraphicsScene def __init__(self): super().__init__(flags=Qt.Widget | Qt.FramelessWindowHint | Qt.BypassWindowManagerHint | Qt.WindowTransparentForInput | Qt.WindowStaysOnTopHint) self.setAttribute(Qt.WA_NoSystemBackground, True) self.setAttribute(Qt.WA_TranslucentBackground, True) self.setAttribute(Qt.WA_DeleteOnClose, True) self._instances = {} virtual_screen = QRect(0, 0, 0, 0) for screen in QGuiApplication.screens(): # TODO: Handle screen change geom = screen.virtualGeometry() virtual_screen = virtual_screen.united(geom) self.scene = QGraphicsScene(0, 0, virtual_screen.width(), virtual_screen.height(), parent=self) self.view = QGraphicsView(self.scene, self) self.view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setStyleSheet("background: transparent") self.view.setGeometry(0, 0, virtual_screen.width(), virtual_screen.height()) self.view.setInteractive(False) self.transparent_pen = QPen() self.transparent_pen.setBrush(Qt.NoBrush) self.setGeometry(virtual_screen) def add_instance( self, instance: "GameInstance" ) -> Tuple[QGraphicsItem, Callable[[], None]]: """Add instance to manage, return a disconnect function and the canvas""" positionChanged = lambda rect: self.on_instance_moved(instance, rect) instance.positionChanged.connect(positionChanged) focusChanged = lambda focus: self.on_instance_focus_change( instance, focus) instance.focusChanged.connect(focusChanged) instance_pos = instance.get_position() gfx = QGraphicsRectItem(rect=instance_pos) gfx.setPen(self.transparent_pen) gfx.setPos(instance_pos.x(), instance_pos.y()) self.scene.addItem(gfx) self._instances[instance.wid] = gfx def disconnect(): gfx.hide() self.scene.removeItem(gfx) instance.positionChanged.disconnect(positionChanged) instance.focusChanged.disconnect(focusChanged) return gfx, disconnect def on_instance_focus_change(self, instance, focus): # self._instances[instance.wid].setVisible(focus) pass def on_instance_moved(self, instance, pos: QRect): rect = self._instances[instance.wid] rect.setRect(0, 0, pos.width(), pos.height()) rect.setPos(pos.x(), pos.y())
class DesktopWideOverlay(QMainWindow): _instances: Dict[int, QGraphicsItem] view: QGraphicsView scene: QGraphicsScene def __init__(self): super().__init__(flags=Qt.Widget | Qt.FramelessWindowHint | Qt.BypassWindowManagerHint | Qt.WindowTransparentForInput | Qt.WindowStaysOnTopHint) self.logger = logging.getLogger(__name__ + "." + self.__class__.__name__) self.setAttribute(Qt.WA_NoSystemBackground, True) self.setAttribute(Qt.WA_TranslucentBackground, True) self.setAttribute(Qt.WA_DeleteOnClose, True) self.setStyleSheet("background: transparent") self._instances = {} virtual_screen = QRect(0, 0, 0, 0) for screen in QGuiApplication.screens(): # TODO: Handle screen change geom = screen.virtualGeometry() virtual_screen = virtual_screen.united(geom) self.scene = QGraphicsScene(0, 0, virtual_screen.width(), virtual_screen.height(), parent=self) self.view = QGraphicsView(self.scene, self) self.view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setStyleSheet("background: transparent") self.view.setGeometry(0, 0, virtual_screen.width(), virtual_screen.height()) self.view.setInteractive(False) self.transparent_pen = QPen() self.transparent_pen.setBrush(Qt.NoBrush) self.setGeometry(virtual_screen) def add_instance( self, instance: "GameInstance" ) -> Tuple[QGraphicsItem, Callable[[], None]]: """Add instance to manage, return a disconnect function and the canvas""" positionChanged = lambda rect: self.on_instance_moved(instance, rect) instance.positionChanged.connect(positionChanged) focusChanged = lambda focus: self.on_instance_focus_change( instance, focus) instance.focusChanged.connect(focusChanged) instance_pos = instance.get_position() gfx = QGraphicsRectItem(rect=instance_pos) gfx.setPen(self.transparent_pen) gfx.setPos(instance_pos.x(), instance_pos.y()) self.scene.addItem(gfx) self._instances[instance.wid] = gfx def disconnect(): gfx.hide() self.scene.removeItem(gfx) instance.positionChanged.disconnect(positionChanged) instance.focusChanged.disconnect(focusChanged) return gfx, disconnect def on_instance_focus_change(self, instance, focus): # self._instances[instance.wid].setVisible(focus) pass def on_instance_moved(self, instance, pos: QRect): rect = self._instances[instance.wid] rect.setRect(0, 0, pos.width(), pos.height()) rect.setPos(pos.x(), pos.y()) def check_compatibility(self): QTimer.singleShot(300, self._check_compatibility) def _check_compatibility(self): # If we cause black screen then hide ourself out of shame... screenshot = QGuiApplication.primaryScreen().grabWindow(0) image = qpixmap_to_np(screenshot) black_pixels = np.count_nonzero( np.all(image[:, :, :3] == [0, 0, 0], axis=-1)) total_pixels = len(image[:, :, 0].flatten()) black_percent = black_pixels / total_pixels self.logger.debug("Screen black ratio %.2f%%", black_percent * 100) if black_percent > 0.95: self.logger.warning( "Detected black screen at %.2f%%. Disabling overlay", black_percent * 100, ) self.hide()
class MyForm(QMainWindow): def __init__(self, width, height, pixel_ratio, path): super().__init__() self.ui = Ui_MainWindow() self.ui.setupUi(self) # overwrite dimesions specified in slideShow3.py, as they are specific to MacBookPro display, and QTDesigner # has its own idea on what they should be. This code should work on any size display self.resize(width, height) self.ui.graphicsView.setGeometry(QtCore.QRect(0, 0, width, height)) self.ui.menubar.setGeometry(QtCore.QRect(0, 0, width, 0)) self.width = width self.height = height self.pixel_ratio = pixel_ratio self.path = path self.imageFiles = [] self.slideIndex = -1 self.random_index_number = 0 self.random = "" self.imageFiles, self.random_index, self.path, self.max_index = self.getImageNames2() self.helpFile = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "instructions.png") #print(self.helpFile) self.scene = QGraphicsScene(self) #self.scene.setAlignment(QtCore.Qt.AlignCenter) self.ui.actionDir.triggered.connect(self.openFileNameDialog) self.ui.actionStart_Slide_Show.triggered.connect(self.slide_show) self.ui.actionRandom_Slide_Show.triggered.connect(self.random_slide_show) eventFilter = MouseEventFilter(self.scene) self.scene.installEventFilter(eventFilter) self.ui.actionHelp.triggered.connect(self.helpWindow) #self.show() self.showFullScreen() extension = staticmethod(lambda f: f.split('.').pop().lower()) filename = staticmethod(lambda f: f.split('/').pop()) def openFileNameDialog(self): options = QFileDialog.Options() options |= QFileDialog.DontUseNativeDialog self.fileName, _ = QFileDialog.getOpenFileName(self,"QFileDialog.getOpenFileName()", "","All Files (*);;Python Files (*.py)", options=options) if self.fileName: self.path = os.path.dirname(self.fileName) self.imageFiles = [] self.random_index = [] self.max_index = [] self.imageFiles, self.random_index, self.path, self.max_index = self.getImageNames2() self.slideIndex = self.imageFiles.index(self.fileName) -1 def getImageNames2(self): "get the names of all images on disc or from the web (which are cached locally)" if not self.path: self.path = os.getcwd() if self.path[-1] != '/': self.path += '/' try: os.listdir(self.path) except: error_dialog = QtWidgets.QErrorMessage() error_dialog.showMessage('Error in path' +self.path) # https://stackoverflow.com/questions/40227047/python-pyqt5-how-to-show-an-error-message-with-pyqt5 return [], self.path for i in GlobDirectoryWalker(self.path, "*.*"): if os.path.isfile(i): if self.checkImageType(i): self.imageFiles.append(i) max_index = len(self.imageFiles) - 1 self.imageFiles.sort() random_index = list(range(max_index + 1)) random.shuffle(random_index) return self.imageFiles, random_index, self.path, max_index def slide(self, i): self.pixmap = QtGui.QPixmap() #self.pixmap.setAlignment(QtCore.Qt.AlignCenter) self.pixmap.load(self.imageFiles[i]) self.pixmap.setDevicePixelRatio(self.pixel_ratio) # https://stackoverflow.com/questions/50127246/pyqt-5-10-enabling-high-dpi-support-for-macos-poor-pixmap-quality #self.pixmap4 = self.pixmap.scaled(self.width * self.pixel_ratio, (self.height * self.pixel_ratio)-45, Qt.KeepAspectRatio) self.pixmap4 = self.pixmap.scaled(self.width * self.pixel_ratio, (self.height * self.pixel_ratio), Qt.KeepAspectRatio) try: self.scene.removeItem(self.item) except: print("failed to remove item") self.item = QGraphicsPixmapItem(self.pixmap4) self.scene.addItem(self.item) #myapp.setWindowTitle(os.path.basename(self.imageFiles[i])) self.setWindowTitle(os.path.basename(self.imageFiles[i])) self.ui.graphicsView.setScene(self.scene) def slide_show(self): self.random = 0 self.next_slide() def random_slide_show(self): self.random = 1 self.next_slide() def next_slide(self): if self.random == 0: self.increment_slide() else: self.random_next() def prev_slide(self): if self.random == 0: self.decrement_slide() else: self.random_prev() def random_next(self): "display the next random slide" self.random_index_number += 1 try: self.slideIndex = self.random_index[self.random_index_number] self.slide(self.slideIndex) except IndexError: self.random_index_number = 0 self.slideIndex = self.random_index[self.random_index_number] self.slide(self.slideIndex) return False def random_prev(self): "display the previous random slide" self.random_index_number -= 1 #self.ImageWindow.clear() try: self.slideIndex = self.random_index[self.random_index_number] self.slide(self.slideIndex) except IndexError: self.random_index_number = self.max_index self.slideIndex = self.random_index[self.random_index_number] self.slide(self.slideIndex) return False def increment_slide(self): "display a higher slide" print("in increment_slide") self.slideIndex += 1 if self.slideIndex > self.max_index: self.slideIndex = 0 print('Max index hit') self.slide(self.slideIndex) return False def decrement_slide(self): "display a lower slide" self.slideIndex -= 1 if self.slideIndex < 0: self.slideIndex = self.max_index self.slide(self.slideIndex) return False def checkImageType(self, f): "check to see if we have an file with an image extension" ext = self.extension(f) chk = [i for i in ['jpg','gif','ppm', 'tif', 'png', 'jpeg'] if i==ext] if chk == []: return False return True def helpWindow(self): self.pixmap = QtGui.QPixmap() #self.pixmap.setAlignment(QtCore.Qt.AlignCenter) self.pixmap.load(self.helpFile) self.pixmap.setDevicePixelRatio(self.pixel_ratio) # https://stackoverflow.com/questions/50127246/pyqt-5-10-enabling-high-dpi-support-for-macos-poor-pixmap-quality #self.pixmap4 = self.pixmap.scaled(self.width * self.pixel_ratio, (self.height * self.pixel_ratio)-45, Qt.KeepAspectRatio) self.pixmap4 = self.pixmap.scaled(self.width * self.pixel_ratio, (self.height * self.pixel_ratio), Qt.KeepAspectRatio) try: self.scene.removeItem(self.item) except: print("failed to remove item") self.item = QGraphicsPixmapItem(self.pixmap4) self.scene.addItem(self.item) #myapp.setWindowTitle(os.path.basename("Instructions")) self.setWindowTitle(os.path.basename("Instructions")) self.ui.graphicsView.setScene(self.scene) def keyPressEvent(self, e): if e.key() == Qt.Key_Escape: self.Quit() if e.key() == Qt.Key_Q: self.Quit() if e.key() == Qt.Key_Space: self.next_slide() if e.key() == Qt.Key_N: self.random_next() if e.key() == Qt.Key_P: self.random_prev() if e.key() == Qt.Key_Comma: self.decrement_slide() if e.key() == Qt.Key_Period: self.increment_slide() if e.key() == Qt.Key_H: self.helpWindow = self.helpWindow() if e.key() == Qt.Key_BracketLeft: self.slideIndex = self.decrement_slide() def mousePressEvent(self, e): if e.button() == QtCore.Qt.LeftButton: print("trapped left mouse click") self.next_slide() if e.button() == QtCore.Qt.RightButton: print("trapped right mouse click") self.prev_slide() def Quit(self): sys.exit(app.exec_())
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 QGameOfLife(QWidget): Games = { "Game of Life": (GameOfLife, { 'fill_rate': 0.50 }), "Bacteria": (GrayScottDiffusion, { 'coeffs': (0.16, 0.08, 0.035, 0.065) }), "Coral": (GrayScottDiffusion, { 'coeffs': (0.16, 0.08, 0.062, 0.062) }), "Fingerprint": (GrayScottDiffusion, { 'coeffs': (0.19, 0.05, 0.060, 0.062) }), "Spirals": (GrayScottDiffusion, { 'coeffs': (0.10, 0.10, 0.018, 0.050) }), "Unstable": (GrayScottDiffusion, { 'coeffs': (0.16, 0.08, 0.020, 0.055) }), "Worms": (GrayScottDiffusion, { 'coeffs': (0.16, 0.08, 0.050, 0.065) }), "Zebrafish": (GrayScottDiffusion, { 'coeffs': (0.16, 0.08, 0.035, 0.060) }), } def __init__(self, size=(400, 400)): super(QGameOfLife, self).__init__() self.size = size self.game = None self.initUI() self.show() def initUI(self): self.setWindowTitle(self.tr("Game of Life")) self.setLayout(QVBoxLayout()) self.layout().setSpacing(0) self.layout().setContentsMargins(0, 0, 0, 0) self.comboBox = QComboBox() self.comboBox.addItems([*QGameOfLife.Games.keys()]) self.comboBox.currentTextChanged.connect(self.select) self.layout().addWidget(self.comboBox) self.scene = QGraphicsScene() self.view = QGraphicsView(self.scene) self.view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setSizePolicy( QSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)) self.view.setFrameShape(QFrame.NoFrame) self.layout().addWidget(self.view) self.item = None self.timer = QTimer() self.timer.setInterval(10) self.timer.timeout.connect(self.tick) initialGame = random.choice([*QGameOfLife.Games.keys()]) self.select(initialGame) self.view.fitInView(self.item, Qt.KeepAspectRatioByExpanding) self.comboBox.setCurrentText(initialGame) def select(self, name: str): self.timer.stop() Game, args = QGameOfLife.Games[name] self.game = Game(self.size, **args) self.tick() self.timer.start() def tick(self): self.game.tick() bitmap = self.game.visualize() image = QImage(bitmap.data, bitmap.shape[1], bitmap.shape[0], QImage.Format_Grayscale8) self.scene.removeItem(self.item) pixmap = QPixmap.fromImage(image) self.item = self.scene.addPixmap(pixmap) def resizeEvent(self, event: QResizeEvent): self.view.fitInView(self.item, Qt.KeepAspectRatioByExpanding) def sizeHint(self) -> QSize: return QSize(self.size[0], self.size[1])
class QGameOfLife(QWidget): Games = { "Game of Life": (GameOfLife, {'fill_rate': 0.50}), "Bacteria": (GrayScottDiffusion, {'coeffs': (0.16, 0.08, 0.035, 0.065)}), "Coral": (GrayScottDiffusion, {'coeffs': (0.16, 0.08, 0.062, 0.062)}), "Fingerprint": (GrayScottDiffusion, {'coeffs': (0.19, 0.05, 0.060, 0.062)}), "Spirals": (GrayScottDiffusion, {'coeffs': (0.10, 0.10, 0.018, 0.050)}), "Unstable": (GrayScottDiffusion, {'coeffs': (0.16, 0.08, 0.020, 0.055)}), "Worms": (GrayScottDiffusion, {'coeffs': (0.16, 0.08, 0.050, 0.065)}), "Zebrafish": (GrayScottDiffusion, {'coeffs': (0.16, 0.08, 0.035, 0.060)}), } def __init__(self, size=(400, 400)): super(QGameOfLife, self).__init__() self.size = size self.game = None self.initUI() self.show() def initUI(self): self.setWindowTitle(self.tr("Game of Life")) self.setLayout(QVBoxLayout()) self.layout().setSpacing(0) self.layout().setContentsMargins(0, 0, 0, 0) self.comboBox = QComboBox() self.comboBox.addItems([*QGameOfLife.Games.keys()]) self.comboBox.currentTextChanged.connect(self.select) self.layout().addWidget(self.comboBox) self.scene = QGraphicsScene() self.view = QGraphicsView(self.scene) self.view.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.view.setSizePolicy(QSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)) self.view.setFrameShape(QFrame.NoFrame) self.layout().addWidget(self.view) self.item = None self.timer = QTimer() self.timer.setInterval(10) self.timer.timeout.connect(self.tick) initialGame = random.choice([*QGameOfLife.Games.keys()]) self.select(initialGame) self.view.fitInView(self.item, Qt.KeepAspectRatioByExpanding) self.comboBox.setCurrentText(initialGame) def select(self, name: str): self.timer.stop() Game, args = QGameOfLife.Games[name] self.game = Game(self.size, **args) self.tick() self.timer.start() def tick(self): self.game.tick() bitmap = self.game.visualize() image = QImage(bitmap.data, bitmap.shape[1], bitmap.shape[0], QImage.Format_Grayscale8) self.scene.removeItem(self.item) pixmap = QPixmap.fromImage(image) self.item = self.scene.addPixmap(pixmap) def resizeEvent(self, event: QResizeEvent): self.view.fitInView(self.item, Qt.KeepAspectRatioByExpanding) def sizeHint(self) -> QSize: return QSize(self.size[0], self.size[1])
class CanvasGraphicsView(QGraphicsView): onSelection = Signal(PickNode) requestEditMode = Signal(bool) def __init__(self, parent=None): super(CanvasGraphicsView, self).__init__(parent) self.setFocusPolicy(Qt.StrongFocus) # Scene properties self.setAcceptDrops(True) self.setMouseTracking(True) self.setRenderHint(QPainter.Antialiasing) self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) self.setResizeAnchor(QGraphicsView.AnchorUnderMouse) self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) self.setBackgroundBrush(QBrush(QColor(51, 51, 51))) self.setFrameShape(QFrame.NoFrame) self.setDragMode(QGraphicsView.RubberBandDrag) self.ViewportUpdateMode(QGraphicsView.BoundingRectViewportUpdate) self.init() def init(self): self.piiPath = str() self._model = {'background': str()} self._isPanning = False self._isZooming = False self._mousePressed = False self._scene = QGraphicsScene() self._scene.selectionChanged.connect(self.update_node_settings) self._backgroundNode = QGraphicsPixmapItem() self._scene.addItem(self._backgroundNode) self._orderSelected = list() self._lastPos = QPoint(0, 0) self.editMode = False self._namespace = str() self._dragMulti = list() self._defaultColor = QColor(255, 255, 255) self._defaultTextColor = QColor(0, 0, 0) self._defaultTextSize = 20 self._defaultText = "New Node" self.workHight = 2160 self.workWidth = 4096 self.setScene(self._scene) self.setContextMenuPolicy(Qt.CustomContextMenu) self.setBackgroundImage(str()) def update_node_settings(self): if self._orderSelected: node = self._orderSelected[-1] self._defaultText = node.toPlainText() self._defaultColor = node.Background self._defaultTextColor = node.defaultTextColor() self._defaultTextSize = node.font().pointSize() def update_maya_selection(self): ''' Update Maya Scene base on active selection. ''' clearSelection() selection = list() for each in self._orderSelected: selection += each.Items if selection: selectObjects(selection) def setBackgroundImage(self, path=str()): ''' Set background image Parameters ---------- path: (str) Path to background image. ''' self._model['background'] = path self.setStatusTip(self._model['background']) pixmap = QPixmap(self._model['background']) self._backgroundNode.setPixmap(pixmap) def getBackgroundImage(self): ''' Get background image ''' return self._model['background'] BackgroundImage = property(getBackgroundImage, setBackgroundImage) def actionMenu(self, QPos): ''' Show action menu. Parameters ---------- QPos: (list) list of x and y location. ''' self.mainMenu = QMenu() add_action = self.mainMenu.addAction('Add A Button') add_action.setEnabled(self.editMode) add_action.triggered.connect(self.add_node) addMany_action = self.mainMenu.addAction('Add Many Buttons') addMany_action.setEnabled(self.editMode) addMany_action.triggered.connect(self.add_multiple_nodes) activeNode = self.mouse_on_node() if activeNode: update_action = self.mainMenu.addAction('Update Button') update_action.setEnabled(self.editMode) update_action.triggered.connect( lambda: self.update_node(activeNode)) delete_action = self.mainMenu.addAction('Delete Button') delete_action.setEnabled(self.editMode) delete_action.setShortcut('Backspace') delete_action.triggered.connect(self.removeSelected) self.mainMenu.addSeparator() # search for selected ButtonNode btnStatus = [ isinstance(n, ButtonNode) for n in self._scene.selectedItems() ] if True in btnStatus: # first ButtonNode activeNode = self._scene.selectedItems()[btnStatus.index(True)] command_action = self.mainMenu.addAction('Edit Command Button...') command_action.setEnabled(self.editMode) command_action.triggered.connect( lambda: self.update_ButtonNode(activeNode)) else: command_action = self.mainMenu.addAction('add Command Button...') command_action.setEnabled(self.editMode) command_action.triggered.connect(self.add_commands) self.mainMenu.addSeparator() reset_action = self.mainMenu.addAction('Reset View') reset_action.setShortcut('H') reset_action.triggered.connect(self.reset_view) frame_action = self.mainMenu.addAction('Frame View') frame_action.setShortcut('F') frame_action.triggered.connect(self.frame_view) self.mainMenu.addSeparator() alignGrp = QMenu('Align') self.mainMenu.addMenu(alignGrp) hac_action = alignGrp.addAction('Horizontal Align Center') hac_action.setIcon(QIconSVG('h_align-01')) hac_action.setEnabled(self.editMode) hac_action.triggered.connect(self.align_horizontal) vac_action = alignGrp.addAction('Vertical Align Center') vac_action.setIcon(QIconSVG('v_align-01')) vac_action.setEnabled(self.editMode) vac_action.triggered.connect(self.align_vertical) hd_action = alignGrp.addAction('Horizontal Distribute') hd_action.setIcon(QIconSVG('h_d_align-01')) hd_action.setEnabled(self.editMode) hd_action.triggered.connect(self.align_horizontal_distribute) vd_action = alignGrp.addAction('Vertical Distribute') vd_action.setIcon(QIconSVG('v_d_align-01')) vd_action.setEnabled(self.editMode) vd_action.triggered.connect(self.align_vertical_distribute) alignGrp.addSeparator() ins_action = alignGrp.addAction('Increase Size') ins_action.setShortcut('+') ins_action.setEnabled(self.editMode) ins_action.triggered.connect(self.increase_size) dis_action = alignGrp.addAction('Decrease Size') dis_action.setShortcut('-') dis_action.setEnabled(self.editMode) dis_action.triggered.connect(self.decrease_size) self.mainMenu.addSeparator() edit_mode = self.mainMenu.addAction('Edit Mode') edit_mode.setCheckable(True) edit_mode.setChecked(self.editMode) edit_mode.triggered.connect( lambda: self.request_edit(not self.editMode)) pos = self.mapToGlobal(QPoint(0, 0)) self.mainMenu.move(pos + QPos) self.mainMenu.show() def mouse_on_node(self): globPosition = self.mapFromGlobal(QCursor.pos()) scenePosition = self.mapToScene(globPosition) for node in self._scene.items(): if isinstance(node, PickNode): if node.mapRectToScene( node.boundingRect()).contains(scenePosition): return node return None def update_node(self, node=PickNode): ''' Update the Node selection base on selection in maya. ''' mayaScene = getActiveItems() # for each in self._scene.selectedItems(): node.Items = mayaScene def update_ButtonNode(self, node=ButtonNode): ''' Update the ButtonNode commands. Parameters ---------- node: (ButtonNode) ButtonNode Node. ''' self.newCommand = CommandDialog(text=node.toPlainText(), cmd=node.Command, cmdType=node.CommandsType) if self.newCommand.exec_() == QDialog.Accepted: data = self.newCommand.Raw node.setPlainText(data[PIIButton.TEXT]) node.Command = data[PIIButton.COMMAND] node.CommandsType = data[PIIButton.COMMANDTYPE] def add_commands(self): ''' Create a new ButtonNode with Commands. ''' globPosition = self.mapFromGlobal(QCursor.pos()) scenePosition = self.mapToScene(globPosition) self.newCommand = CommandDialog() if self.newCommand.exec_() == QDialog.Accepted: data = self.newCommand.Raw self.create_button(position=scenePosition, text=data[PIIButton.TEXT], size=self._defaultTextSize, textColor=self._defaultTextColor, bgColor=self._defaultColor, cmd=data[PIIButton.COMMAND], cmdType=data[PIIButton.COMMANDTYPE]) def align_horizontal(self): ''' Align the selection to center horizontally. ''' selected = self._scene.selectedItems() if len(selected) > 1: minValue = selected[0].y() maxValue = selected[0].y() whole = None for each in selected: y = each.y() value = y + each.boundingRect().height() # finding lowest value minValue = y if y < minValue else minValue minValue = value if value < minValue else minValue # finding highest value maxValue = y if y > maxValue else maxValue maxValue = value if value > maxValue else maxValue total = maxValue - minValue if total != 0: middle = (maxValue + minValue) / 2 for each in selected: center = each.shape().boundingRect().center() start_y = each.y() offset = start_y + center.y() - middle each.setY(each.y() - offset) def align_vertical(self): ''' Align the selection to center vertically. ''' selected = self._scene.selectedItems() if len(selected) > 1: # sort it based on x position + width selected = sorted(selected, key=lambda x: x.x() + x.boundingRect().width()) leftNode = selected[0] rightNode = selected[-1] # total length of x axis total = rightNode.boundingRect().width() + rightNode.x( ) - leftNode.x() if total != 0: middle = (total / 2) + leftNode.x() for each in selected: center = each.shape().boundingRect().center() start_x = each.x() offset = start_x + center.x() - middle each.setX(each.x() - offset) def align_horizontal_distribute(self): ''' Disturbute the selected nodes evenly between first node on the left and last node on the right horizontally. ''' selected = self._scene.selectedItems() if len(selected) > 2: # sort it based on x position + width selected = sorted(selected, key=lambda x: x.x() + x.boundingRect().width()) startItem = selected.pop(0) endItem = selected.pop(-1) # total length of items itemsLength = int() for each in selected: itemsLength += each.boundingRect().width() startPoint = startItem.x() + startItem.boundingRect().width() total = endItem.x() - startPoint section_num = len(selected) + 1 extraSpace = total - itemsLength # nicly divide if extraSpace > 0: gap = extraSpace / section_num nextPlace = startPoint for each in selected: newLoc = nextPlace + gap nextPlace += gap + each.boundingRect().width() each.setX(newLoc) else: total = endItem.x() - startPoint gap = total / section_num nextPlace = startPoint for each in selected: nextPlace += gap each.setX(nextPlace) else: errorMes("PUPPETMASTER-INFO: Select more than 2 nodes.") def align_vertical_distribute(self): ''' Disturbute the selected nodes evenly between first node on the top and last node on the bottom vertically. ''' selected = self._scene.selectedItems() if len(selected) > 2: # sort it based on y position + width selected = sorted( selected, key=lambda node: node.y() + node.boundingRect().height()) startItem = selected.pop(0) endItem = selected.pop(-1) # total length of items itemsLength = int() for each in selected: itemsLength += each.boundingRect().height() startPoint = startItem.y() + startItem.boundingRect().height() total = endItem.y() - startPoint section_num = len(selected) + 1 extraSpace = total - itemsLength # nicly divide if extraSpace > 0: gap = extraSpace / section_num nextPlace = startPoint for each in selected: newLoc = nextPlace + gap nextPlace += gap + each.boundingRect().height() each.setY(newLoc) else: total = endItem.y() - startPoint gap = total / section_num nextPlace = startPoint for each in selected: nextPlace += gap each.setY(nextPlace) else: errorMes("PUPPETMASTER-INFO: Select more than 2 nodes.") def reset_view(self): ''' Fit all the items to the view. ''' items = self._scene.items() if items: rects = [ item.mapToScene(item.boundingRect()).boundingRect() for item in items ] rect = self.min_bounding_rect(rects) self._scene.setSceneRect(rect) self.fitInView(rect, Qt.KeepAspectRatio) def frame_view(self): ''' Fit selected items to the view. ''' items = self._scene.selectedItems() if items: rects = [ item.mapToScene(item.boundingRect()).boundingRect() for item in items ] rect = self.min_bounding_rect(rects) self.fitInView(rect, Qt.KeepAspectRatio) def fit_contents(self): ''' Update the scene boundery. ''' items = self._scene.items() if items: rects = [ item.mapToScene(item.boundingRect()).boundingRect() for item in items ] rect = self.min_bounding_rect(rects) self._scene.setSceneRect(rect) def request_edit(self, value=bool): self.requestEditMode.emit(value) def min_bounding_rect(self, rectList=list()): ''' Get the minimum boundry based on objects in the scene. Parameters ---------- rectList: (list) List of QRectF (boundry of objects) Return ------ out: (QRectF) Get the minimum boundry ''' minX = rectList[0].left() minY = rectList[0].top() maxX = rectList[0].right() maxY = rectList[0].bottom() for k in range(1, len(rectList)): minX = min(minX, rectList[k].left()) minY = min(minY, rectList[k].top()) maxX = max(maxX, rectList[k].right()) maxY = max(maxY, rectList[k].bottom()) return QRectF(minX, minY, maxX - minX, maxY - minY) def increase_size(self): ''' Increase the size of selected items by 1 unit. ''' selected = self._scene.selectedItems() for each in selected: font = each.font() fontSize = font.pointSize() if fontSize < 99: fontSize += 1 font.setPointSize(fontSize) each.setFont(font) def decrease_size(self): ''' Decrease the size of selected items by 1 unit. ''' selected = self._scene.selectedItems() for each in selected: font = each.font() fontSize = font.pointSize() if fontSize > 1: fontSize -= 1 font.setPointSize(fontSize) each.setFont(font) def is_texture(self, path=str): ''' Check if the texture path is valid. Return ------ out: (bool) True if texture is valide, otherwise False. ''' if path.lower().endswith(IMAGE_FORMATS): return True return False def _QMimeDataToFile(self, data=QMimeData): ''' Get all the filepath from drag file. Parameters ---------- data: (QMimeData) QMimeData of dragged file. ''' files = list() if data.hasUrls: for each in data.urls(): files.append(each.toLocalFile()) return files def _is_dragValid(self, event): ''' Check for draged file validation ''' dragedItems = self._QMimeDataToFile(event.mimeData()) if dragedItems: first_path = dragedItems[0] if self.is_texture(first_path) and self.editMode: return True return False def dragEnterEvent(self, event): event.accept() if self._is_dragValid(event) else event.ignore() def dragMoveEvent(self, event): event.accept() if self._is_dragValid(event) else event.ignore() def dropEvent(self, event): dragedItems = self._QMimeDataToFile(event.mimeData()) if dragedItems: first_path = dragedItems[0] if self.is_texture(first_path): self.setBackgroundImage(path=first_path) event.accept() else: event.ignore() def mousePressEvent(self, event): self._lastPos = event.pos() self._lastScenePos = self.mapToScene(event.pos()) if self._dragMulti: for each in self._dragMulti: each.setSelected(True) self._dragMulti = list() if event.button() == Qt.MiddleButton: self._isPanning = True self.setCursor(QPixmap(iconSVG('nav-pan-02'))) self._dragPos = event.pos() event.accept() elif event.button() == Qt.RightButton: if event.modifiers() == Qt.AltModifier: self._isZooming = True self.setCursor(QPixmap(iconSVG('nav-zoom-02'))) self._dragPos = event.pos() self._dragPos2 = self.mapToScene(event.pos()) else: self.actionMenu(event.pos()) event.accept() else: super(CanvasGraphicsView, self).mousePressEvent(event) def mouseMoveEvent(self, event): if self._dragMulti and len(self._dragMulti) > 1: start = self._lastScenePos end = self.mapToScene(event.pos()) total = len(self._dragMulti) - 1 xLength = start.x() - end.x() yLength = start.y() - end.y() xStep = 0 if xLength == 0 else -(xLength / total) yStep = 0 if yLength == 0 else -(yLength / total) num = 0 for each in self._dragMulti: position = QPointF(start.x() + (num * xStep), start.y() + (num * yStep)) each.setPos(position) num += 1 if self._isPanning: newPos = event.pos() diff = newPos - self._dragPos self._dragPos = newPos self.horizontalScrollBar().setValue( self.horizontalScrollBar().value() - diff.x()) self.verticalScrollBar().setValue( self.verticalScrollBar().value() - diff.y()) event.accept() elif self._isZooming: newPos = event.pos() diff = newPos - self._dragPos self._dragPos = newPos factor = 1.000 if diff.x() < 0: factor = 0.98 else: factor = 1.02 self.scale(factor, factor) event.accept() else: if event.modifiers() == Qt.ShiftModifier: diff = event.pos() - self._lastPos x = event.x() if abs(diff.x()) > abs( diff.y()) else self._lastPos.x() y = event.y() if abs(diff.y()) > abs( diff.x()) else self._lastPos.y() event = QMouseEvent(QEvent.MouseMove, QPoint(x, y), self.mapToGlobal(QPoint(x, y)), Qt.LeftButton, Qt.LeftButton, Qt.NoModifier) super(CanvasGraphicsView, self).mouseMoveEvent(event) def mouseReleaseEvent(self, event): self._isPanning = False self._isZooming = False self.setCursor(Qt.ArrowCursor) super(CanvasGraphicsView, self).mouseReleaseEvent(event) self.fit_contents() self.update_maya_selection() def keyPressEvent(self, event): if event.key() == Qt.Key_Backspace or event.key() == Qt.Key_Delete: if self.editMode: self.removeSelected() elif event.key() == Qt.Key_Plus: if self.editMode: self.increase_size() elif event.key() == Qt.Key_Minus: if self.editMode: self.decrease_size() elif event.key() == Qt.Key_H: self.reset_view() elif event.key() == Qt.Key_F: self.frame_view() else: super(CanvasGraphicsView, self).keyPressEvent(event) def removeSelected(self): ''' Remove selected Items. ''' for each in self._scene.selectedItems(): self._scene.removeItem(each) self.remove_stack(each) def wheelEvent(self, event): factor = 1.05 if event.delta() < 0: # factor = .2 / factor factor = 0.95 self.scale(factor, factor) self.update() def add_node(self): ''' Add a new PickNode to the scene. ''' # Cursor Position on Scene globPosition = self.mapFromGlobal(QCursor.pos()) scenePosition = self.mapToScene(globPosition) self.create_node(text=self._defaultText, size=self._defaultTextSize, textColor=self._defaultTextColor, bgColor=self._defaultColor, position=scenePosition, items=getActiveItems(), shape=PickShape.SQUARE) def add_multiple_nodes(self): ''' Add group of PickNode bellow each other to the scene. ''' # Cursor Position on Scene globPosition = self.mapFromGlobal(QCursor.pos()) scenePosition = self.mapToScene(globPosition) self._dragMulti = list() for each in getActiveItems(): node = self.create_node(text=self._defaultText, size=self._defaultTextSize, textColor=self._defaultTextColor, bgColor=self._defaultColor, position=scenePosition, items=[each], shape=PickShape.SQUARE) self._dragMulti.append(node) # scenePosition = QPointF(scenePosition.x(), node.y() + node.boundingRect().height() + 5) def create_node(self, position=list, text=str, size=int, textColor=QColor, bgColor=QColor, items=list, shape=PickShape.SQUARE): ''' Create a new PickNode. Parameters ---------- position: (list) List of x and y location. text: (str) Name of the text. size: (int) Size of the text. textColor: (QColor) Color of the text. bgColor: (QColor) Background Color of the node. items: (list) List of selected Maya object. Return ------ out: (PickNode) Reference of created Node. ''' textNode = PickNode() font = QFont("SansSerif", size) font.setStyleHint(QFont.Helvetica) textNode.setFont(font) textNode.setDefaultTextColor(textColor) textNode.setFlag(QGraphicsItem.ItemIsMovable, self.editMode) textNode.setFlag(QGraphicsItem.ItemIsSelectable) # textNode.setFlag(QGraphicsItem.ItemIsFocusable, self.editMode) textNode.Background = bgColor textNode.Items = items textNode.Shape = shape textNode.onSelected.connect(lambda: self.onSelection.emit(textNode)) textNode.onAddToStack.connect(lambda: self.add_stack(textNode)) textNode.onRemoveFromStack.connect(lambda: self.remove_stack(textNode)) textNode.setPos(position) textNode.setPlainText(text) self._scene.addItem(textNode) return textNode def create_button(self, position=list, text=str, size=int, textColor=QColor, bgColor=QColor, cmd=str, cmdType=str): ''' Create a new ButtonNode. Parameters ---------- position: (list) List of x and y location. text: (str) Name of the text. size: (int) Size of the text. textColor: (QColor) Color of the text. bgColor: (QColor) Background Color of the node. cmd: (str) Command to run when it's pressed. cmdType: (str) Type of command.("python"/"mel") ''' btnNode = ButtonNode() font = QFont("SansSerif", size) font.setStyleHint(QFont.Helvetica) btnNode.setFont(font) btnNode.setDefaultTextColor(textColor) btnNode.setFlag(QGraphicsItem.ItemIsMovable, self.editMode) btnNode.setFlag(QGraphicsItem.ItemIsSelectable) btnNode.Background = bgColor btnNode.CommandsType = cmdType btnNode.Command = cmd # btnNode.onSelected.connect(lambda: self.onSelection.emit(textNode)) btnNode.onSelected.connect(lambda: self.onSelection.emit(btnNode)) btnNode.onClicked.connect(self.scriptJob) btnNode.setPos(position) btnNode.setPlainText(text) self._scene.addItem(btnNode) def scriptJob(self, cmdType=str, cmd=str): ''' Run a command. Parameters ---------- cmd: (str) Command to run. cmdType: (str) Type of command.("python"/"mel") ''' if not self.editMode: if cmdType == CommandType.PYTHON: runPython(cmd) elif cmdType == CommandType.MEL: runMel(cmd) def add_stack(self, node=PickNode): ''' Add a node selection in right order into the stack. Parameters ---------- node: (PickNode) Selected node. ''' self._orderSelected.append(node) def remove_stack(self, node=PickNode): ''' Remove a node from the stack. Parameters ---------- node: (PickNode) Selected node. ''' if node in self._orderSelected: index = self._orderSelected.index(node) self._orderSelected.pop(index) def get_edit(self): return self.editMode def set_edit(self, value=bool): self.editMode = value for each in self._scene.items(): if type(each) == PickNode: each.setFlag(QGraphicsItem.ItemIsMovable, self.editMode) elif type(each) == ButtonNode: each.setFlag(QGraphicsItem.ItemIsMovable, self.editMode) Edit = property(get_edit, set_edit) def get_path(self): return self.piiPath def set_path(self, path=str): self.piiPath = path Path = property(get_path, set_path) def get_raw(self): ''' Get the scene information. (can be be save in .pii) Return ------ out: (dict) Dictionary of scene date to be save in .pii file. ''' image_data = str() pixmap = self._backgroundNode.pixmap() # Extract Image Data if not pixmap.isNull(): buffer = QBuffer() buffer.open(QIODevice.WriteOnly) pixmap.save(buffer, "PNG") # Image Data image_data = bytes(buffer.data().toBase64()).decode('ascii') nodeList = [] for each in self._scene.items(): if type(each) == PickNode: textColor = each.defaultTextColor() bgColor = each.Background item = { PIIPick.TYPE: PIINode.PICK, PIIPick.TEXT: each.toPlainText(), PIIPick.SIZE: each.font().pointSize(), PIIPick.POSITION: (each.pos().x(), each.pos().y()), PIIPick.COLOR: (textColor.red(), textColor.green(), textColor.blue()), PIIPick.BACKGROUND: (bgColor.red(), bgColor.green(), bgColor.blue()), PIIPick.SELECTION: each.Items, PIIPick.SHAPE: each.Shape } nodeList.append(item) elif type(each) == ButtonNode: textColor = each.defaultTextColor() bgColor = each.Background item = { PIIButton.TYPE: PIINode.BUTTON, PIIButton.TEXT: each.toPlainText(), PIIButton.SIZE: each.font().pointSize(), PIIButton.POSITION: (each.pos().x(), each.pos().y()), PIIButton.COLOR: (textColor.red(), textColor.green(), textColor.blue()), PIIButton.BACKGROUND: (bgColor.red(), bgColor.green(), bgColor.blue()), PIIButton.COMMAND: each.Command, PIIButton.COMMANDTYPE: each.CommandsType } nodeList.append(item) rawData = { PII.VERSION: "1.0.0", PII.BACKGROUND: image_data, PII.NODES: nodeList } return rawData def set_raw(self, data=dict): ''' set the scene information. (information from .pii) Parameters ---------- data: (dict) Dictionary of date from .pii file. ''' if data: if data[PII.VERSION] == "1.0.0": self.load_1_0_0(data) Raw = property(get_raw, set_raw) def get_namespace(self): ''' Get namespace of all PickNode. Return ------ out: (list) List of namespaces. ''' namespaceList = [] for each in self._scene.items(): if type(each) == PickNode: valueList = each.Items for sObj in valueList: if ":" in sObj: group = sObj.split(":")[:-1] for index in range(len(group)): namespaceList.append(":".join(group[:index + 1])) return list(set(namespaceList)) def set_namespace(self, data=dict): ''' Set namespace of all PickNode. Parameters ---------- data: (dict) Dictionary of namespace with value of new namespace. ''' for each in self._scene.items(): if type(each) == PickNode: valueList = each.Items newValue = list() for sObj in valueList: if ":" in sObj: # namesapce nameS = ":".join(sObj.split(":")[:-1]) # object name object_name = sObj.split(":")[-1] keys = data.keys() keys.sort(reverse=True) for key in keys: if key in nameS: nameS = nameS.replace(key, data[key], 1) # making sure doesn't start with ':' nameS = nameS[1:] if nameS.startswith(":") else nameS # add the object to namespace nameS = ":".join([nameS, object_name ]) if nameS else object_name newValue.append(nameS) else: newValue.append(sObj) each.Items = newValue Namespace = property(get_namespace, set_namespace) def get_NSHistory(self): return self._namespace def set_NSHistory(self, name=str): self._namespace = name NamespaceHistory = property(get_NSHistory, set_NSHistory) def get_highlight(self): return def set_highlight(self, data=list): if data: for each in self._scene.items(): # QApplication.processEvents() if type(each) == PickNode: for item in data: if item in each.Items: each.Highlight = True break else: each.Highlight = False else: for each in self._scene.items(): if type(each) == PickNode: each.Highlight = False Highlight = property(get_highlight, set_highlight) def clear_scene(self): ''' Clear the scene. ''' self._orderSelected = list() self._scene.clear() self._backgroundNode = QGraphicsPixmapItem() self._scene.addItem(self._backgroundNode) self.reset_view() def is_changed(self): ''' Check for the scene changes. ''' if self._backgroundNode.pixmap(): return True elif len(self._scene.items()) > 1: return True return False def load_1_0_0(self, data=dict): ''' Load v1.0.0 of .pii version file. Parameters ---------- data: (dict) Dictionary of date from .pii file. ''' if data[PII.BACKGROUND]: # Import Image Data newPix = QPixmap() newPix.loadFromData( QByteArray.fromBase64(data[PII.BACKGROUND].encode('ascii')), "PNG") self._backgroundNode.setPixmap(newPix) for each in data[PII.NODES]: if each["type"] == PIINode.PICK: self.create_node(text=each[PIIPick.TEXT], size=each[PIIPick.SIZE], textColor=QColor(*each[PIIPick.COLOR]), bgColor=QColor(*each[PIIPick.BACKGROUND]), position=QPointF(*each[PIIPick.POSITION]), items=each[PIIPick.SELECTION], shape=each[PIIPick.SHAPE]) elif each["type"] == PIINode.BUTTON: self.create_button(position=QPointF(*each[PIIButton.POSITION]), text=each[PIIButton.TEXT], size=each[PIIButton.SIZE], textColor=QColor(*each[PIIButton.COLOR]), bgColor=QColor(*each[PIIButton.BACKGROUND]), cmd=each[PIIButton.COMMAND], cmdType=each[PIIButton.COMMANDTYPE]) def set_nodes_bg_color(self, color=QColor): ''' Set background color of selected nodes. Parameters ---------- color: (QColor) QColor value. ''' self._defaultColor = color for each in self._scene.selectedItems(): each.Background = color self.update() def set_nodes_font_color(self, color=QColor): ''' Set font color of selected nodes. Parameters ---------- color: (QColor) QColor value. ''' self._defaultTextColor = color for each in self._scene.selectedItems(): each.setDefaultTextColor(color) def set_nodes_font_size(self, size=int): ''' Set font size of selected nodes. Parameters ---------- size: (int) font size. ''' self._defaultTextSize = size for each in self._scene.selectedItems(): font = each.font() font.setPointSize(size) each.setFont(font) def set_nodes_text(self, text=str): ''' Set text for selected nodes. Parameters ---------- text: (str) text for the node. ''' self._defaultText = text for each in self._scene.selectedItems(): each.setPlainText(text) def set_nodes_shape(self, shape=str): ''' Set shape for selected nodes. Parameters ---------- shape: (str) name for the shape. ''' for each in self._scene.selectedItems(): if isinstance(each, PickNode): each.Shape = shape
class MainWindow(QMainWindow): def __init__(self): # load setting json file with open('setting.json') as f: self.app_setting = json.load(f) # Status of window super().__init__() self.title = 'Image Editor' self.left = 70 self.top = 70 self.width = 800 self.height = 700 # Status of view image self.org_qimg = None self.org_img_width = 0 self.org_img_height = 0 self.layer_pixmap = None self.layer_width = 0 self.layer_height = 0 self.layer_alpha = 50.0 # Prepare color bar data self.colormap_gain = self.app_setting["SoftwareSetting"]["process"][ "colormap"]["gain"] self.colormap_offset_x = self.app_setting["SoftwareSetting"][ "process"]["colormap"]["offset_x"] self.colormap_offset_green = self.app_setting["SoftwareSetting"][ "process"]["colormap"]["offset_green"] self.colormap_data = [ colormap.colorBarRGB(x * 0.001, self.colormap_offset_x, self.colormap_offset_green, self.colormap_gain) for x in range(1000) ] self.img_edit_mode = 'cursor' self.draw_color = QColor(255, 0, 0) self.draw_tool_size = 5 self.eraser_color = QColor(0, 0, 0, 0) # setup user interface components self.setup_ui() # Setup user interface components def setup_ui(self): # Set main window title self.setWindowTitle(self.title) # Set main wiodow initial position self.setGeometry(self.left, self.top, self.width, self.height) # Set up mainWindow's layout self.mainWidget = QWidget(self) # Note not to forget this code. self.main_layout = QVBoxLayout() # Set menu for main window self.main_menu = self.menuBar() self.file_menu = self.main_menu.addMenu('File') self.edit_menu = self.main_menu.addMenu('Edit') self.help_menu = self.main_menu.addMenu('Help') self.main_layout.addWidget(self.main_menu) # Set "Original Image Open" menu self.org_img_open_button = QAction( self.style().standardIcon(getattr(QStyle, 'SP_FileDialogStart')), 'Open Orginal Image', self) self.org_img_open_button.setShortcut('Ctrl+O') self.org_img_open_button.triggered.connect(self.open_org_img_dialog) self.file_menu.addAction(self.org_img_open_button) # Set "Save layer image" menu self.layer_img_save_button = QAction( self.style().standardIcon(getattr(QStyle, 'SP_FileDialogEnd')), 'Save Layer Image', self) self.layer_img_save_button.setShortcut('Ctrl+S') self.layer_img_save_button.triggered.connect(self.save_layer_image) self.file_menu.addAction(self.layer_img_save_button) # Set "Save compose image(original + layer image)" menu self.compose_img_save_button = QAction( self.style().standardIcon(getattr(QStyle, 'SP_FileDialogEnd')), 'Save Compose Image', self) self.compose_img_save_button.setShortcut('Ctrl+D') self.compose_img_save_button.triggered.connect(self.save_compose_image) self.file_menu.addAction(self.compose_img_save_button) # Set "exit software" menu self.exit_button = QAction( self.style().standardIcon(getattr(QStyle, 'SP_DialogCloseButton')), 'Exit', self) self.exit_button.setShortcut('Ctrl-Q') self.exit_button.setStatusTip('Exit software') self.exit_button.triggered.connect(self.close) self.file_menu.addAction(self.exit_button) self.upper_layout = QHBoxLayout() self.main_layout.addLayout(self.upper_layout) # Set image display area self.gview_default_size = 500 self.graphics_view = QGraphicsView() self.graphics_view.setFixedSize(self.gview_default_size, self.gview_default_size) self.graphics_view.setObjectName("imageDisplayArea") self.upper_layout.addWidget(self.graphics_view) # image display area's contents self.scene = GraphicsSceneForMainView(self.graphics_view, self) self.imgs_pixmap = [] self.imgs = [] self.img_status_layout = QVBoxLayout() self.upper_layout.addLayout(self.img_status_layout) # Set tranparency value of layer image self.transparency_title_label = QLabel('layer transparency value') self.img_status_layout.addWidget(self.transparency_title_label) transparency = round((1.0 - self.layer_alpha / 255.0) * 100) self.img_transparency_edit = QLineEdit(str(transparency)) self.img_transparency_sld = QSlider(Qt.Horizontal) self.img_transparency_sld.setFocusPolicy(Qt.NoFocus) self.img_transparency_sld.setRange(0, 100) self.img_transparency_sld.setValue(transparency) self.transparency_layout = QFormLayout() self.transparency_layout.addRow(self.img_transparency_sld, self.img_transparency_edit) self.img_status_layout.addLayout(self.transparency_layout) # Signal of transparency value changed self.img_transparency_sld.valueChanged.connect( self.transparency_change_sld) self.img_transparency_edit.textChanged.connect( self.transparency_change_edit) self.img_editor_layout = QVBoxLayout() self.img_status_layout.addLayout(self.img_editor_layout) # Set layer image editor tool self.img_editor_tool1_layout = QHBoxLayout() self.img_editor_layout.addLayout(self.img_editor_tool1_layout) self.tool_button_size = 64 # Set Mouse cursor self.mouse_cursor_button = QPushButton() self.mouse_cursor_button.setIcon(QIcon('icon/cursor.png')) self.mouse_cursor_button.setCheckable(True) self.mouse_cursor_button.setIconSize( QSize(self.tool_button_size, self.tool_button_size)) self.img_editor_tool1_layout.addWidget(self.mouse_cursor_button) # Set Pen self.pen_button = QPushButton() self.pen_button.setIcon(QIcon('icon/pen.png')) self.pen_button.setCheckable(True) self.pen_button.setIconSize( QSize(self.tool_button_size, self.tool_button_size)) self.img_editor_tool1_layout.addWidget(self.pen_button) # Set Eraser self.eraser_button = QPushButton() self.eraser_button.setIcon(QIcon('icon/eraser.png')) self.eraser_button.setCheckable(True) self.eraser_button.setIconSize( QSize(self.tool_button_size, self.tool_button_size)) self.img_editor_tool1_layout.addWidget(self.eraser_button) # Group button of mouse cursor, pen, eraser self.img_editor_tool1_group1 = QButtonGroup() self.img_editor_tool1_group1.addButton(self.mouse_cursor_button, 1) self.img_editor_tool1_group1.addButton(self.pen_button, 2) self.img_editor_tool1_group1.addButton(self.eraser_button, 3) # Set signal-slot of image editor button self.mouse_cursor_button.toggled.connect( self.mouse_cursor_button_toggled) self.pen_button.toggled.connect(self.pen_button_toggled) self.eraser_button.toggled.connect(self.eraser_button_toggled) # Set color bar self.color_bar_width = 64 self.color_bar_height = 256 self.color_bar_view = QGraphicsView() self.color_bar_view.setFixedSize(self.color_bar_width + 3, self.color_bar_height + 3) self.color_bar_scene = GraphicsSceneForTools() self.color_bar_img = QImage(self.color_bar_width, self.color_bar_height, QImage.Format_RGB888) for i in range(self.color_bar_height): # Set drawing pen for colormap ii = round(i * (1000 / 256)) color = QColor(self.colormap_data[ii][0], self.colormap_data[ii][1], self.colormap_data[ii][2]) pen = QPen(color, 1, Qt.SolidLine, \ Qt.SquareCap, Qt.RoundJoin) self.color_bar_scene.addLine(0, self.color_bar_height - i - 1, self.color_bar_width, self.color_bar_height - i - 1, pen=pen) for j in range(self.color_bar_width): self.color_bar_img.setPixelColor(j, self.color_bar_height - i - 1, color) self.color_bar_scene.set_img_content(self.color_bar_img) self.color_bar_view.setScene(self.color_bar_scene) # Connect signal to slot of color_bar_scene self.color_bar_scene.img_info.connect(self.set_selected_color) self.img_editor_tool1_layout.addWidget(self.color_bar_view) # Set thickness of Pen or Eraser self.draw_status_layout = QVBoxLayout() self.draw_thick_title_label = QLabel('thickness of pen or eraser') self.draw_status_layout.addWidget(self.draw_thick_title_label) self.draw_thick_edit = QLineEdit(str(self.draw_tool_size)) self.draw_thick_sld = QSlider(Qt.Horizontal) self.draw_thick_sld.setFocusPolicy(Qt.NoFocus) self.draw_thick_sld.setRange(1, 30) self.draw_thick_sld.setValue(self.draw_tool_size) self.draw_thick_layout = QFormLayout() self.draw_thick_layout.addRow(self.draw_thick_sld, self.draw_thick_edit) self.draw_status_layout.addLayout(self.draw_thick_layout) self.img_editor_layout.addLayout(self.draw_status_layout) # Signal of draw thickness value changed self.draw_thick_sld.valueChanged.connect(self.draw_thick_change_sld) self.draw_thick_edit.textChanged.connect(self.draw_thick_change_edit) # Set view area of selected color self.select_color_view_size = 64 self.select_color_view = QGraphicsView() self.select_color_view.setFixedSize(self.select_color_view_size + 3, self.select_color_view_size + 3) self.select_color_scene = QGraphicsScene() brush = QBrush(self.draw_color) self.select_color_rect = self.select_color_scene.addRect(QRect(0, 0, self.select_color_view_size, self.select_color_view_size), \ brush=brush) self.select_color_view.setScene(self.select_color_scene) self.select_color_title_label = QLabel('Selected color') self.selected_color_layout = QFormLayout() self.selected_color_layout.addRow(self.select_color_title_label, self.select_color_view) self.img_editor_layout.addLayout(self.selected_color_layout) # Set save button self.save_button_layout = QHBoxLayout() self.img_status_layout.addLayout(self.save_button_layout) self.layer_save_button = QPushButton('Save layer image') self.layer_save_button.setIcon(QIcon('icon/layer_save.png')) self.layer_save_button.setIconSize( QSize(self.tool_button_size, self.tool_button_size)) self.compose_save_button = QPushButton( 'Save composed original and layer image') self.compose_save_button.setIcon(QIcon('icon/compose_save.png')) self.compose_save_button.setIconSize( QSize(self.tool_button_size, self.tool_button_size)) self.save_button_layout.addWidget(self.layer_save_button) self.save_button_layout.addWidget(self.compose_save_button) self.layer_save_button.clicked.connect(self.save_layer_image) self.compose_save_button.clicked.connect(self.save_compose_image) # Set display area of selected file path self.org_img_path_title_label = QLabel('original image file: ') self.org_img_path_label = QLabel('') self.file_path_layout = QFormLayout() self.file_path_layout.addRow(self.org_img_path_title_label, self.org_img_path_label) self.bottom_layout = QVBoxLayout() self.bottom_layout.addLayout(self.file_path_layout) self.main_layout.addLayout(self.bottom_layout) self.mainWidget.setLayout(self.main_layout) self.setCentralWidget(self.mainWidget) # Original image select Function def open_org_img_dialog(self): options = QFileDialog.Options() org_img_default_path = self.app_setting["SoftwareSetting"][ "file_path"]["org_img_dir"] self.org_img_file_path, selected_filter = QFileDialog.getOpenFileName(self, 'Select original image', org_img_default_path, \ 'Image files(*.jpg *jpeg *.png)', options=options) org_img_dir_path, org_img_file = os.path.split(self.org_img_file_path) org_img_bare_name, org_img_ext = os.path.splitext(org_img_file) self.org_img_path_label.setText(self.org_img_file_path) self.set_image_on_viewer() def set_image_on_viewer(self): # Delete existing image item if len(self.imgs_pixmap) != 0: for item in self.imgs_pixmap: self.scene.removeItem(item) self.scene.clear_contents() self.imgs_pixmap.clear() self.imgs.clear() # load original image self.org_qimg = QImage(self.org_img_file_path) self.org_pixmap = QPixmap.fromImage(self.org_qimg) org_img_size = self.org_qimg.size() self.org_img_width = org_img_size.width() self.org_img_height = org_img_size.height() # Set layer image self.layer_qimg = QImage(self.org_img_width, self.org_img_height, QImage.Format_RGBA8888) self.layer_qimg.fill(QColor(0, 0, 0, self.layer_alpha)) self.layer_pixmap = QPixmap.fromImage(self.layer_qimg) self.imgs.append(self.org_qimg) self.imgs.append(self.layer_qimg) # Set image to scene self.imgs_pixmap.append(QGraphicsPixmapItem(self.org_pixmap)) self.scene.addItem(self.imgs_pixmap[-1]) self.imgs_pixmap.append(QGraphicsPixmapItem(self.layer_pixmap)) self.scene.addItem(self.imgs_pixmap[-1]) self.scene.set_img_contents(self.imgs) # Set scene to graphics view self.graphics_view.setScene(self.scene) self.graphics_view.setAlignment(Qt.AlignHCenter | Qt.AlignVCenter) self.show() # Slot function of transparency slider changed def transparency_change_sld(self, value): self.img_transparency_edit.setText(str(value)) self.layer_alpha = int(255 * (1.0 - (value / 100))) # Change layer image's transparency(alpha value) for y in range(self.org_img_height): for x in range(self.org_img_width): self.layer_qimg.setPixelColor( QPoint(x, y), QColor(0, 0, 0, self.layer_alpha)) self.layer_pixmap = QPixmap.fromImage(self.layer_qimg) # remove previous layer image self.scene.removeItem(self.imgs_pixmap[-1]) self.imgs_pixmap.pop(-1) # add new layer image to scene self.imgs_pixmap.append(QGraphicsPixmapItem(self.layer_pixmap)) self.scene.addItem(self.imgs_pixmap[-1]) self.show() # Slot function of transparency text edit changed def transparency_change_edit(self, value): if int(value) < 0 or int(value) > 100: return self.img_transparency_sld.setValue(int(value)) self.layer_alpha = int(255 * (1.0 - (int(value) / 100.0))) # Change layer image's transparency(alpha value) for y in range(self.org_img_height): for x in range(self.org_img_width): self.layer_qimg.setPixelColor( QPoint(x, y), QColor(0, 0, 0, self.layer_alpha)) self.layer_pixmap = QPixmap.fromImage(self.layer_qimg) # remove previous layer image self.scene.removeItem(self.imgs_pixmap[-1]) self.imgs_pixmap.pop(-1) # add new layer image to scene self.imgs_pixmap.append(QGraphicsPixmapItem(self.layer_pixmap)) self.scene.addItem(self.imgs_pixmap[-1]) self.show() # slot(receiver of signal) of mouse_cursor_button toggled def mouse_cursor_button_toggled(self, checked): if checked: self.img_edit_mode = 'cursor' self.scene.set_mode(self.img_edit_mode) self.color_bar_scene.set_mode(self.img_edit_mode) # slot(receiver of signal) of pen_button toggled def pen_button_toggled(self, checked): if checked: self.img_edit_mode = 'pen' self.scene.set_mode(self.img_edit_mode) self.color_bar_scene.set_mode(self.img_edit_mode) # slot(receiver of signal) of eraser_button toggled def eraser_button_toggled(self, checked): if checked: self.img_edit_mode = 'eraser' self.scene.set_mode(self.img_edit_mode) self.color_bar_scene.set_mode(self.img_edit_mode) self.draw_color = self.eraser_color self.select_color_scene.removeItem(self.select_color_rect) brush = QBrush(self.draw_color) self.select_color_rect = self.select_color_scene.addRect(QRect(0, 0, self.select_color_view_size, self.select_color_view_size), \ brush=brush) self.select_color_view.setScene(self.select_color_scene) # Slot of color bar clicked for selection color def set_selected_color(self, color): # Delete existng image item self.select_color_scene.removeItem(self.select_color_rect) self.draw_color = color brush = QBrush(self.draw_color) self.select_color_rect = self.select_color_scene.addRect(QRect(0, 0, self.select_color_view_size, self.select_color_view_size), \ brush=brush) self.select_color_view.setScene(self.select_color_scene) # Slot function of draw thicikeness slider changed def draw_thick_change_sld(self, value): self.draw_thickness_edit.setText(str(value)) self.draw_tool_size = value # Slot function of draw thicikeness text editor changed def draw_thick_change_edit(self, value): if int(value) < 1 or int(value) > 30: return self.draw_thickness_sld.setValue(int(value)) def make_layer_image(self): for i, line in enumerate(self.scene.lines): pen = self.scene.pens[i] pen_size = int(pen.width()) pen_color = pen.color() # start pixel of line x1 = int(line.x1()) y1 = int(line.y1()) # end pixel of line x2 = int(line.x2()) y2 = int(line.y2()) dx = int(line.dx()) dy = int(line.dy()) # When only 1pixl line if dx <= 1 and dy <= 1: draw_pix_x1_s = max(x1 - int(pen_size / 2), 0) draw_pix_x1_e = min(x1 + int(pen_size / 2), self.org_img_width - 1) draw_pix_y1_s = max(y1 - int(pen_size / 2), 0) draw_pix_y1_e = min(y1 + int(pen_size / 2), self.org_img_height - 1) # for Pen's size for y in range(draw_pix_y1_s, draw_pix_y1_e): for x in range(draw_pix_x1_s, draw_pix_x1_e): self.layer_qimg.setPixelColor(x, y, pen_color) draw_pix_x2_s = max(x2 - int(pen_size / 2), 0) draw_pix_x2_e = min(x2 + int(pen_size / 2), self.org_img_width - 1) draw_pix_y2_s = max(y2 - int(pen_size / 2), 0) draw_pix_y2_e = min(y2 + int(pen_size / 2), self.org_img_height - 1) # for Pen's size for y in range(draw_pix_y2_s, draw_pix_y2_e): for x in range(draw_pix_x2_s, draw_pix_x2_e): self.layer_qimg.setPixelColor(x, y, pen_color) else: # For avoid devide by 0 if dx == 0: for y in range(y1, y2 + 1): draw_pix_y_s = y - int(pen_size / 2) draw_pix_y_e = y + int(pen_size / 2) # for Pen's size for yy in range(draw_pix_y_s, draw_pix_y_e): self.layer_qimg.setPixelColor(x1, yy, pen_color) else: grad = dy / dx # Choose coordinates with small slope not to skip pixels if grad >= 1.0: for x in range(dx): y = y1 + int(grad * x + 0.5) draw_pix_x_s = max(x1 + x - int(pen_size / 2), 0) draw_pix_x_e = min(x1 + x + int(pen_size / 2), self.org_img_width - 1) draw_pix_y_s = max(y - int(pen_size / 2), 0) draw_pix_y_e = min(y + int(pen_size / 2), self.org_img_height - 1) # for Pen's size for yy in range(draw_pix_y_s, draw_pix_y_e + 1): for xx in range(draw_pix_x_s, draw_pix_x_e + 1): self.layer_qimg.setPixelColor( xx, yy, pen_color) else: for y in range(dy): x = x1 + int(1 / grad * y + 0.5) draw_pix_y_s = max(y1 + y - int(pen_size / 2), 0) draw_pix_y_e = min(y1 + y + int(pen_size / 2), self.org_img_height - 1) draw_pix_x_s = max(x - int(pen_size / 2), 0) draw_pix_x_e = min(x + int(pen_size / 2), self.org_img_width - 1) # for Pen's size for yy in range(draw_pix_y_s, draw_pix_y_e + 1): for xx in range(draw_pix_x_s, draw_pix_x_e + 1): self.layer_qimg.setPixelColor( xx, yy, pen_color) # Slot function of save layer image button clicked def save_layer_image(self): self.make_layer_image() layer_img_default_path = self.app_setting["SoftwareSetting"][ "file_path"]["layer_img_dir"] options = QFileDialog.Options() file_name, selected_filete = QFileDialog.getSaveFileName(self, 'Save layer image', layer_img_default_path, \ 'image files(*.png, *jpg)', options=options) #print('layer image save name:{file}'.format(file=file_name)) self.layer_qimg.save(file_name) ret = QMessageBox(self, 'Success', 'layer image is saved successfully', QMessageBox.Ok) # Make composed orignal and layered image def make_compose_image(self): self.make_layer_image() self.compose_qimg = QImage(self.org_img_width, self.org_img_height, QImage.Format_RGBA8888) painter = QPainter(self.compose_qimg) painter.setCompositionMode(QPainter.CompositionMode_SourceOver) painter.drawImage(0, 0, self.org_qimg) painter.setCompositionMode(QPainter.CompositionMode_SourceOver) painter.drawImage(0, 0, self.layer_qimg) painter.end() # Slot function of save composer original and layer image button clicked def save_compose_image(self): self.make_compose_image() compose_img_default_path = self.app_setting["SoftwareSetting"][ "file_path"]["compose_img_dir"] options = QFileDialog.Options() file_name, selected_fileter = QFileDialog.getSaveFileName(self, 'Save composed image', compose_img_default_path, \ 'image files(*.png, *jpg)', options=options) #print('compose image save name:{file}'.format(file=file_name)) self.compose_qimg.save(file_name) ret = QMessageBox(self, 'Success', 'compose image is saved successfully', QMessageBox.Ok)
class HumanVisualizationWidget(QGraphicsView): def __init__(self, parent=None): super(HumanVisualizationWidget, self).__init__(parent) self._scene = QGraphicsScene(self) self.setScene(self._scene) # circle = QGraphicsEllipseItem( 10, 10, 10 ,10) # self._scene.addItem(circle) self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) self.setResizeAnchor(QGraphicsView.AnchorUnderMouse) self._boxes = [] self._humans = {} def load_inner_model(self, file): import xml.etree.cElementTree as ET tree = ET.ElementTree(file=file) root = tree.getroot() transforms = tree.findall(".//transform[plane]") walls = {} for trans in transforms: if 'id' in trans.attrib and 'pared' in trans.attrib['id']: print("Pared:", trans.attrib['id']) current_wall = [0] * 7 # "wall5": [x, y, width, height, posx, posy, 0] if 'tx' in trans.attrib: # print trans.attrib['tx'] current_wall[4] = int(float(trans.attrib['tx'])) if 'ty' in trans.attrib: # print trans.attrib['ty'] current_wall[5] = int(float(trans.attrib['tz'])) # current_wall = planes = trans.findall('plane') for plane in planes: if 'id' in plane.attrib and 'muro' in plane.attrib['id']: # if 'nx' in plane.attrib: # print plane.attrib['nx'] # if 'nz' in plane.attrib: # print plane.attrib['nz'] if 'size' in plane.attrib: # print int(float(plane.attrib['size'].split(',')[0]) current_wall[2] = int( float(plane.attrib['size'].split(',')[0])) / 2. # print int(float(plane.attrib['size'].split(',')[1]) current_wall[3] = int( float(plane.attrib['size'].split(',')[1])) / 2. if current_wall[2] < current_wall[3]: current_wall[2] = 200 else: current_wall[3] = 200 walls[trans.attrib['id']] = current_wall for id in sorted(walls.keys()): object = walls[id] # rect = QRectF(-float(object[2]) / 2, -float(object[3]) / 2, float(object[2]), float(object[3])) rect = QRectF(0, 0, float(object[2]), float(object[3])) border = QPen(QColor("black")) fill = QBrush(QColor("black")) box = self._scene.addRect(rect, border, fill) self._scene.addEllipse( QRectF(float(object[4]), float(object[5]), 10, 10), QPen(QColor("green")), QBrush(QColor("green"))) box.setPos(float(object[4]), float(object[5])) box.setRotation(float(object[6])) self._boxes.append(box) self._scene.update() QApplication.processEvents() sleep(1) def load_custom_json_world(self, file): if not os.path.isfile(file): print("Error reading world file, check config params:", file) return False with open(file, "r") as read_file: json_data = json.load(read_file) types_colors = { "tables": "SandyBrown", "roundTables": "Khaki", "walls": "Brown", "points": "Blue" } self.clear() for type, color in types_colors.items(): if type in json_data: tables = json_data[type] for object in tables.values(): rect = QRectF(-float(object[2]) / 2, -float(object[3]) / 2, float(object[2]), float(object[3])) border = QPen(QColor(color)) fill = QBrush(QColor(color)) if type == "roundTables": box = self._scene.addEllipse(rect, border, fill) else: box = self._scene.addRect(rect, border, fill) box.setPos(float(object[4]), float(object[5])) box.setRotation(float(object[6])) self._boxes.append(box) def load_json_world(self, file): if not os.path.isfile(file): print("Error reading world file, check config params:", file) return False with open(file, "r") as read_file: json_data = json.load(read_file) polygon_points = [] paths_count = 0 for item in json_data: if 'json_geometry' in item: geometry = item['json_geometry'] if geometry['type'] == 'Polygon': for coord in geometry['coordinates'][0]: if isinstance(coord, list) and ( (isinstance(coord, list) and len(coord) == 2) or (len(coord) == 3 and coord[3] == 0)): current_point = QPointF(coord[0], coord[1]) polygon_points.append(current_point) else: print("Unknown coord", geometry["coordinates"][0]) polygon = QPolygonF(polygon_points) path = QPainterPath() path.addPolygon(polygon) contour = QGraphicsPathItem(path) # r = lambda: random.randint(0, 255) # next_color = '#%02X%02X%02X' % (r(), r(), r()) contour.setPen(QPen(QColor("red"), 0.1)) contour.setBrush(QBrush(Qt.transparent)) # if paths_count == 4: print(item['json_featuretype']) self._scene.addItem(contour) paths_count += 1 self.update() def clear(self): for human in self._humans.values(): self._scene.removeItem(human) self._humans = {} # self._scene.setSceneRect(QRectF(0,0,400,400)) def resizeEvent(self, event): # skip initial entry self.own_resize() super(HumanVisualizationWidget, self).resizeEvent(event) def own_resize(self): self.fitInView(self._scene.itemsBoundingRect(), Qt.KeepAspectRatio) self._scene.setSceneRect(self._scene.itemsBoundingRect()) def add_human_by_pos(self, id, pos): x, y = pos human = QGraphicsEllipseItem(0, 0, 200, 200) self._scene.addItem(human) human.setBrush(QBrush(Qt.black, style=Qt.SolidPattern)) human_text = QGraphicsTextItem(str(pos)) font = QFont("Helvetica [Cronyx]", 40, QFont.Bold) human_text.setFont(font) human_text.setParentItem(human) human.setPos(pos[0], pos[1]) self._humans[id] = human human.setZValue(30) def move_human(self, id, pos): x, y = pos human = self._humans[id] human.setPos(x, y) def set_human_color(self, id, color): if id in self._humans: self._humans[id].setBrush(color)
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 BarrenLandsWindow(QMainWindow): """User interface to interact with the barren lands library.""" def __init__(self, width=400, height=600, zones=None): """Initialization of BarrenLands GUI Args: width (int): The width of the canvas/field. height (int): The height of the canvas/field. zones (str): Text data to add to raw input if provided. """ QMainWindow.__init__(self) self.app = QApplication.instance() self.resources = pathlib.Path(__file__).parent self.results = set() self.canvas_width = width self.canvas_height = height self.user_input = zones # Colors for drawing onto the QGraphicsScene self.brush_barren = QBrush(QColor(120, 120, 75), Qt.Dense2Pattern) self.brush_overlay = QBrush(QColor(20, 20, 20, 35), Qt.SolidPattern) self.brush_innerlay = QBrush(QColor(32, 37, 44), Qt.SolidPattern) self.brush_fertile = QBrush(QColor(90, 220, 90, 150), Qt.Dense5Pattern) self.brush_fertile_no_go = QBrush(QColor(90, 220, 90, 150), Qt.BDiagPattern) # Setup the filed class that will handle our calculations. self.field = land.Field(width=self.canvas_width, height=self.canvas_height) self.barren_zones = list() self.fertile_zones = list() self.setup_window() self.build_core_ui() self.build_controls() def setup_window(self): """Setup all the window related settings/flags.""" self.setWindowTitle("Barren Lands") self.setMinimumHeight(600) self.setMinimumWidth(700) self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint) self.setWindowIcon(QIcon(utils.get_icon())) self.setStyleSheet(utils.get_css()) def build_core_ui(self): """Sets up the core elements of this window. Notes: Setting up QGraphicsScene: https://stackoverflow.com/questions/23174481/pyside-qt-layout-not-working """ self.base_widget = QWidget(self) self.layout_base = QHBoxLayout(self.base_widget) self.layout_canvas = QVBoxLayout() self.setContentsMargins(0, 0, 0, 0) # Setup the QGraphics items to display out zones as active elements. scene_zone = QRect(0, 0, self.canvas_width, self.canvas_height) self.scene = QGraphicsScene(scene_zone, self) self.canvas = QGraphicsView(self.scene, self) self.canvas.setFixedSize(self.canvas_width + 1, self.canvas_height + 1) # Draw faint bullseye for target. self.scene.addEllipse(self.get_center_rect(300), Qt.NoPen, self.brush_overlay) self.scene.addEllipse(self.get_center_rect(200), Qt.NoPen, self.brush_innerlay) self.scene.addEllipse(self.get_center_rect(100), Qt.NoPen, self.brush_overlay) self.layout_canvas.addWidget(self.canvas) self.lbl_island = QLabel("Island Areas:") self.layout_canvas.addWidget(self.lbl_island) self.layout_base.addLayout(self.layout_canvas) # Set the core widget to the window. self.setCentralWidget(self.base_widget) def get_center_rect(self, size): """Generates a QRect that is at the center of the canvas and it's width/height set to size. Args: size (int): The size to set the bounds to. Returns: (QRectF): A rectangle coordinate that is center on the canvas with a given size. """ return QRectF(self.canvas_width * 0.5 - (size * 0.5), self.canvas_height * 0.5 - (size * 0.5), size, size) def build_controls(self): """Build the control panel that contains all the inputs.""" self.controls = QWidget(self) self.layout_controls = QVBoxLayout(self.controls) self.layout_controls.setContentsMargins(0, 0, 0, 0) # Initialize the text box for user input self.input = self.build_coord_input() # Create the buttons self.btn_add_bzone = BarrenButton(label="Add", parent=self, cmd=self.ctl_add_input) self.btn_run = BarrenButton(label="Analyze", parent=self, cmd=self.ctl_run) # Build the utility buttons self.layout_btm_btn = QHBoxLayout() self.layout_btm_btn.setContentsMargins(0, 0, 0, 0) self.btn_reset = BarrenButton(label="Reset All", parent=self, cmd=self.ctl_clear_zones) self.btn_debug = BarrenButton(label="Debug", parent=self, cmd=self.ctr_debug) self.layout_btm_btn.addWidget(self.btn_debug) self.layout_btm_btn.addWidget(self.btn_reset) self.results_grp = self.build_results_group() self.lbl_area = QLabel("Total Area: 0") self.lbl_area.setAlignment(Qt.AlignCenter) # Add everything to the controls layout self.layout_controls.addLayout(self.input) self.layout_controls.addWidget(self.btn_add_bzone) self.layout_controls.addWidget(self.btn_run) self.layout_controls.addWidget(self.results_grp) self.layout_controls.addWidget(self.lbl_area) self.layout_controls.addLayout(self.layout_btm_btn) # Add to the windows base layout. self.layout_base.addWidget(self.controls) def build_results_group(self): """Build the results group to contain the results of analysis. Returns: results_grp (QGroupBox): The group box that contains the results layout. """ grp_results = QGroupBox("Results: Largest island zone.") grp_results.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding) # Add a layout to hold all the results from the Analysis. self.layout_results = QVBoxLayout() self.layout_results.setContentsMargins(0, 0, 0, 0) grp_results.setLayout(self.layout_results) return grp_results def build_coord_input(self): """Build the necessary elements for the user input text box. Returns: (QVBoxLayout): The base layout of the coord input. """ layout_base = QVBoxLayout() layout_base.setContentsMargins(0, 0, 0, 0) self.raw_input = QPlainTextEdit() self.raw_input.setMaximumHeight(75) self.raw_input.setPlaceholderText('example: "0 292 399 307"') if self.user_input: self.raw_input.setPlainText(self.user_input) self.layout_input = QVBoxLayout() self.layout_input.setContentsMargins(0, 0, 0, 0) layout_base.addWidget(self.raw_input) layout_base.addLayout(self.layout_input) return layout_base def ctr_debug(self): """CONTROL: Called by the 'debug' button to randomly assign colors to the zones.""" for result in self.results: rand_color = [random.randint(0, 255) for i in range(3)] rand_color = QColor(*rand_color) brush = QBrush(rand_color, Qt.Dense5Pattern) result.rectangle.setBrush(brush) result.base = brush def ctl_add_input(self): """CONTROL: Called by the 'Add' button to handle parsing the raw input.""" raw_data = self.raw_input.toPlainText() if raw_data: parsed_input = utils.format_raw_input(raw_data) for user_input in parsed_input: start_coord = land.Coord(user_input[0], user_input[1]) end_coord = land.Coord(user_input[2], user_input[3]) zone = land.Zone(start_coord, end_coord) self.field.add_zone(zone, barren=True) self.draw_zone(zone, barren=True) def ctl_run(self): """CONTROL: Called by the 'Run' button to initialize the final analysis.""" # Make sure some values are cleared ahead of the analysis. self.clear_results() for rectangle in self.fertile_zones: self.scene.removeItem(rectangle) self.field.fertile_zones = set() # Run the analysis self.field.check_zones() total_area = 0 # Format and draw the results in the QGraphicsScene. for i, island in enumerate(self.field.islands): for zone in island: rectangle = self.draw_zone(zone) size = zone.get_size() if i == 0: # Only add (largest) island to results. Always at 0. total_area += size self.results.add( ResultLabel(label=str(size), rectangle=rectangle, zone=zone)) else: # This zone is fertile, but Inaccessible. rectangle.setBrush(self.brush_fertile_no_go) self.canvas.update() # Pain the update for the user to see the new zone. self.app.processEvents() # Dont show all at the same time. (For more pleasing visual) time.sleep(.015) # Print islands, smallest to largest. print("Islands", self.field.islands_as_area()) island_areas = " ".join(str(i) for i in self.field.islands_as_area()) self.lbl_island.setText(f"Island Areas: {island_areas}") # Set the label as the total area as the largest island. self.lbl_area.setText(f"Total Area: {total_area}") # Sort the results by their zones area. for result in sorted(self.results, key=lambda x: x.zone.get_size()): # Add the result to the results layout in the UI self.layout_results.addWidget(result) def ctl_clear_zones(self): """CONTROL: Called by the 'Reset' button to handle resetting the data and graphics.""" for rectangle in self.barren_zones + self.fertile_zones: self.scene.removeItem(rectangle) self.field.fertile_zones = set() self.fertile_zones = list() self.field.barren_zones = set() self.barren_zones = list() self.clear_results() def clear_results(self): """Clears the results set in preparation for incoming new results.""" for result in self.results: result.deleteLater() self.results = set() self.lbl_area.setText("Total Area:") self.lbl_island.setText("Island Areas:") def draw_zone(self, zone, barren=False): """Creates a QGraphicsRecItem from a Zone and adds it to the scene. Args: zone (Zone): The zone to generate the rectangle from. barren (bool): barren colors if True else fertile color. Returns: (QGraphicsRecItem): The drown rectangle that is added to the canvas. """ # Build the QGraphicsRecItem from the zone data. rectangle = self.scene.addRect(zone.start.x, zone.start.y, zone.end.x - zone.start.x + 1, zone.end.y - zone.start.y + 1) if barren: self.barren_zones.append(rectangle) brush = self.brush_barren else: self.fertile_zones.append(rectangle) brush = self.brush_fertile # Apply the color to the rectangle item. rectangle.setBrush(brush) # Remove the default border rectangle.setPen(Qt.NoPen) return rectangle
class ImageView(QGraphicsView): # Tell PageWidget that a file is dropped onto view. dropped_relay = Signal(QDropEvent) def __init__(self, image_path): super(ImageView, self).__init__(None) self.scene = QGraphicsScene() self.setScene(self.scene) self.pixmapitem = self.scene.addPixmap(QPixmap.fromImage(QImage(image_path))) self.last_release_time = 0 self.watcher = QFileSystemWatcher() self.watcher.fileChanged.connect(self.refresh_image) # Register file watcher self.watcher.addPath(image_path) def dragEnterEvent(self, drag_enter_event): # QDragEnterEvent if drag_enter_event.mimeData().hasUrls(): drag_enter_event.acceptProposedAction() # https://stackoverflow.com/a/4421835/4112667 def dragMoveEvent(self, event): pass def dropEvent(self, drop_event): # QDropEvent self.dropped_relay.emit(drop_event) ''' When overwriting an image file, I guess Windows will delete it and then create a new file with the same name. So this function will be called twice. The first round is triggered by deleting. In this case, the image file doesn't exist, so QImage and QPixmap are all invalid and as a result, the view will become white background. Only after the image being created and the function is called for the second time, will the view show the image normally. The User will notice a white flicker because of two rounds of callings. To resolve this problem, we need to detect the invalid QImage or QPixmap and skip the unintended round. ''' def refresh_image(self, image_path): qimage = QImage(image_path) if qimage.isNull(): return pixmap = QPixmap.fromImage(qimage) self.scene.removeItem(self.pixmapitem) self.pixmapitem = self.scene.addPixmap(pixmap) # This will make scrollbar fit the image self.setSceneRect(QRectF(pixmap.rect())) def mousePressEvent(self, mouse_event): # QMouseEvent if mouse_event.button() == Qt.LeftButton: self.setDragMode(QGraphicsView.ScrollHandDrag) elif mouse_event.button() == Qt.RightButton: self.setDragMode(QGraphicsView.RubberBandDrag) QGraphicsView.mousePressEvent(self, mouse_event) def mouseReleaseEvent(self, mouse_event): # QMouseEvent QGraphicsView.mouseReleaseEvent(self, mouse_event) if mouse_event.button() == Qt.LeftButton: self.setDragMode(QGraphicsView.NoDrag) elif mouse_event.button() == Qt.RightButton: self.setDragMode(QGraphicsView.NoDrag) now = time.time() delta = now - self.last_release_time self.last_release_time = now if delta < 0.3: # fast double click self.resetTransform() # Reset to original size (reset scale matrix) return # Maybe a selection selection = self.scene.selectionArea().boundingRect() self.scene.setSelectionArea(QPainterPath()) if selection.isValid(): self.fitInView(selection, Qt.KeepAspectRatio) def wheelEvent(self, wheel_event): # QWheelEvent num_degrees = wheel_event.angleDelta().y() / 8 num_steps = num_degrees / 15 coefficient = 1 + (num_steps * 0.25) self.scale(coefficient, coefficient)
class QtImageViewer(QGraphicsView): """ PyQt image viewer widget for a QPixmap in a QGraphicsView scene with mouse zooming and panning. Displays a QImage or QPixmap (QImage is internally converted to a QPixmap). To display any other image format, you must first convert it to a QImage or QPixmap. Some useful image format conversion utilities: qimage2ndarray: NumPy ndarray <==> QImage (https://github.com/hmeine/qimage2ndarray) ImageQt: PIL Image <==> QImage (https://github.com/python-pillow/Pillow/blob/master/PIL/ImageQt.py) Mouse interaction: Left mouse button drag: Pan image. Right mouse button drag: Zoom box. Right mouse button doubleclick: Zoom to show entire image. """ # Mouse button signals emit image scene (x, y) coordinates. # !!! For image (row, column) matrix indexing, row = y and column = x. leftMouseButtonPressed = Signal(float, float) rightMouseButtonPressed = Signal(float, float) leftMouseButtonReleased = Signal(float, float) rightMouseButtonReleased = Signal(float, float) leftMouseButtonDoubleClicked = Signal(float, float) rightMouseButtonDoubleClicked = Signal(float, float) def __init__(self): QGraphicsView.__init__(self) # Image is displayed as a QPixmap in a QGraphicsScene attached to this QGraphicsView. self.scene = QGraphicsScene() self.setScene(self.scene) # Store a local handle to the scene's current image pixmap. self._pixmapHandle = None # Image aspect ratio mode. # !!! ONLY applies to full image. Aspect ratio is always ignored when zooming. # Qt.IgnoreAspectRatio: Scale image to fit viewport. # Qt.KeepAspectRatio: Scale image to fit inside viewport, preserving aspect ratio. # Qt.KeepAspectRatioByExpanding: Scale image to fill the viewport, preserving aspect ratio. self.aspectRatioMode = Qt.KeepAspectRatio # Scroll bar behaviour. # Qt.ScrollBarAlwaysOff: Never shows a scroll bar. # Qt.ScrollBarAlwaysOn: Always shows a scroll bar. # Qt.ScrollBarAsNeeded: Shows a scroll bar only when zoomed. self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn) self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn) self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse) # Stack of QRectF zoom boxes in scene coordinates. self.zoomStack = [] # Flags for enabling/disabling mouse interaction. self.canZoom = True self.canPan = True def hasImage(self): """ Returns whether or not the scene contains an image pixmap. """ return self._pixmapHandle is not None def clearImage(self): """ Removes the current image pixmap from the scene if it exists. """ if self.hasImage(): self.scene.removeItem(self._pixmapHandle) self._pixmapHandle = None def pixmap(self): """ Returns the scene's current image pixmap as a QPixmap, or else None if no image exists. :rtype: QPixmap | None """ if self.hasImage(): return self._pixmapHandle.pixmap() return None def image(self): """ Returns the scene's current image pixmap as a QImage, or else None if no image exists. :rtype: QImage | None """ if self.hasImage(): return self._pixmapHandle.pixmap().toImage() return None def setImage(self, image): """ Set the scene's current image pixmap to the input QImage or QPixmap. Raises a RuntimeError if the input image has type other than QImage or QPixmap. :type image: QImage | QPixmap """ if type(image) is QPixmap: pixmap = image elif type(image) is QImage: pixmap = QPixmap.fromImage(image) else: raise RuntimeError( "ImageViewer.setImage: Argument must be a QImage or QPixmap." ) if self.hasImage(): self._pixmapHandle.setPixmap(pixmap) else: self._pixmapHandle = self.scene.addPixmap(pixmap) self.setSceneRect(QRectF( pixmap.rect())) # Set scene size to image size. self.updateViewer() def loadImageFromFile(self, fileName): """ Load an image from file. Without any arguments, loadImageFromFile() will popup a file dialog to choose the image file. With a fileName argument, loadImageFromFile(fileName) will attempt to load the specified image file directly. """ if len(fileName) and os.path.isfile(fileName): image = QImage(fileName) self.setImage(image) def updateViewer(self): """ Show current zoom (if showing entire image, apply current aspect ratio mode). """ if not self.hasImage(): return if len(self.zoomStack) and self.sceneRect().contains( self.zoomStack[-1]): self.fitInView(self.zoomStack[-1], Qt.IgnoreAspectRatio ) # Show zoomed rect (ignore aspect ratio). else: self.zoomStack = [ ] # Clear the zoom stack (in case we got here because of an invalid zoom). self.fitInView( self.sceneRect(), self.aspectRatioMode ) # Show entire image (use current aspect ratio mode). def resizeEvent(self, event): """ Maintain current zoom on resize. """ self.updateViewer() def wheelEvent(self, event): scale_factor = 1.1 if event.delta() > 0: self.scale(scale_factor, scale_factor) else: self.scale(1 / scale_factor, 1 / scale_factor) def mousePressEvent(self, event): """ Start mouse pan or zoom mode. """ scenePos = self.mapToScene(event.pos()) if event.button() == Qt.LeftButton: if self.canPan: self.setDragMode(QGraphicsView.ScrollHandDrag) self.leftMouseButtonPressed.emit(scenePos.x(), scenePos.y()) QGraphicsView.mousePressEvent(self, event) def mouseReleaseEvent(self, event): """ Stop mouse pan or zoom mode (apply zoom if valid). """ QGraphicsView.mouseReleaseEvent(self, event) scenePos = self.mapToScene(event.pos()) if event.button() == Qt.LeftButton: self.setDragMode(QGraphicsView.NoDrag) self.leftMouseButtonReleased.emit(scenePos.x(), scenePos.y()) def mouseDoubleClickEvent(self, event): """ Show entire image. """ scenePos = self.mapToScene(event.pos()) if event.button() == Qt.LeftButton: self.leftMouseButtonDoubleClicked.emit(scenePos.x(), scenePos.y()) elif event.button() == Qt.RightButton: if self.canZoom: self.zoomStack = [] # Clear zoom stack. self.updateViewer() self.rightMouseButtonDoubleClicked.emit( scenePos.x(), scenePos.y()) QGraphicsView.mouseDoubleClickEvent(self, event)