diff --git a/mal_gui/app.py b/mal_gui/app.py index 312773d..358dd3b 100644 --- a/mal_gui/app.py +++ b/mal_gui/app.py @@ -24,10 +24,11 @@ QPushButton, QDialogButtonBox, QFileDialog, - QMessageBox + QMessageBox, ) from .main_window import MainWindow + class FileSelectionDialog(QDialog): def __init__(self, parent=None): super().__init__(parent) @@ -46,7 +47,7 @@ def __init__(self, parent=None): # Load the config file containing latest lang file path config_file_dir = user_config_dir("mal-gui", "mal-lang") - self.config_file_path = config_file_dir + '/config.ini' + self.config_file_path = config_file_dir + "/config.ini" # Make sure config file exists os.makedirs(os.path.dirname(self.config_file_path), exist_ok=True) @@ -54,7 +55,8 @@ def __init__(self, parent=None): self.config = configparser.ConfigParser() self.config.read(self.config_file_path) self.selected_lang_file = self.config.get( - 'Settings', 'langFilePath', fallback=None) + "Settings", "langFilePath", fallback=None + ) print(f"Initial langFilePath path: {self.selected_lang_file}") self.lang_file_path_text.setText(self.selected_lang_file) @@ -98,27 +100,25 @@ def save_lang_file_path(self): selected_lang_file = self.lang_file_path_text.text() - if selected_lang_file.endswith('.mar') or \ - selected_lang_file.endswith('.mal'): + if selected_lang_file.endswith(".mar") or selected_lang_file.endswith(".mal"): self.selected_lang_file = selected_lang_file # Remember language choice in user settings try: - self.config.add_section('Settings') + self.config.add_section("Settings") except configparser.DuplicateSectionError: pass - self.config.set('Settings', 'langFilePath', - self.selected_lang_file) + self.config.set("Settings", "langFilePath", self.selected_lang_file) - with open(self.config_file_path, 'w', encoding='utf-8') as conf_file: + with open(self.config_file_path, "w", encoding="utf-8") as conf_file: self.config.write(conf_file) self.accept() # Close the dialog and return accepted else: QMessageBox.warning( - self, "Invalid File", - "Please select a valid .mal or .mar file.") + self, "Invalid File", "Please select a valid .mal or .mar file." + ) def get_selected_file(self): return self.selected_lang_file diff --git a/mal_gui/assets_container/__init__.py b/mal_gui/assets_container/__init__.py index 3d3f596..151dfc3 100644 --- a/mal_gui/assets_container/__init__.py +++ b/mal_gui/assets_container/__init__.py @@ -1,7 +1,4 @@ from .assets_container_rectangle_box import AssetsContainerRectangleBox from .assets_container import AssetsContainer -__all__ = [ - "AssetsContainerRectangleBox", - "AssetsContainer" -] +__all__ = ["AssetsContainerRectangleBox", "AssetsContainer"] diff --git a/mal_gui/assets_container/assets_container.py b/mal_gui/assets_container/assets_container.py index 0b0b73f..666fa74 100644 --- a/mal_gui/assets_container/assets_container.py +++ b/mal_gui/assets_container/assets_container.py @@ -1,4 +1,4 @@ -from PySide6.QtCore import QRectF, Qt,QPointF,QSize,QSizeF,QTimer +from PySide6.QtCore import QRectF, Qt, QPointF, QSize, QSizeF, QTimer from PySide6.QtGui import ( QPixmap, QFont, @@ -8,51 +8,52 @@ QPainterPath, QFontMetrics, QLinearGradient, - QImage + QImage, ) -from PySide6.QtWidgets import QGraphicsItem +from PySide6.QtWidgets import QGraphicsItem from ..object_explorer.editable_text_item import EditableTextItem from .assets_container_rectangle_box import AssetsContainerRectangleBox + class AssetsContainer(QGraphicsItem): # Starting Sequence Id with normal start at 100 (randomly taken) container_sequence_id = 100 def __init__( - self, - container_type, - container_name, - image_path, - plus_symbol_image_path, - minus_symbol_image_path, - parent=None - ): + self, + container_type, + container_name, + image_path, + plus_symbol_image_path, + minus_symbol_image_path, + parent=None, + ): super().__init__(parent) self.setZValue(1) # rect items are on top self.container_type = container_type self.container_name = container_name - self.container_sequence_id = \ - AssetsContainer.generate_next_sequence_id() + self.container_sequence_id = AssetsContainer.generate_next_sequence_id() self.image_path = image_path self.plus_symbol_image_path = plus_symbol_image_path self.minus_symbol_image_path = minus_symbol_image_path self.plus_or_minus_image_rect = QRectF() self.is_plus_symbol_visible = True # Track the current symbol state self.container_box = None - print("image path = "+ self.image_path) + print("image path = " + self.image_path) - self.image = self.load_image_with_quality( - self.image_path, QSize(512, 512)) + self.image = self.load_image_with_quality(self.image_path, QSize(512, 512)) self.plus_symbol_image = self.load_image_with_quality( - self.plus_symbol_image_path, QSize(512, 512)) + self.plus_symbol_image_path, QSize(512, 512) + ) self.minus_symbol_image = self.load_image_with_quality( - self.minus_symbol_image_path, QSize(512, 512)) + self.minus_symbol_image_path, QSize(512, 512) + ) self.setFlags( - QGraphicsItem.ItemIsSelectable | - QGraphicsItem.ItemIsMovable | - QGraphicsItem.ItemSendsGeometryChanges + QGraphicsItem.ItemIsSelectable + | QGraphicsItem.ItemIsMovable + | QGraphicsItem.ItemSendsGeometryChanges ) # Create the editable text item for block type @@ -67,8 +68,8 @@ def __init__( self.height = 70 self.size = QRectF(-self.width / 2, -self.height / 2, self.width, self.height) - self.container_type_bg_color = QColor(0, 200, 255) #Blue - self.container_name_bg_color = QColor(20, 20, 20, 200) # Gray + self.container_type_bg_color = QColor(0, 200, 255) # Blue + self.container_name_bg_color = QColor(20, 20, 20, 200) # Gray self.icon_path = None self.icon_visible = True @@ -80,12 +81,12 @@ def __init__( self.horizontal_margin = 15 # Horizontal margin self.vertical_margin = 15 # Vertical margin - + self.timer = QTimer() - self.status_color = QColor(0, 255, 0) + self.status_color = QColor(0, 255, 0) self.attacker_toggle_state = False self.timer.timeout.connect(self.update_status_color) - #timer to trigger every 500ms (0.5 seconds) + # timer to trigger every 500ms (0.5 seconds) self.timer.start(500) self.build() @@ -97,8 +98,7 @@ def boundingRect(self): def mouseDoubleClickEvent(self, event): """Overrides base method""" if event.button() == Qt.LeftButton: - self.type_text_item.setTextInteractionFlags( - Qt.TextEditorInteraction) + self.type_text_item.setTextInteractionFlags(Qt.TextEditorInteraction) self.type_text_item.setFocus() # Select all text when activated self.type_text_item.select_all_text() @@ -121,7 +121,9 @@ def mousePressEvent(self, event): # self.is_plus_symbol_visible = not self.is_plus_symbol_visible # self.update() self.toggle_container_expansion() - elif self.type_text_item.hasFocus() and not self.type_text_item.contains(event.pos()): + elif self.type_text_item.hasFocus() and not self.type_text_item.contains( + event.pos() + ): self.type_text_item.clearFocus() elif not self.type_text_item.contains(event.pos()): self.type_text_item.deselect_text() @@ -139,7 +141,7 @@ def setIcon(self, icon_path=None): def itemChange(self, change, value): """Overrides base method""" if change == QGraphicsItem.ItemPositionChange and self.scene(): - if hasattr(self, 'item_moved') and callable(self.item_moved): + if hasattr(self, "item_moved") and callable(self.item_moved): self.item_moved() return super().itemChange(change, value) @@ -178,21 +180,23 @@ def paint(self, painter, option, widget=None): # Qt.KeepAspectRatio, Qt.SmoothTransformation) resized_image_icon = self.image - # Calculate the position and size for the icon background icon_rect = QRectF( -self.width / 2 + 10, -self.height / 2 + 10, target_icon_size.width(), - target_icon_size.height() + target_icon_size.height(), ) margin = 5 # Margin around the icon # Draw the background for the icon with additional margin background_rect = QRectF( icon_rect.topLeft() - QPointF(margin, margin), - QSizeF(target_icon_size.width() + 2 * margin, - target_icon_size.height() + 2 * margin)) + QSizeF( + target_icon_size.width() + 2 * margin, + target_icon_size.height() + 2 * margin, + ), + ) painter.setBrush(Qt.white) # Set the brush color to white @@ -203,9 +207,11 @@ def paint(self, painter, option, widget=None): painter.drawPixmap(icon_rect.toRect(), resized_image_icon) # Determine which symbol to draw based on the current state - current_symbol_image = self.plus_symbol_image\ - if self.is_plus_symbol_visible\ - else self.minus_symbol_image + current_symbol_image = ( + self.plus_symbol_image + if self.is_plus_symbol_visible + else self.minus_symbol_image + ) if not current_symbol_image.isNull(): # Desired size for the second icon target_symbol_image_size = QSize(12, 12) @@ -216,28 +222,23 @@ def paint(self, painter, option, widget=None): # Calculate the position for the symbol at the # bottom-right corner of title_bg_path self.plus_or_minus_image_rect = QRectF( - title_bg_rect.right() - - target_symbol_image_size.width() - 10, - title_bg_rect.bottom() - - target_symbol_image_size.height() - 5, + title_bg_rect.right() - target_symbol_image_size.width() - 10, + title_bg_rect.bottom() - target_symbol_image_size.height() - 5, target_symbol_image_size.width(), - target_symbol_image_size.height() + target_symbol_image_size.height(), ) painter.setBrush(Qt.white) painter.drawRect(self.plus_or_minus_image_rect) resized_symbol_image = current_symbol_image.scaled( - target_symbol_image_size, - Qt.KeepAspectRatio, - Qt.SmoothTransformation + target_symbol_image_size, Qt.KeepAspectRatio, Qt.SmoothTransformation ) # Draw the plus or minus symbol with a white background painter.drawPixmap( - self.plus_or_minus_image_rect.toRect(), resized_symbol_image) - - + self.plus_or_minus_image_rect.toRect(), resized_symbol_image + ) # Draw the highlight if selected if self.isSelected(): @@ -267,12 +268,7 @@ def build(self): # Draw the background of the node self.path = QPainterPath() self.path.addRoundedRect( - -fixed_width / 2, - -fixed_height / 2, - fixed_width, - fixed_height, - 6, - 6 + -fixed_width / 2, -fixed_height / 2, fixed_width, fixed_height, 6, 6 ) self.title_bg_path = QPainterPath() @@ -282,18 +278,13 @@ def build(self): fixed_width, title_font.pointSize() + 2 * self.vertical_margin, 6, - 6 + 6, ) # Draw status path self.status_path.setFillRule(Qt.WindingFill) self.status_path.addRoundedRect( - fixed_width / 2 - 12, - -fixed_height / 2 + 2, - 10, - 10, - 2, - 2 + fixed_width / 2 - 12, -fixed_height / 2 + 2, 10, 10, 2, 2 ) # Center title in the upper half @@ -302,10 +293,9 @@ def build(self): # Center horizontally -title_font_metrics.horizontalAdvance(self.title_text) / 2, # Center vertically within its section - -fixed_height / 2 + self.vertical_margin - + title_font_metrics.ascent(), + -fixed_height / 2 + self.vertical_margin + title_font_metrics.ascent(), title_font, - self.title_text + self.title_text, ) # Set the font and default color for type_text_item @@ -316,9 +306,7 @@ def build(self): self.update_type_text_item_position() # Connect lostFocus signal to update position when text loses focus - self.type_text_item.lostFocus.connect( - self.update_type_text_item_position - ) + self.type_text_item.lostFocus.connect(self.update_type_text_item_position) # self.widget.move(-self.widget.size().width() / 2, # fixed_height / 2 - self.widget.size().height() + 5) @@ -334,14 +322,10 @@ def update_type_text_item_position(self): # Calculate the new position for type_text_item type_text_item_pos_x = ( - -type_font_metrics.horizontalAdvance( - self.type_text_item.toPlainText() - ) / 2 + -type_font_metrics.horizontalAdvance(self.type_text_item.toPlainText()) / 2 ) type_text_item_pos_y = ( - -fixed_height / 2 - + title_font_metrics.height() - + 2 * self.vertical_margin + -fixed_height / 2 + title_font_metrics.height() + 2 * self.vertical_margin ) # Update position @@ -361,7 +345,7 @@ def get_item_attribute_vakues(self): return { "Container Sequence ID": self.container_sequence_id, "Container Name": self.container_name, - "Container Type": self.container_type + "Container Type": self.container_type, } def toggle_icon_visibility(self): @@ -369,18 +353,14 @@ def toggle_icon_visibility(self): self.update() def update_status_color(self): - self.status_color = QColor(0, 255, 0) + self.status_color = QColor(0, 255, 0) self.update() def load_image_with_quality(self, path, size): image = QImage(path) if not image.isNull(): return QPixmap.fromImage( - image.scaled( - size, - Qt.KeepAspectRatio, - Qt.SmoothTransformation - ) + image.scaled(size, Qt.KeepAspectRatio, Qt.SmoothTransformation) ) return QPixmap() @@ -409,12 +389,8 @@ def hide_container_box(self): def update_container_box_position(self): if self.container_box: - container_bottom_left = ( - self.pos() + QPointF( - -self.width / 2, - self.boundingRect().height() / 2 - ) + container_bottom_left = self.pos() + QPointF( + -self.width / 2, self.boundingRect().height() / 2 ) - container_box_position = \ - container_bottom_left + QPointF(0, self.height / 2) + container_box_position = container_bottom_left + QPointF(0, self.height / 2) self.container_box.setPos(container_box_position) diff --git a/mal_gui/assets_container/assets_container_rectangle_box.py b/mal_gui/assets_container/assets_container_rectangle_box.py index 90ad748..f6ca775 100644 --- a/mal_gui/assets_container/assets_container_rectangle_box.py +++ b/mal_gui/assets_container/assets_container_rectangle_box.py @@ -1,9 +1,10 @@ from PySide6.QtWidgets import QGraphicsRectItem from PySide6.QtGui import QBrush, QColor, QPen + class AssetsContainerRectangleBox(QGraphicsRectItem): def __init__(self, rect, parent=None): super().__init__(rect, parent) - self.setBrush(QBrush(QColor(0, 0, 255, 50))) #Blue - self.setPen(QPen(QColor(0, 0, 255, 50))) #Blue + self.setBrush(QBrush(QColor(0, 0, 255, 50))) # Blue + self.setPen(QPen(QColor(0, 0, 255, 50))) # Blue diff --git a/mal_gui/association_table_view.py b/mal_gui/association_table_view.py index c75dd49..b0a3f51 100644 --- a/mal_gui/association_table_view.py +++ b/mal_gui/association_table_view.py @@ -1,7 +1,9 @@ -from PySide6.QtWidgets import QWidget,QTableView,QVBoxLayout -from PySide6.QtGui import QStandardItemModel,QStandardItem +from PySide6.QtWidgets import QWidget, QTableView, QVBoxLayout +from PySide6.QtGui import QStandardItemModel, QStandardItem from .main_window import MainWindow + + class AssociationDefinitions(QWidget): def __init__(self, parent: MainWindow): super().__init__(parent) @@ -12,10 +14,15 @@ def __init__(self, parent: MainWindow): self.table_association_view = QTableView(self) self.association_info_model = QStandardItemModel() - #headers for the columns + # headers for the columns self.association_info_model.setHorizontalHeaderLabels( - ['AssocLeftAsset', 'AssocLeftField', 'AssocName', - 'AssocRightField','AssocRightAsset'] + [ + "AssocLeftAsset", + "AssocLeftField", + "AssocName", + "AssocRightField", + "AssocRightAsset", + ] ) self.association_info_model.removeRows( @@ -28,7 +35,7 @@ def __init__(self, parent: MainWindow): QStandardItem(assoc.left_field.fieldname), QStandardItem(assoc.name), QStandardItem(assoc.right_field.fieldname), - QStandardItem(assoc.right_field.asset.name) + QStandardItem(assoc.right_field.asset.name), ] self.association_info_model.appendRow(items) diff --git a/mal_gui/connection_dialog.py b/mal_gui/connection_dialog.py index 8d5a60e..98a4560 100644 --- a/mal_gui/connection_dialog.py +++ b/mal_gui/connection_dialog.py @@ -15,9 +15,10 @@ if TYPE_CHECKING: from .object_explorer import AssetItem, AttackerItem - from maltoolbox.language import LanguageGraph, LanguageGraphAsset + from maltoolbox.language import LanguageGraph from maltoolbox.model import Model, ModelAsset + class ConnectionDialog(QDialog): def filter_items(self, text): pass @@ -28,13 +29,13 @@ def ok_button_clicked(self): class AssociationConnectionDialog(ConnectionDialog): def __init__( - self, - start_item: AssetItem, - end_item: AssetItem, - lang_graph: LanguageGraph, - model: Model, - parent=None - ): + self, + start_item: AssetItem, + end_item: AssetItem, + lang_graph: LanguageGraph, + model: Model, + parent=None, + ): super().__init__(parent) self.lang_graph: LanguageGraph = lang_graph @@ -43,8 +44,8 @@ def __init__( self.setWindowTitle("Select Association Type") self.setMinimumWidth(300) - print(f'START ITEM TYPE {start_item.asset_type}') - print(f'END ITEM TYPE {end_item.asset_type}') + print(f"START ITEM TYPE {start_item.asset_type}") + print(f"END ITEM TYPE {end_item.asset_type}") self.association_list_widget = QListWidget() @@ -53,9 +54,7 @@ def __init__( self.field_name = None self.layout = QVBoxLayout() - self.label = ( - QLabel(f"{self.start_asset.name} -> {self.end_asset.name}") - ) + self.label = QLabel(f"{self.start_asset.name} -> {self.end_asset.name}") self.layout.addWidget(self.label) self.filter_edit = QLineEdit() self.filter_edit.setPlaceholderText("Type to filter...") @@ -69,14 +68,20 @@ def __init__( # If assoc ends with end_assets type, give that assoc # as option in list widget if field.asset == self.end_asset.lg_asset: - assoc_list_item = QListWidgetItem(self.start_asset.name + "." + fieldname + " = " + self.end_asset.name) + assoc_list_item = QListWidgetItem( + self.start_asset.name + + "." + + fieldname + + " = " + + self.end_asset.name + ) assoc_list_item.setData( Qt.UserRole, { - 'from': self.start_asset, - 'to': self.end_asset, - 'fieldname': fieldname - } + "from": self.start_asset, + "to": self.end_asset, + "fieldname": fieldname, + }, ) self.association_list_widget.addItem(assoc_list_item) @@ -111,24 +116,24 @@ def ok_button_clicked(self): if selected_item: data = selected_item.data(Qt.UserRole) - from_asset: ModelAsset = data.get('from') - to_asset: ModelAsset = data.get('to') - self.field_name: str = data.get('fieldname') - - print(f'{from_asset}.{self.field_name} = {to_asset} chosen') + from_asset: ModelAsset = data.get("from") + to_asset: ModelAsset = data.get("to") + self.field_name: str = data.get("fieldname") + print(f"{from_asset}.{self.field_name} = {to_asset} chosen") self.accept() + class EntrypointConnectionDialog(ConnectionDialog): def __init__( - self, - attacker_item: AttackerItem, - asset_item: AssetItem, - lang_graph: LanguageGraph, - model, - parent=None - ): + self, + attacker_item: AttackerItem, + asset_item: AssetItem, + lang_graph: LanguageGraph, + model, + parent=None, + ): super().__init__(parent) self.lang_graph = lang_graph @@ -146,7 +151,7 @@ def __init__( already_attached_entrypoints = set(attacker_item.entry_points) for attack_step in asset_type.attack_steps.values(): - if attack_step.type not in ['or', 'and']: + if attack_step.type not in ["or", "and"]: continue attack_step_full_name = attack_step.asset.name + ":" + attack_step.name if attack_step_full_name not in already_attached_entrypoints: @@ -155,9 +160,7 @@ def __init__( self.layout = QVBoxLayout() - self.label = QLabel( - f"{attacker_item.name}:{asset_item.asset.name}" - ) + self.label = QLabel(f"{attacker_item.name}:{asset_item.asset.name}") self.layout.addWidget(self.label) self.filter_edit = QLineEdit() @@ -174,8 +177,7 @@ def __init__( self.cancel_button = QPushButton("Cancel") self.cancel_button.clicked.connect(self.reject) - self.cancel_button.setSizePolicy( - QSizePolicy.Expanding, QSizePolicy.Fixed) + self.cancel_button.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) button_layout.addWidget(self.cancel_button) self.layout.addLayout(button_layout) @@ -195,13 +197,13 @@ def ok_button_clicked(self): class GoalConnectionDialog(ConnectionDialog): def __init__( - self, - attacker_item: AttackerItem, - asset_item: AssetItem, - lang_graph: LanguageGraph, - model, - parent=None - ): + self, + attacker_item: AttackerItem, + asset_item: AssetItem, + lang_graph: LanguageGraph, + model, + parent=None, + ): super().__init__(parent) self.lang_graph = lang_graph @@ -219,7 +221,7 @@ def __init__( already_attached_goals = set(attacker_item.goals) for attack_step in asset_type.attack_steps.values(): - if attack_step.type not in ['or', 'and']: + if attack_step.type not in ["or", "and"]: continue attack_step_full_name = attack_step.asset.name + ":" + attack_step.name if attack_step_full_name not in already_attached_goals: @@ -229,9 +231,7 @@ def __init__( self.layout = QVBoxLayout() - self.label = QLabel( - f"{attacker_item.name}:{asset_item.asset.name}" - ) + self.label = QLabel(f"{attacker_item.name}:{asset_item.asset.name}") self.layout.addWidget(self.label) self.filter_edit = QLineEdit() @@ -248,8 +248,7 @@ def __init__( self.cancel_button = QPushButton("Cancel") self.cancel_button.clicked.connect(self.reject) - self.cancel_button.setSizePolicy( - QSizePolicy.Expanding, QSizePolicy.Fixed) + self.cancel_button.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed) button_layout.addWidget(self.cancel_button) self.layout.addLayout(button_layout) diff --git a/mal_gui/connection_item.py b/mal_gui/connection_item.py index 75360e8..1e43814 100644 --- a/mal_gui/connection_item.py +++ b/mal_gui/connection_item.py @@ -1,13 +1,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from PySide6.QtCore import Qt, QPointF, QLineF -from PySide6.QtGui import QBrush, QColor,QPen +from PySide6.QtGui import QBrush, QColor, QPen from PySide6.QtWidgets import ( QGraphicsLineItem, QGraphicsTextItem, QGraphicsRectItem, - QGraphicsItemGroup + QGraphicsItemGroup, ) if TYPE_CHECKING: @@ -15,8 +15,10 @@ from .model_scene import ModelScene from .object_explorer import AssetItem, AttackerItem + class IConnectionItem(QGraphicsLineItem): """'interface' for Connection Item""" + start_item: AssetItem end_item: AssetItem association_details: list[str] @@ -44,7 +46,7 @@ def __init__( start_item: AssetItem, end_item: AssetItem, scene: ModelScene, - parent = None + parent=None, ): super().__init__(parent) @@ -61,9 +63,9 @@ def __init__( self.end_item.add_connection(self) # Fetch the association and the fieldnames - lg_assoc: LanguageGraphAssociation = ( - start_item.asset.lg_asset.associations[fieldname] - ) + lg_assoc: LanguageGraphAssociation = start_item.asset.lg_asset.associations[ + fieldname + ] opposite_fieldname = lg_assoc.get_opposite_fieldname(fieldname) # Get left field name @@ -74,14 +76,11 @@ def __init__( self.right_fieldname = fieldname # Create labels with background color - self.label_assoc_left_field = \ - self.create_label(self.left_fieldname) + self.label_assoc_left_field = self.create_label(self.left_fieldname) - self.label_assoc_middle_name = \ - self.create_label(self.assoc_name) + self.label_assoc_middle_name = self.create_label(self.assoc_name) - self.label_assoc_right_field = \ - self.create_label(self.right_fieldname) + self.label_assoc_right_field = self.create_label(self.right_fieldname) self.update_path() @@ -115,32 +114,37 @@ def update_path(self): label_assoc_left_field_pos = self.line().pointAt(0.2) self.label_assoc_left_field.setPos( - label_assoc_left_field_pos - QPointF( + label_assoc_left_field_pos + - QPointF( self.label_assoc_left_field.boundingRect().width() / 2, - self.label_assoc_left_field.boundingRect().height() / 2 + self.label_assoc_left_field.boundingRect().height() / 2, ) ) label_assoc_middle_name_pos = self.line().pointAt(0.5) self.label_assoc_middle_name.setPos( - label_assoc_middle_name_pos - QPointF( + label_assoc_middle_name_pos + - QPointF( self.label_assoc_middle_name.boundingRect().width() / 2, - self.label_assoc_middle_name.boundingRect().height() / 2 + self.label_assoc_middle_name.boundingRect().height() / 2, ) ) labelassoc_right_field_pos = self.line().pointAt(0.8) self.label_assoc_right_field.setPos( - labelassoc_right_field_pos - QPointF( + labelassoc_right_field_pos + - QPointF( self.label_assoc_right_field.boundingRect().width() / 2, - self.label_assoc_right_field.boundingRect().height() / 2 + self.label_assoc_right_field.boundingRect().height() / 2, ) ) self.label_assoc_left_field.setVisible( - self._scene.get_show_assoc_checkbox_status()) + self._scene.get_show_assoc_checkbox_status() + ) self.label_assoc_right_field.setVisible( - self._scene.get_show_assoc_checkbox_status()) + self._scene.get_show_assoc_checkbox_status() + ) def calculate_offset(self, rect, label_pos, angle): """Calculate the offset to position the label @@ -182,6 +186,7 @@ def delete(self): class AttackerConnectionBase(IConnectionItem): """For both entrypoints and goals""" + COLOR = QColor(0, 0, 0) LINE_STYLE = Qt.SolidLine ICON_TEXT = "?" @@ -192,7 +197,7 @@ def __init__( attacker_item: AttackerItem, asset_item: AssetItem, scene: ModelScene, - parent=None + parent=None, ): super().__init__(parent) @@ -213,7 +218,6 @@ def __init__( self.update_path() - def create_label(self, text) -> QGraphicsItemGroup: label = QGraphicsTextItem(self.ICON_TEXT + " " + text) label.setDefaultTextColor(Qt.black) @@ -243,9 +247,10 @@ def update_path(self): # label in the middle mid = self.line().pointAt(0.5) self.label.setPos( - mid - QPointF( + mid + - QPointF( self.label.boundingRect().width() / 2, - self.label.boundingRect().height() / 2 + self.label.boundingRect().height() / 2, ) ) diff --git a/mal_gui/docked_windows/__init__.py b/mal_gui/docked_windows/__init__.py index 498f1fa..228a4f2 100644 --- a/mal_gui/docked_windows/__init__.py +++ b/mal_gui/docked_windows/__init__.py @@ -14,5 +14,5 @@ "EditableDelegate", "CustomDialog", "CustomDialogGlobal", - "Visibility" + "Visibility", ] diff --git a/mal_gui/docked_windows/asset_relations_window.py b/mal_gui/docked_windows/asset_relations_window.py index 8791a83..7d285a5 100644 --- a/mal_gui/docked_windows/asset_relations_window.py +++ b/mal_gui/docked_windows/asset_relations_window.py @@ -1,5 +1,6 @@ from PySide6.QtWidgets import QListWidget + class AssetRelationsWindow(QListWidget): def __init__(self, parent=None): super(AssetRelationsWindow, self).__init__(parent) diff --git a/mal_gui/docked_windows/attack_steps_window.py b/mal_gui/docked_windows/attack_steps_window.py index 4dea683..17ef8ad 100644 --- a/mal_gui/docked_windows/attack_steps_window.py +++ b/mal_gui/docked_windows/attack_steps_window.py @@ -1,5 +1,6 @@ from PySide6.QtWidgets import QListWidget + class AttackStepsWindow(QListWidget): def __init__(self, parent=None): super(AttackStepsWindow, self).__init__(parent) diff --git a/mal_gui/docked_windows/draggable_tree_view.py b/mal_gui/docked_windows/draggable_tree_view.py index 9e80f0a..b0e82ac 100644 --- a/mal_gui/docked_windows/draggable_tree_view.py +++ b/mal_gui/docked_windows/draggable_tree_view.py @@ -14,14 +14,9 @@ if TYPE_CHECKING: from ..object_explorer.asset_item import AssetItem + class DraggableTreeView(QTreeWidget): - def __init__( - self, - scene, - eye_unhide_icon, - eve_hide_icon, - rgb_color_icon - ): + def __init__(self, scene, eye_unhide_icon, eve_hide_icon, rgb_color_icon): super().__init__() self.scene = scene @@ -48,7 +43,9 @@ def __init__( def startDrag(self, supported_actions): """Overrides base method""" item = self.currentItem() - if item and item.parent() is None: # Only start drag if the item is a top-level item (parent) + if ( + item and item.parent() is None + ): # Only start drag if the item is a top-level item (parent) drag = QDrag(self) mime_data = QMimeData() mime_data.setText(item.text(0)) @@ -92,7 +89,7 @@ def set_parent_item_text(self, text, icon=None): if icon: parent_item.setIcon(0, QIcon(icon)) self.addTopLevelItem(parent_item) - self.add_button_to_item(parent_item,"",is_parent=True) + self.add_button_to_item(parent_item, "", is_parent=True) return parent_item def add_child_item(self, parent_item, child_item_asset, text): @@ -100,7 +97,7 @@ def add_child_item(self, parent_item, child_item_asset, text): child_item = QTreeWidgetItem([text, ""]) child_item.assetItemReference = child_item_asset parent_item.addChild(child_item) - self.add_button_to_item(child_item,"",is_parent=False) + self.add_button_to_item(child_item, "", is_parent=False) return child_item def add_button_to_item(self, item, text, is_parent, icon_path=None): @@ -118,14 +115,16 @@ def add_button_to_item(self, item, text, is_parent, icon_path=None): left_eye_button = QPushButton(text) left_eye_button.setIcon(QIcon(self.eye_unhide_icon)) left_eye_button.clicked.connect( - lambda: self.hide_unhide_asset_item(left_eye_button, item)) + lambda: self.hide_unhide_asset_item(left_eye_button, item) + ) # Place the left button in the second column self.setItemWidget(item, 1, left_eye_button) right_color_button = QPushButton(text) right_color_button.setIcon(QIcon(self.rgb_color_icon)) right_color_button.clicked.connect( - lambda: self.show_local_asset_edit_form(item)) + lambda: self.show_local_asset_edit_form(item) + ) # Place the right button in the third column self.setItemWidget(item, 2, right_color_button) @@ -133,35 +132,34 @@ def hide_unhide_asset_item(self, eye_button, item): if self.eye_visibility == Visibility.UNHIDE: self.eye_visibility = Visibility.HIDE - #First Hide the connections associtaed with the asset item + # First Hide the connections associtaed with the asset item asset_item: AssetItem = item.assetItemReference - if hasattr(asset_item, 'connections'): + if hasattr(asset_item, "connections"): connections = asset_item.connections for connection in connections: connection.remove_labels() connection.setVisible(False) - #Then hide the asset item itself + # Then hide the asset item itself asset_item.setVisible(False) eye_button.setIcon(QIcon(self.eve_hide_icon)) else: self.eye_visibility = Visibility.UNHIDE - #First unhide the connections associtaed with the asset item + # First unhide the connections associtaed with the asset item asset_item = item.assetItemReference - if hasattr(asset_item, 'connections'): + if hasattr(asset_item, "connections"): connections = asset_item.connections for connection in connections: connection.restore_labels() connection.setVisible(True) - #Then unhide the asset item itself + # Then unhide the asset item itself asset_item.setVisible(True) - eye_button.setIcon(QIcon(self.eye_unhide_icon)) def adjust_button_width(self): @@ -194,7 +192,6 @@ def adjust_button_width(self): # if font2: # item.setFont(1, QFont(font2)) - def update_color_callback(self, color1, color2): item = self.selected_item item.asset_type_background_color = color1 @@ -211,14 +208,14 @@ def show_global_asset_edit_form(self, item): # self.globalAssetStyleHandlerDialog(self.scene) # print("globalAssetStyleHandlerDialog executing") # self.globalAssetStyleHandlerDialog.exec() - self.dialog = CustomDialogGlobal(self.scene,item) + self.dialog = CustomDialogGlobal(self.scene, item) self.dialog.exec() def check_and_get_if_parent_asset_type_exists(self, child_asset_type): for i in range(self.topLevelItemCount()): parent_item = self.topLevelItem(i) if parent_item.text(0) == child_asset_type: - return parent_item,child_asset_type + return parent_item, child_asset_type return None, None def remove_children(self, parent_item): diff --git a/mal_gui/docked_windows/item_details_window.py b/mal_gui/docked_windows/item_details_window.py index 3c65ecb..60b4159 100644 --- a/mal_gui/docked_windows/item_details_window.py +++ b/mal_gui/docked_windows/item_details_window.py @@ -1,7 +1,6 @@ from PySide6.QtWidgets import QTreeWidget, QTreeWidgetItem from PySide6.QtCore import Qt -from ..object_explorer import ItemBase from PySide6.QtWidgets import QStyledItemDelegate, QComboBox @@ -78,4 +77,4 @@ def _on_item_changed(self, item, column): attr_name = item.data(0, Qt.UserRole) new_value = item.text(1) - self._current_object.set_item_attribute_value(attr_name, new_value) \ No newline at end of file + self._current_object.set_item_attribute_value(attr_name, new_value) diff --git a/mal_gui/docked_windows/properties_window.py b/mal_gui/docked_windows/properties_window.py index 9896e8e..5d9ada5 100644 --- a/mal_gui/docked_windows/properties_window.py +++ b/mal_gui/docked_windows/properties_window.py @@ -1,19 +1,20 @@ -from PySide6.QtCore import Qt, QLocale,QObject +from PySide6.QtCore import Qt, QLocale, QObject from PySide6.QtGui import QDoubleValidator from PySide6.QtWidgets import ( QLineEdit, QStyledItemDelegate, QMessageBox, QTableWidget, - QHeaderView + QHeaderView, ) + class FloatValidator(QDoubleValidator): def __init__(self, parent=None): super(FloatValidator, self).__init__(0.0, 1.0, 2, parent) self.setNotation(QDoubleValidator.StandardNotation) - - #Without US Locale, decimal point was not appearing even when typed from keyboard + + # Without US Locale, decimal point was not appearing even when typed from keyboard self.setLocale(QLocale(QLocale.English, QLocale.UnitedStates)) def validate(self, input, pos): @@ -21,6 +22,7 @@ def validate(self, input, pos): return QDoubleValidator.Intermediate, input, pos return super(FloatValidator, self).validate(input, pos) + class EditableDelegate(QStyledItemDelegate): def __init__(self, asset_item, parent=None): super(EditableDelegate, self).__init__(parent) @@ -42,11 +44,13 @@ def setEditorData(self, editor, index): def setModelData(self, editor, model, index): """Overrides base""" value = editor.text() - print("Value Entered: "+ value) + print("Value Entered: " + value) # setattr(selected_item.asset, row[0],value) state = editor.validator().validate(value, 0) if state[0] != QDoubleValidator.Acceptable: - QMessageBox.warning(editor, "Input Error", "Value must be a float between 0.0 and 1.0.") + QMessageBox.warning( + editor, "Input Error", "Value must be a float between 0.0 and 1.0." + ) # Revert to previous valid value (optional) # editor.setText(index.model().data(index, Qt.EditRole)) else: @@ -65,13 +69,12 @@ def validate_editor(self): state = editor.validator().validate(editor.text(), 0) if state[0] != QDoubleValidator.Acceptable: QMessageBox.warning( - editor, - "Input Error", - "Value must be a float between 0.0 and 1.0." + editor, "Input Error", "Value must be a float between 0.0 and 1.0." ) # Revert to previous valid value (optional) # editor.setText(self.oldValue) + class PropertiesWindow(QObject): def __init__(self): super().__init__() @@ -80,20 +83,26 @@ def __init__(self): self.properties_table = QTableWidget() self.properties_table.setColumnCount(3) self.properties_table.setHorizontalHeaderLabels( - ["Defense Property", "Value", "Default Value"]) + ["Defense Property", "Value", "Default Value"] + ) # self.properties_table.setRowCount(10) # Example: setting 10 rows self.properties_table.horizontalHeader().setSectionResizeMode( - QHeaderView.Stretch) + QHeaderView.Stretch + ) self.properties_table.horizontalHeader().setSectionResizeMode( - 0, QHeaderView.ResizeToContents) # Adjust the first column + 0, QHeaderView.ResizeToContents + ) # Adjust the first column self.properties_table.horizontalHeader().setSectionResizeMode( - 1, QHeaderView.ResizeToContents) # Adjust the second column + 1, QHeaderView.ResizeToContents + ) # Adjust the second column self.properties_table.horizontalHeader().setSectionResizeMode( - 2, QHeaderView.ResizeToContents) # Adjust the third column + 2, QHeaderView.ResizeToContents + ) # Adjust the third column # Hide the vertical header (row numbers) self.properties_table.verticalHeader().setVisible(False) self.properties_table.setItemDelegateForColumn( - 1, EditableDelegate(self.properties_table)) + 1, EditableDelegate(self.properties_table) + ) diff --git a/mal_gui/docked_windows/style_configuration.py b/mal_gui/docked_windows/style_configuration.py index 79760b3..edf4c03 100644 --- a/mal_gui/docked_windows/style_configuration.py +++ b/mal_gui/docked_windows/style_configuration.py @@ -7,17 +7,19 @@ QColorDialog, QFormLayout, QDialogButtonBox, - QLabel + QLabel, ) from PySide6.QtCore import Signal from PySide6.QtGui import QColor from ..object_explorer.asset_item import AssetItem + class Visibility(Enum): HIDE = 1 UNHIDE = 2 + class CustomDialog(QDialog): color_changed_1 = Signal(QColor) color_changed_2 = Signal(QColor) @@ -36,12 +38,16 @@ def __init__(self, item, update_color_callback, parent=None): layout.addRow("Name:", self.name_edit) self.color_button_1 = QPushButton("Select AssetType background color") - print("type(self.selected_item.childItemObj) = " + - str(type(self.selected_item.asset_item_reference))) + print( + "type(self.selected_item.childItemObj) = " + + str(type(self.selected_item.asset_item_reference)) + ) self.color_button_1.setStyleSheet( - f"background-color: {self.selected_item.asset_item_reference.asset_type_background_color.name()}") + f"background-color: {self.selected_item.asset_item_reference.asset_type_background_color.name()}" + ) self.color_button_1.clicked.connect( - lambda: self.open_color_dialog(1,self.selected_item.asset_item_reference)) + lambda: self.open_color_dialog(1, self.selected_item.asset_item_reference) + ) layout.addRow("Color 1:", self.color_button_1) self.rgb_label_1 = QLabel("RGB: ") @@ -49,15 +55,19 @@ def __init__(self, item, update_color_callback, parent=None): self.color_button_2 = QPushButton("Select AssetName background color") self.color_button_2.setStyleSheet( - f"background-color: {self.selected_item.asset_item_reference.asset_name_background_color.name()}") + f"background-color: {self.selected_item.asset_item_reference.asset_name_background_color.name()}" + ) self.color_button_2.clicked.connect( - lambda: self.open_color_dialog(2,self.selected_item.asset_item_reference)) + lambda: self.open_color_dialog(2, self.selected_item.asset_item_reference) + ) layout.addRow("Color 2:", self.color_button_2) self.rgb_label_2 = QLabel("RGB: ") layout.addRow("RGB Values for Color 2:", self.rgb_label_2) - self.button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) + self.button_box = QDialogButtonBox( + QDialogButtonBox.Ok | QDialogButtonBox.Cancel + ) self.button_box.accepted.connect(self.accept) self.button_box.rejected.connect(self.reject) layout.addRow(self.button_box) @@ -89,14 +99,16 @@ def open_color_dialog(self, item_number, asset_item_reference): # QMessageBox.information(self, "Color Selected", f"Selected color for Item 2: {color.name()}") asset_item_reference.asset_name_background_color = color asset_item_reference.update() - + def update_color_label_1(self, color): self.rgb_label_1.setText( - f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}") - + f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}" + ) + def update_color_label_2(self, color): self.rgb_label_2.setText( - f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}") + f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}" + ) def get_name(self): return self.name_edit.text() @@ -112,15 +124,13 @@ def accept(self): self.update_color_callback(self.get_color_1(), self.get_color_2()) - - class CustomDialogGlobal(QDialog): color_changed_1 = Signal(QColor) color_changed_2 = Signal(QColor) - def __init__(self,scene, item, parent=None): + def __init__(self, scene, item, parent=None): super().__init__(parent) - + self.scene = scene self.selectedAssetType = item @@ -136,21 +146,21 @@ def __init__(self,scene, item, parent=None): # self.color_button_1.setStyleSheet(f"background-color: {self.selected_item.asset_item_reference.asset_type_background_color.name()}") self.color_button_1.clicked.connect(lambda: self.open_color_dialog(1)) layout.addRow("Color 1:", self.color_button_1) - + self.rgb_label_1 = QLabel("RGB: ") layout.addRow("RGB Values for Color 1:", self.rgb_label_1) - self.color_button_2 = QPushButton("Select AssetName background color") # self.color_button_2.setStyleSheet(f"background-color: {self.selected_item.asset_item_reference.asset_name_background_color.name()}") self.color_button_2.clicked.connect(lambda: self.open_color_dialog(2)) layout.addRow("Color 2:", self.color_button_2) - self.rgb_label_2 = QLabel("RGB: ") layout.addRow("RGB Values for Color 2:", self.rgb_label_2) - self.button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) + self.button_box = QDialogButtonBox( + QDialogButtonBox.Ok | QDialogButtonBox.Cancel + ) self.button_box.accepted.connect(self.accept) self.button_box.rejected.connect(self.reject) layout.addRow(self.button_box) @@ -179,12 +189,16 @@ def open_color_dialog(self, item_number): self.color_changed_2.emit(self.selected_color_2) self.color_button_2.setStyleSheet(f"background-color: {color.name()}") # QMessageBox.information(self, "Color Selected", f"Selected color for Item 2: {color.name()}") - + def update_color_label_1(self, color): - self.rgb_label_1.setText(f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}") - + self.rgb_label_1.setText( + f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}" + ) + def update_color_label_2(self, color): - self.rgb_label_2.setText(f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}") + self.rgb_label_2.setText( + f"RGB: {color.red()}, {color.green()}, {color.blue()}, {color.alpha()}" + ) def get_name(self): return self.name_edit.text() @@ -203,4 +217,4 @@ def accept(self): if item.asset_type == self.selectedAssetType.text(0): item.asset_type_background_color = self.get_color_1() item.asset_name_background_color = self.get_color_2() - item.update() \ No newline at end of file + item.update() diff --git a/mal_gui/file_utils.py b/mal_gui/file_utils.py index cd1e09e..2f56e53 100644 --- a/mal_gui/file_utils.py +++ b/mal_gui/file_utils.py @@ -4,6 +4,7 @@ PACKAGE_DIR = Path(__file__).resolve().parent + def image_path(filename): """From a filename, return the absolute path of the image""" return str(PACKAGE_DIR / "images" / filename) diff --git a/mal_gui/main_window.py b/mal_gui/main_window.py index 0a76556..8f7dbf1 100644 --- a/mal_gui/main_window.py +++ b/mal_gui/main_window.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Any, Optional +from typing import Optional from PySide6.QtWidgets import ( QWidget, @@ -17,12 +17,12 @@ QFileDialog, QMessageBox, QTableWidgetItem, - QApplication + QApplication, ) from PySide6.QtGui import QDrag, QAction, QIcon, QIntValidator from PySide6.QtCore import Qt, QMimeData, QByteArray, QSize, Signal, QPointF -from qt_material import apply_stylesheet,list_themes +from qt_material import apply_stylesheet, list_themes from maltoolbox import __version__ as maltoolbox_version from maltoolbox.language import LanguageGraph @@ -46,12 +46,13 @@ PropertiesWindow, EditableDelegate, AttackStepsWindow, - AssetRelationsWindow + AssetRelationsWindow, ) # Used to create absolute paths of assets PACKAGE_DIR = Path(__file__).resolve().parent + class DraggableListWidget(QListWidget): def mousePressEvent(self, event): if event.buttons() == Qt.LeftButton: @@ -60,7 +61,8 @@ def mousePressEvent(self, event): drag = QDrag(self) mime_data = QMimeData() mime_data.setData( - "application/x-qabstractitemmodeldatalist", QByteArray()) + "application/x-qabstractitemmodeldatalist", QByteArray() + ) mime_data.setData("text/plain", item.text().encode()) drag.setMimeData(mime_data) drag.exec() @@ -72,7 +74,7 @@ class MainWindow(QMainWindow): def __init__(self, app: QApplication, lang_file_path: str): super().__init__() self.setWindowTitle("MAL GUI") - self.app = app # declare an app member + self.app = app # declare an app member self.scenario_file_name = None self.model_file_name = None @@ -108,20 +110,15 @@ def clear_window(self): self.removeDockWidget(dock_widget) def load_scene( - self, - lang_file_path: str, - model: Model, - scenario: Optional[Scenario] = None - ): + self, lang_file_path: str, model: Model, scenario: Optional[Scenario] = None + ): """Load scene with given language and model""" print("LOADING SCENE!") self.clear_window() self.lang_file_path = lang_file_path lang_graph = LanguageGraph.load_from_file(lang_file_path) self.asset_factory = self.create_asset_factory(lang_graph) - self.scene = self.create_scene( - lang_graph, self.asset_factory, model, scenario - ) + self.scene = self.create_scene(lang_graph, self.asset_factory, model, scenario) self.create_menu_bar() self.create_actions(self.scene) @@ -149,37 +146,32 @@ def create_asset_factory(self, lang_graph: LanguageGraph): "RoutingFirewall": image_path("routingFirewall.png"), "SoftwareProduct": image_path("softwareProduct.png"), "SoftwareVulnerability": image_path("softwareVulnerability.png"), - "User": image_path("user.png") + "User": image_path("user.png"), } # Create a registry as a dictionary containing # name as key and class as value asset_factory = AssetFactory() - asset_factory.register_asset( - "Attacker", image_path("attacker.png") - ) + asset_factory.register_asset("Attacker", image_path("attacker.png")) for asset in lang_graph.assets.values(): if not asset.is_abstract: asset_factory.register_asset( - asset.name, - asset_images.get(asset.name, image_path('unknown.png')) + asset.name, asset_images.get(asset.name, image_path("unknown.png")) ) return asset_factory def create_scene( - self, - lang_graph: LanguageGraph, - asset_factory: AssetFactory, - model: Model, - scenario: Optional[Scenario] = None - ): + self, + lang_graph: LanguageGraph, + asset_factory: AssetFactory, + model: Model, + scenario: Optional[Scenario] = None, + ): """Create and initialize scene from language""" - model_scene = ModelScene( - asset_factory, lang_graph, model, self, scenario - ) + model_scene = ModelScene(asset_factory, lang_graph, model, self, scenario) return model_scene @@ -198,7 +190,6 @@ def create_view(self, scene: ModelScene): ) return view - def create_side_panels(self, asset_factory: AssetFactory): """Add side panel objects""" @@ -211,10 +202,7 @@ def create_side_panels(self, asset_factory: AssetFactory): rgb_color_icon_image = image_path("rgbColor.png") self.object_explorer_tree = DraggableTreeView( - self.scene, - eye_unhide_icon_image, - eye_hide_icon_image, - rgb_color_icon_image + self.scene, eye_unhide_icon_image, eye_hide_icon_image, rgb_color_icon_image ) for _, values in asset_factory.asset_registry.items(): @@ -227,45 +215,46 @@ def create_side_panels(self, asset_factory: AssetFactory): self.addDockWidget(Qt.LeftDockWidgetArea, dock_object_explorer) dock_widgets.append(dock_object_explorer) - #EDOC Tab with treeview + # EDOC Tab with treeview component_tab_tree = QTreeWidget() component_tab_tree.setHeaderLabel(None) - #ItemDetails with treeview + # ItemDetails with treeview self.item_details_window = ItemDetailsWindow() - dock_item_details = QDockWidget("Item Details",self) + dock_item_details = QDockWidget("Item Details", self) dock_item_details.setWidget(self.item_details_window) self.addDockWidget(Qt.LeftDockWidgetArea, dock_item_details) dock_widgets.append(dock_item_details) - #Properties Tab with tableview + # Properties Tab with tableview self.properties_docked_window = PropertiesWindow() self.properties_table = self.properties_docked_window.properties_table - dock_properties = QDockWidget("Properties",self) + dock_properties = QDockWidget("Properties", self) dock_properties.setWidget(self.properties_table) self.addDockWidget(Qt.LeftDockWidgetArea, dock_properties) dock_widgets.append(dock_properties) - #AttackSteps Tab with ListView + # AttackSteps Tab with ListView self.attack_steps_docked_window = AttackStepsWindow() - dock_attack_steps = QDockWidget("Attack Steps",self) + dock_attack_steps = QDockWidget("Attack Steps", self) dock_attack_steps.setWidget(self.attack_steps_docked_window) self.addDockWidget(Qt.LeftDockWidgetArea, dock_attack_steps) dock_widgets.append(dock_attack_steps) - #AssetRelations Tab with ListView + # AssetRelations Tab with ListView self.asset_relations_docker_window = AssetRelationsWindow() - dock_asset_relations = QDockWidget("Asset Relations",self) - dock_asset_relations.setFeatures(QDockWidget.DockWidgetFloatable | - QDockWidget.DockWidgetMovable) + dock_asset_relations = QDockWidget("Asset Relations", self) + dock_asset_relations.setFeatures( + QDockWidget.DockWidgetFloatable | QDockWidget.DockWidgetMovable + ) dock_asset_relations.setWidget(self.asset_relations_docker_window) self.addDockWidget(Qt.LeftDockWidgetArea, dock_asset_relations) dock_widgets.append(dock_asset_relations) - #Keep Propeties Window and Attack Step Window Tabbed + # Keep Propeties Window and Attack Step Window Tabbed self.tabifyDockWidget(dock_properties, dock_attack_steps) - #Keep the properties Window highlighted and raised + # Keep the properties Window highlighted and raised dock_properties.raise_() return dock_widgets @@ -282,7 +271,7 @@ def show_image_icon_checkbox_changed(self, checked): """Called on button click""" print("self.show_image_icon_checkbox_changed clicked") for item in self.scene.items(): - if isinstance(item, (AssetItem,AssetsContainer)): + if isinstance(item, (AssetItem, AssetsContainer)): item.toggle_icon_visibility() def fit_to_view_button_clicked(self): @@ -302,8 +291,8 @@ def update_properties_window(self, asset_item: AssetItem): for attack_step_name, value in asset.defenses.items(): # Add defenses that are set in model attack_step = asset.lg_asset.attack_steps[attack_step_name] - if attack_step.ttc and len(attack_step.ttc['arguments']) > 0: - default_value = attack_step.ttc['arguments'][0] + if attack_step.ttc and len(attack_step.ttc["arguments"]) > 0: + default_value = attack_step.ttc["arguments"][0] else: default_value = 0.0 properties.append((attack_step_name, str(value), str(default_value))) @@ -313,8 +302,8 @@ def update_properties_window(self, asset_item: AssetItem): if attack_step.name in asset.defenses: continue if attack_step.type == "defense": - if attack_step.ttc and len(attack_step.ttc['arguments']) > 0: - default_value = attack_step.ttc['arguments'][0] + if attack_step.ttc and len(attack_step.ttc["arguments"]) > 0: + default_value = attack_step.ttc["arguments"][0] else: default_value = 0.0 properties.append((attack_step.name, "", str(default_value))) @@ -324,22 +313,32 @@ def update_properties_window(self, asset_item: AssetItem): self.properties_table.setRowCount(num_rows) self.properties_table.currentItem = asset_item - for row, (property_key, property_value, property_default) in enumerate(properties): + for row, (property_key, property_value, property_default) in enumerate( + properties + ): col_property_name = QTableWidgetItem(property_key) - col_property_name.setFlags(Qt.ItemIsEnabled) # Make the property name read-only + col_property_name.setFlags( + Qt.ItemIsEnabled + ) # Make the property name read-only col_value = QTableWidgetItem(property_value) - col_value.setFlags(Qt.ItemIsEditable | Qt.ItemIsEnabled) # Make the value editable + col_value.setFlags( + Qt.ItemIsEditable | Qt.ItemIsEnabled + ) # Make the value editable col_default_value = QTableWidgetItem(property_default) - col_default_value.setFlags(Qt.ItemIsEnabled) # Make the default value read-only + col_default_value.setFlags( + Qt.ItemIsEnabled + ) # Make the default value read-only self.properties_table.setItem(row, 0, col_property_name) self.properties_table.setItem(row, 1, col_value) self.properties_table.setItem(row, 2, col_default_value) # Set the item delegate and pass asset_item - based on Andrei's input - self.properties_table.setItemDelegateForColumn(1, EditableDelegate(asset_item)) + self.properties_table.setItemDelegateForColumn( + 1, EditableDelegate(asset_item) + ) else: self.properties_table.currentItem = None @@ -389,25 +388,29 @@ def create_actions(self, scene: ModelScene): self.cut_action = QAction(QIcon(cut_icon), "Cut", self) self.cut_action.setShortcut("Ctrl+x") self.cut_action.triggered.connect( - lambda: self.scene.cut_assets(scene.selectedItems())) + lambda: self.scene.cut_assets(scene.selectedItems()) + ) copy_icon = image_path("copyIcon.png") self.copy_action = QAction(QIcon(copy_icon), "Copy", self) self.copy_action.setShortcut("Ctrl+c") self.copy_action.triggered.connect( - lambda: self.scene.copy_assets(scene.selectedItems())) + lambda: self.scene.copy_assets(scene.selectedItems()) + ) paste_icon = image_path("pasteIcon.png") self.paste_action = QAction(QIcon(paste_icon), "Paste", self) self.paste_action.setShortcut("Ctrl+v") self.paste_action.triggered.connect( - lambda: self.scene.paste_assets(QPointF(0,0))) + lambda: self.scene.paste_assets(QPointF(0, 0)) + ) delete_icon = image_path("deleteIcon.png") self.delete_action = QAction(QIcon(delete_icon), "Delete", self) self.delete_action.setShortcut("Delete") self.delete_action.triggered.connect( - lambda: self.scene.delete_assets(scene.selectedItems())) + lambda: self.scene.delete_assets(scene.selectedItems()) + ) def create_menu_bar(self): """Create the menu and add to the GUI""" @@ -417,8 +420,12 @@ def create_menu_bar(self): self.file_menu_new_action = self.file_menu.addAction("New") self.file_menu_open_action = self.file_menu.addAction("Load Model/Scenario") self.file_menu_save_as_action = self.file_menu.addAction("Export Model..") - self.file_menu_export_scenario_action = self.file_menu.addAction("Export Scenario..") - self.file_menu_save_as_drawio = self.file_menu.addAction("Export draw.io file..") + self.file_menu_export_scenario_action = self.file_menu.addAction( + "Export Scenario.." + ) + self.file_menu_save_as_drawio = self.file_menu.addAction( + "Export draw.io file.." + ) self.file_menu_quit_action = self.file_menu.addAction("Quit") self.file_menu_open_action.triggered.connect(self.load_model_or_scenario) self.file_menu_save_as_action.triggered.connect(self.save_as_model) @@ -456,17 +463,19 @@ def create_toolbar(self): toolbar.addWidget(show_association_checkbox_label) toolbar.addWidget(show_association_checkbox) show_association_checkbox.stateChanged.connect( - self.show_association_checkbox_changed) + self.show_association_checkbox_changed + ) toolbar.addSeparator() - show_image_icon_checkbox_label = QLabel("Show Image Icon") + show_image_icon_checkbox_label = QLabel("Show Image Icon") show_image_icon_checkbox = QCheckBox() show_image_icon_checkbox.setCheckState(Qt.CheckState.Checked) toolbar.addWidget(show_image_icon_checkbox_label) toolbar.addWidget(show_image_icon_checkbox) - show_image_icon_checkbox.stateChanged\ - .connect(self.show_image_icon_checkbox_changed) + show_image_icon_checkbox.stateChanged.connect( + self.show_image_icon_checkbox_changed + ) toolbar.addSeparator() @@ -478,8 +487,7 @@ def create_toolbar(self): # No limit on zoom level, but should be an integer self.zoom_line_edit.setValidator(QIntValidator()) self.zoom_line_edit.setText("100") - self.zoom_line_edit.returnPressed.connect( - self.set_zoom_level_from_line_edit) + self.zoom_line_edit.returnPressed.connect(self.set_zoom_level_from_line_edit) self.zoom_line_edit.setFixedWidth(40) toolbar.addWidget(self.zoom_label) toolbar.addWidget(self.zoom_line_edit) @@ -493,25 +501,22 @@ def create_toolbar(self): toolbar.addAction(self.delete_action) toolbar.addSeparator() fit_to_view_icon = image_path("fitToView.png") - fit_to_view_button = QPushButton( - QIcon(fit_to_view_icon), "Fit To View") + fit_to_view_button = QPushButton(QIcon(fit_to_view_icon), "Fit To View") toolbar.addWidget(fit_to_view_button) fit_to_view_button.clicked.connect(self.fit_to_view_button_clicked) toolbar.addSeparator() - #Material Theme - https://pypi.org/project/qt-material/ - material_theme_label = QLabel("Theme") + # Material Theme - https://pypi.org/project/qt-material/ + material_theme_label = QLabel("Theme") self.theme_combo_box = QComboBox() - self.theme_combo_box.addItem('None') + self.theme_combo_box.addItem("None") inbuilt_theme_list_from_package = list_themes() self.theme_combo_box.addItems(inbuilt_theme_list_from_package) toolbar.addWidget(material_theme_label) toolbar.addWidget(self.theme_combo_box) - self.theme_combo_box.currentIndexChanged.connect( - self.on_theme_selection_change - ) + self.theme_combo_box.currentIndexChanged.connect(self.on_theme_selection_change) toolbar.addSeparator() return toolbar @@ -537,10 +542,10 @@ def update_zoom_label(self): def load_model_or_scenario(self): """Load a file, either model or scenario""" - file_extension_filter = \ - "YAML Files (*.yaml *.yml);;JSON Files (*.json)" + file_extension_filter = "YAML Files (*.yaml *.yml);;JSON Files (*.json)" file_path, _ = QFileDialog.getOpenFileName( - None, "Select model or scenario File", "", file_extension_filter) + None, "Select model or scenario File", "", file_extension_filter + ) if not file_path: print("No valid path detected for loading") @@ -552,11 +557,11 @@ def load_model_or_scenario(self): "Load New Project", "Loading a new project will delete current work (if any). " "Do you want to continue ?", - QMessageBox.Ok | QMessageBox.Cancel + QMessageBox.Ok | QMessageBox.Cancel, ) if open_project_user_confirmation == QMessageBox.Ok: - #clear scene so that canvas becomes blank + # clear scene so that canvas becomes blank self.scene.clear() try: self.load_model(file_path) @@ -568,7 +573,6 @@ def load_model_or_scenario(self): print("User cancelled 'Load'") return - def load_scenario(self, file_path: str): """Load model and agents from a scenario""" scenario = Scenario.load_from_file(file_path) @@ -580,9 +584,7 @@ def load_scenario(self, file_path: str): def load_model(self, file_path: str): """Load a MAL model from a file""" - model = Model.load_from_file( - file_path, self.scene.lang_graph - ) + model = Model.load_from_file(file_path, self.scene.lang_graph) self.load_scene(self.lang_file_path, model, None) self.scenario_file_name = None self.model_file_name = file_path @@ -590,15 +592,12 @@ def load_model(self, file_path: str): def add_positions_to_model(self): """Add x/y positions to asset extras of model""" for asset in self.scene.model.assets.values(): - print(f'ASSET NAME:{asset.name} ID:{asset.id} TYPE:{asset.type}') + print(f"ASSET NAME:{asset.name} ID:{asset.id} TYPE:{asset.type}") item = self.scene._asset_id_to_item[int(asset.id)] position = item.pos() extras_dict = asset.extras if asset.extras else {} - extras_dict["position"] = { - "x": position.x(), - "y": position.y() - } + extras_dict["position"] = {"x": position.x(), "y": position.y()} asset.extras = extras_dict def save_model(self): @@ -610,7 +609,7 @@ def save_model(self): self.save_as_model() def save_as_model(self): - """ `Save as`. Let user select target file and save model.""" + """`Save as`. Let user select target file and save model.""" self.add_positions_to_model() file_dialog = QFileDialog() file_dialog.setAcceptMode(QFileDialog.AcceptSave) @@ -632,7 +631,8 @@ def save_as_model(self): return def save_as_drawio(self): - """ `Save as`. Let user select target file and save .drawio file.""" + """`Save as`. Let user select target file and save .drawio file.""" + def versiontuple(v): return tuple(map(int, (v.split(".")))) @@ -653,7 +653,7 @@ def versiontuple(v): None, "Save As Draw.io file", default_name, - "DrawIO Files (*.drawio);;All Files (*)" + "DrawIO Files (*.drawio);;All Files (*)", ) if not file_path: @@ -673,7 +673,7 @@ def versiontuple(v): return def save_as_scenario(self): - """ `Save as`. Let user select target file and save scenario.""" + """`Save as`. Let user select target file and save scenario.""" file_dialog = QFileDialog() file_dialog.setAcceptMode(QFileDialog.AcceptSave) file_dialog.setDefaultSuffix("yaml") @@ -716,9 +716,7 @@ def save_as_scenario(self): if self.scene.scenario: if self.scene.scenario.rewards: - rewards = ( - self.scene.scenario.rewards.to_dict() - ) + rewards = self.scene.scenario.rewards.to_dict() if self.scene.scenario.false_negative_rates: false_negative_rates = ( self.scene.scenario.false_negative_rates.to_dict() @@ -728,13 +726,9 @@ def save_as_scenario(self): self.scene.scenario.false_positive_rates.to_dict() ) if self.scene.scenario.is_actionable: - is_actionable = ( - self.scene.scenario.is_actionable.to_dict() - ) + is_actionable = self.scene.scenario.is_actionable.to_dict() if self.scene.scenario.is_observable: - is_observable = ( - self.scene.scenario.is_observable.to_dict() - ) + is_observable = self.scene.scenario.is_observable.to_dict() try: scenario = Scenario( @@ -749,20 +743,19 @@ def save_as_scenario(self): ) scenario.save_to_file(file_path) - if hasattr(self, '_lang_file'): + if hasattr(self, "_lang_file"): scenario_dict = yaml.safe_load(open(file_path, "r")) scenario_dict["lang_file"] = self._lang_file yaml.safe_dump(scenario_dict, open(file_path, "w")) except Exception as e: self.show_error_popup("Could not save scenario: " + str(e)) - def quitApp(self): self.app.quit() def show_information_popup(self, message_text): """Show a popup with given message""" - parent_widget = QWidget() #To maintain object lifetim + parent_widget = QWidget() # To maintain object lifetim message_box = QMessageBox(parent_widget) message_box.setIcon(QMessageBox.Information) message_box.setWindowTitle("Information") @@ -773,7 +766,7 @@ def show_information_popup(self, message_text): def show_error_popup(self, message_text): """Show error popup with given message""" - parent_widget = QWidget() #To maintain object lifetim + parent_widget = QWidget() # To maintain object lifetim message_box = QMessageBox(parent_widget) message_box.setIcon(QMessageBox.Critical) message_box.setWindowTitle("Error") @@ -788,19 +781,19 @@ def update_explorer_docked_window(self): """ self.object_explorer_tree.clear_all_object_explorer_child_items() - #Fill all the items from Scene one by one + # Fill all the items from Scene one by one for child_asset_item in self.scene.items(): - if isinstance(child_asset_item,AssetItem): + if isinstance(child_asset_item, AssetItem): # Check if parent exists before adding child - parent_item, parent_asset_type = self.object_explorer_tree\ - .check_and_get_if_parent_asset_type_exists( + parent_item, parent_asset_type = ( + self.object_explorer_tree.check_and_get_if_parent_asset_type_exists( child_asset_item.asset_type ) + ) if parent_asset_type: self.object_explorer_tree.add_child_item( - parent_item,child_asset_item, - str(child_asset_item.asset_name) + parent_item, child_asset_item, str(child_asset_item.asset_name) ) def on_theme_selection_change(self): diff --git a/mal_gui/model_scene.py b/mal_gui/model_scene.py index fa2da6a..1ebe305 100644 --- a/mal_gui/model_scene.py +++ b/mal_gui/model_scene.py @@ -2,7 +2,7 @@ import pickle import base64 -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Optional from PySide6.QtWidgets import ( QGraphicsScene, @@ -11,7 +11,7 @@ QGraphicsLineItem, QDialog, QGraphicsRectItem, - QGraphicsTextItem + QGraphicsTextItem, ) from PySide6.QtGui import QTransform, QAction, QUndoStack, QPen from PySide6.QtCore import QLineF, Qt, QPointF, QRectF @@ -19,10 +19,21 @@ from maltoolbox.model import Model from malsim.config.agent_settings import AttackerSettings -from mal_gui.undo_redo_commands.create_goal_connection_command import CreateGoalConnectionCommand +from mal_gui.undo_redo_commands.create_goal_connection_command import ( + CreateGoalConnectionCommand, +) -from .connection_item import AssociationConnectionItem, AttackerConnectionBase,EntrypointConnectionItem, GoalConnectionItem -from .connection_dialog import AssociationConnectionDialog,EntrypointConnectionDialog, GoalConnectionDialog +from .connection_item import ( + AssociationConnectionItem, + AttackerConnectionBase, + EntrypointConnectionItem, + GoalConnectionItem, +) +from .connection_dialog import ( + AssociationConnectionDialog, + EntrypointConnectionDialog, + GoalConnectionDialog, +) from .object_explorer import AssetItem, AttackerItem, EditableTextItem, ItemBase from .assets_container import AssetsContainer, AssetsContainerRectangleBox @@ -37,7 +48,7 @@ CreateAssociationConnectionCommand, CreateEntrypointConnectionCommand, DeleteConnectionCommand, - ContainerizeAssetsCommand + ContainerizeAssetsCommand, ) if TYPE_CHECKING: @@ -47,6 +58,7 @@ from .connection_item import IConnectionItem from malsim.scenario import Scenario + class ModelScene(QGraphicsScene): def __init__( self, @@ -160,9 +172,7 @@ def mousePressEvent(self, event): def _start_connection(self, event, item): self.start_item = item - self.line_item = QGraphicsLineItem( - QLineF(event.scenePos(), event.scenePos()) - ) + self.line_item = QGraphicsLineItem(QLineF(event.scenePos(), event.scenePos())) self.addItem(self.line_item) def _handle_left_press(self, event, item): @@ -190,7 +200,8 @@ def _prepare_move(self, item, multi): if multi: self.dragged_items = [ - i for i in self.selectedItems() + i + for i in self.selectedItems() if isinstance(i, (AssetItem, AttackerItem)) ] else: @@ -220,9 +231,7 @@ def mouseMoveEvent(self, event): super().mouseMoveEvent(event) def _update_connection_line(self, event): - self.line_item.setLine( - QLineF(self.line_item.line().p1(), event.scenePos()) - ) + self.line_item.setLine(QLineF(self.line_item.line().p1(), event.scenePos())) def _update_dragged_items(self, event): delta = event.scenePos() - self.start_pos @@ -258,9 +267,7 @@ def _finalize_drag_or_selection(self): self._reset_drag_state() def _finalize_selection_rect(self): - items = self.items( - self.selection_rect.rect(), Qt.IntersectsItemShape - ) + items = self.items(self.selection_rect.rect(), Qt.IntersectsItemShape) for item in items: if isinstance(item, (AssetItem, AttackerItem)): item.setSelected(True) @@ -293,9 +300,7 @@ def _reset_drag_state(self): def _finalize_connection(self, released_item): self.removeItem(self.line_item) - self.end_item = ( - released_item if isinstance(released_item, ItemBase) else None - ) + self.end_item = released_item if isinstance(released_item, ItemBase) else None if self.end_item: if isinstance(self.start_item, AssetItem) and isinstance( @@ -345,18 +350,14 @@ def _create_attacker_connection(self): if isinstance(self.start_item, AttackerItem) and isinstance( self.end_item, AttackerItem ): - raise TypeError( - "Start and end item can not both be type 'Attacker'" - ) + raise TypeError("Start and end item can not both be type 'Attacker'") attacker = ( self.start_item if isinstance(self.start_item, AttackerItem) else self.end_item ) - asset = ( - self.end_item if attacker is self.start_item else self.start_item - ) + asset = self.end_item if attacker is self.start_item else self.start_item # Ask user what type of attacker connection they want scene_pos = self.views()[0].mapToGlobal( @@ -394,9 +395,7 @@ def _create_entrypoint_connection(self, attacker, asset): ) def _create_goal_connection(self, attacker, asset): - dialog = GoalConnectionDialog( - attacker, asset, self.lang_graph, self.model - ) + dialog = GoalConnectionDialog(attacker, asset, self.lang_graph, self.model) if dialog.exec() != QDialog.Accepted: return @@ -442,33 +441,35 @@ def contextMenuEvent(self, event): print("Found text box", item) item = item.parentItem() item = item.parentItem() if item else None - if isinstance(item, (AssociationConnectionItem, AttackerConnectionBase)): + if isinstance( + item, (AssociationConnectionItem, AttackerConnectionBase) + ): print("Found parent of text box, a connection item") self.show_connection_item_context_menu(event.screenPos(), item) elif isinstance(item, AssetsContainer): - print("Found Assets Container item",item) + print("Found Assets Container item", item) self.show_assets_container_context_menu(event.screenPos(), item) elif isinstance(item, AssetsContainerRectangleBox): - print("Found Assets Container Box",item) + print("Found Assets Container Box", item) self.show_assets_container_box_context_menu(event.screenPos(), item) else: - self.show_scene_context_menu(event.screenPos(),event.scenePos()) + self.show_scene_context_menu(event.screenPos(), event.scenePos()) def assign_position_to_assets_without_positions( - self, assets_without_position, x_max, y_max - ): + self, assets_without_position, x_max, y_max + ): """Assign position to assets that don't have any""" distance_between_two_assets_vertically = 200 for i, asset in enumerate(assets_without_position): x_pos = x_max - y_pos = y_max + (i* distance_between_two_assets_vertically) - print("In x_pos= "+ str(x_pos)) - print("In y_pos= "+ str(y_pos)) - asset.setPos(QPointF(x_pos,y_pos)) + y_pos = y_max + (i * distance_between_two_assets_vertically) + print("In x_pos= " + str(x_pos)) + print("In y_pos= " + str(y_pos)) + asset.setPos(QPointF(x_pos, y_pos)) def draw_model(self): """Draw all assets in the model""" @@ -477,35 +478,33 @@ def draw_model(self): x_max = 0 y_max = 0 for asset in self.model.assets.values(): - - if 'position' in asset.extras: + if "position" in asset.extras: pos = QPointF( - asset.extras['position']['x'], - asset.extras['position']['y'] + asset.extras["position"]["x"], asset.extras["position"]["y"] ) # Storing x_max and y_max to be used at the end # for moving the assets without position - if x_max< asset.extras['position']['x']: - x_max = asset.extras['position']['x'] - print("x_max = "+ str(x_max)) - if y_max < asset.extras['position']['y']: - y_max = asset.extras['position']['y'] - print("y_max = "+ str(y_max)) + if x_max < asset.extras["position"]["x"]: + x_max = asset.extras["position"]["x"] + print("x_max = " + str(x_max)) + if y_max < asset.extras["position"]["y"]: + y_max = asset.extras["position"]["y"] + print("y_max = " + str(y_max)) else: - pos = QPointF(0,0) + pos = QPointF(0, 0) new_item = self.asset_factory.create_asset_item(asset, pos) self._asset_id_to_item[asset.id] = new_item self.addItem(new_item) # extract assets without position - if 'position' not in asset.extras: + if "position" not in asset.extras: assets_without_position.append(new_item) self.assign_position_to_assets_without_positions( - assets_without_position,x_max, y_max + assets_without_position, x_max, y_max ) # Draw associations between assets @@ -515,14 +514,13 @@ def draw_model(self): self.add_association_connection( self._asset_id_to_item[asset.id], self._asset_id_to_item[associated_asset.id], - fieldname + fieldname, ) # Draw attackers if they exists in scenario if self.scenario: agents = self.scenario.agent_settings for name, agent_info in agents.items(): - if isinstance(agent_info, AttackerSettings): attacker_item = self.create_attacker( QPointF(0, 0), name, agent_info.entry_points @@ -530,18 +528,18 @@ def draw_model(self): for entry_point in agent_info.entry_points: entrypoint_full_name = ( - entry_point if isinstance(entry_point, str) else entry_point.full_name + entry_point + if isinstance(entry_point, str) + else entry_point.full_name ) attack_step = entrypoint_full_name.split(":")[-1] - asset_name = ( - entrypoint_full_name.removesuffix(":" + attack_step) + asset_name = entrypoint_full_name.removesuffix( + ":" + attack_step ) asset = self.model.get_asset_by_name(asset_name) assert asset, "Asset does not exist" self.add_entrypoint_connection( - attack_step, - attacker_item, - self._asset_id_to_item[asset.id] + attack_step, attacker_item, self._asset_id_to_item[asset.id] ) for goal in agent_info.goals: @@ -549,51 +547,31 @@ def draw_model(self): goal if isinstance(goal, str) else goal.full_name ) attack_step = goal_full_name.split(":")[-1] - asset_name = ( - goal_full_name.removesuffix(":" + attack_step) - ) + asset_name = goal_full_name.removesuffix(":" + attack_step) asset = self.model.get_asset_by_name(asset_name) assert asset, "Asset does not exist" self.add_goal_connection( - attack_step, - attacker_item, - self._asset_id_to_item[asset.id] + attack_step, attacker_item, self._asset_id_to_item[asset.id] ) + # based on connectionType use attacker or + # add_association_connection -# based on connectionType use attacker or -# add_association_connection - - def add_association_connection( - self, - start_item, - end_item, - fieldname: str - ): + def add_association_connection(self, start_item, end_item, fieldname: str): """Add associations to the scene""" - connection = AssociationConnectionItem( - fieldname, start_item, end_item, self - ) + connection = AssociationConnectionItem(fieldname, start_item, end_item, self) self.addItem(connection) connection.restore_labels() connection.update_path() return connection - def add_entrypoint_connection( - self, - attack_step_name, - attacker_item, - asset_item - ): + def add_entrypoint_connection(self, attack_step_name, attacker_item, asset_item): """Add attacker entrypoints to the scene""" connection = EntrypointConnectionItem( - attack_step_name, - attacker_item, - asset_item, - self + attack_step_name, attacker_item, asset_item, self ) self.addItem(connection) @@ -601,19 +579,11 @@ def add_entrypoint_connection( connection.update_path() return connection - def add_goal_connection( - self, - attack_step_name, - attacker_item, - asset_item - ): + def add_goal_connection(self, attack_step_name, attacker_item, asset_item): """Add attacker goals to the scene""" connection = GoalConnectionItem( - attack_step_name, - attacker_item, - asset_item, - self + attack_step_name, attacker_item, asset_item, self ) self.addItem(connection) @@ -621,12 +591,7 @@ def add_goal_connection( connection.update_path() return connection - - def recreate_asset( - self, - asset_item: AssetItem, - position: QPointF - ) -> AssetItem: + def recreate_asset(self, asset_item: AssetItem, position: QPointF) -> AssetItem: """Rebuild existing asset and add to model and scene""" # Create new asset from the old asset object # and add it to the model @@ -643,15 +608,13 @@ def recreate_asset( return asset_item def create_asset( - self, asset_type: str, position: QPointF, name=None, asset_id=None - ) -> AssetItem: + self, asset_type: str, position: QPointF, name=None, asset_id=None + ) -> AssetItem: """Add new asset to model and to scene""" new_asset = self.model.add_asset( asset_type=asset_type, name=name, asset_id=asset_id ) - new_asset_item = self.asset_factory.create_asset_item( - new_asset, position - ) + new_asset_item = self.asset_factory.create_asset_item(new_asset, position) self.addItem(new_asset_item) print("Added asset item", new_asset_item, "to scene", self) self._asset_id_to_item[new_asset.id] = new_asset_item @@ -676,9 +639,7 @@ def remove_association(self, association_item: AssociationConnectionItem): def create_attacker(self, position, name, entry_points=None): """Add new attacker to the model and scene""" - new_item = self.asset_factory.create_attacker_item( - name, position, entry_points - ) + new_item = self.asset_factory.create_attacker_item(name, position, entry_points) self.attacker_items.append(new_item) self.addItem(new_item) return new_item @@ -691,7 +652,8 @@ def remove_entrypoint(self, entrypoint_item: EntrypointConnectionItem): """Remove attacker entrypoint and entrypoint item""" full_name = ( - entrypoint_item.asset_item.asset.name + ":" + entrypoint_item.asset_item.asset.name + + ":" + entrypoint_item.attack_step_name ) @@ -702,10 +664,7 @@ def remove_entrypoint(self, entrypoint_item: EntrypointConnectionItem): def remove_goal(self, goal_item: GoalConnectionItem): """Remove attacker goal and goal item""" - full_name = ( - goal_item.asset_item.asset.name + ":" - + goal_item.attack_step_name - ) + full_name = goal_item.asset_item.asset.name + ":" + goal_item.attack_step_name print("Remove goal", goal_item.attack_step_name) try: @@ -737,7 +696,7 @@ def delete_assets(self, selected_assets: list[AssetItem]): def containerize_assets(self, selected_assets: list[AssetItem]): print("Containerization of assets requested..") - command = ContainerizeAssetsCommand(self,selected_assets) + command = ContainerizeAssetsCommand(self, selected_assets) self.undo_stack.push(command) def decontainerize_assets(self, currently_selected_container: AssetsContainer): @@ -746,12 +705,14 @@ def decontainerize_assets(self, currently_selected_container: AssetsContainer): available_connections_in_item: list[IConnectionItem] = [] for item_entry in currently_selected_container.containerized_assets_list: - item: AssetItem = item_entry['item'] - original_position_of_item = current_position_of_container + item_entry['offset'] + item: AssetItem = item_entry["item"] + original_position_of_item = ( + current_position_of_container + item_entry["offset"] + ) self.addItem(item) item.setPos(original_position_of_item) - if hasattr(item, 'connections'): + if hasattr(item, "connections"): available_connections_in_item.extend(item.connections.copy()) # Restore connections @@ -766,22 +727,22 @@ def expand_container(self, currently_selected_container): """Expand container, move the assets out of it""" contained_item_for_bounding_rect_calc = [] - #copied below logic from containerize_assetsCommandUndo + # copied below logic from containerize_assetsCommandUndo - current_centroid_position_of_container = \ - currently_selected_container.scenePos() + current_centroid_position_of_container = currently_selected_container.scenePos() for item_entry in currently_selected_container.containerized_assets_list: - item = item_entry['item'] - offset_position_of_item = \ - current_centroid_position_of_container + item_entry['offset'] + item = item_entry["item"] + offset_position_of_item = ( + current_centroid_position_of_container + item_entry["offset"] + ) self.addItem(item) item.setPos(offset_position_of_item) # Update connections so association lines are visible properly self.update_connections(item) - #Store the item in a list for later bounding rect calculation + # Store the item in a list for later bounding rect calculation contained_item_for_bounding_rect_calc.append(item) # # Restore connections - Avoiding this because container and asset @@ -791,22 +752,24 @@ def expand_container(self, currently_selected_container): # connection.restore_labels() # connection.update_path() - rectangle_bounding_all_containerized_assets = self\ - .calc_surrounding_rect_for_grouped_assets_in_container( + rectangle_bounding_all_containerized_assets = ( + self.calc_surrounding_rect_for_grouped_assets_in_container( contained_item_for_bounding_rect_calc ) + ) self.container_box = AssetsContainerRectangleBox( rectangle_bounding_all_containerized_assets ) self.addItem(self.container_box) - self.container_box.associatied_compressed_container = \ + self.container_box.associatied_compressed_container = ( currently_selected_container + ) # self.removeItem(currently_selected_container) currently_selected_container.setVisible(False) - #MAKE COMPRESSED CONTAINER BOX HEADER FOR EXPANDED CONTAINER BOX - START + # MAKE COMPRESSED CONTAINER BOX HEADER FOR EXPANDED CONTAINER BOX - START # currently_selected_container.setVisible(True) # containerBoxRect = self.container_box.rect() @@ -815,14 +778,15 @@ def expand_container(self, currently_selected_container): # currentlySelectedContainerWidth = containerBoxRect.width() # currently_selected_container.setScale(currentlySelectedContainerWidth / currently_selected_container.boundingRect().width()) - #MAKE COMPRESSED CONTAINER BOX HEADER FOR EXPANDED CONTAINER BOX - END + # MAKE COMPRESSED CONTAINER BOX HEADER FOR EXPANDED CONTAINER BOX - END - def compress_container(self,currently_selected_container_box): - compressed_container = \ + def compress_container(self, currently_selected_container_box): + compressed_container = ( currently_selected_container_box.associatied_compressed_container + ) current_centroid_position_of_container = compressed_container.scenePos() for item_entry in compressed_container.containerized_assets_list: - item = item_entry['item'] + item = item_entry["item"] item.setPos(current_centroid_position_of_container) self.update_connections(item) compressed_container.setVisible(True) @@ -830,10 +794,10 @@ def compress_container(self,currently_selected_container_box): self.removeItem(currently_selected_container_box) def serialize_entrypoints( - self, - entrypoints: list[IConnectionItem], - selected_asset_ids: set[int], - selected_attacker_names: set[str] + self, + entrypoints: list[IConnectionItem], + selected_asset_ids: set[int], + selected_attacker_names: set[str], ): """Serialize selected attacker entrypoints""" @@ -845,8 +809,8 @@ def serialize_entrypoints( # If entry points both_items_selected = ( - conn.asset_item.asset.id in selected_asset_ids and - conn.attacker_item.name in selected_attacker_names + conn.asset_item.asset.id in selected_asset_ids + and conn.attacker_item.name in selected_attacker_names ) if not both_items_selected: continue @@ -855,7 +819,7 @@ def serialize_entrypoints( ( conn.attacker_item.name, conn.asset_item.asset.id, - conn.attack_step_name + conn.attack_step_name, ) ) @@ -867,14 +831,13 @@ def serialize_graphics_items(self, items: list[ItemBase], cut_intended): # Set of selected item IDs selected_attacker_names = { - item.name for item - in items if isinstance(item, AttackerItem) + item.name for item in items if isinstance(item, AttackerItem) } selected_asset_ids = { - item.asset.id for item - in items if isinstance(item, AssetItem) - and item.asset.id is not None + item.asset.id + for item in items + if isinstance(item, AssetItem) and item.asset.id is not None } for item in items: @@ -885,34 +848,31 @@ def serialize_graphics_items(self, items: list[ItemBase], cut_intended): item_id = item.name item_details = { - 'title': item.title, - 'id': item_id, - 'position': (item.pos().x(), item.pos().y()), + "title": item.title, + "id": item_id, + "position": (item.pos().x(), item.pos().y()), } if isinstance(item, AttackerItem): - item_details['type'] = "attacker" - item_details['entrypoints'] = self.serialize_entrypoints( + item_details["type"] = "attacker" + item_details["entrypoints"] = self.serialize_entrypoints( item.connections, selected_asset_ids, selected_attacker_names ) elif isinstance(item, AssetItem): - item_details['type'] = "asset" - item_details['properties'] = next( - iter(item.asset._to_dict().values()) - ) + item_details["type"] = "asset" + item_details["properties"] = next(iter(item.asset._to_dict().values())) serialized_items.append(item_details) serialized_data = pickle.dumps(serialized_items) - base64_serialized_data = \ - base64.b64encode(serialized_data).decode('utf-8') + base64_serialized_data = base64.b64encode(serialized_data).decode("utf-8") return base64_serialized_data def deserialize_graphics_items(self, asset_text): # Fix padding if necessary - I was getting padding error padding_needed = len(asset_text) % 4 if padding_needed: - asset_text += '=' * (4 - padding_needed) + asset_text += "=" * (4 - padding_needed) serialized_data = base64.b64decode(asset_text) deserialized_data = pickle.loads(serialized_data) @@ -962,8 +922,8 @@ def show_connection_item_context_menu(self, position, connection_item): self.delete_connection(connection_item) def show_assets_container_context_menu( - self, position, currently_selected_container - ): + self, position, currently_selected_container + ): print("Assets Container Context menu activated") menu = QMenu() assets_ungroup_action = QAction("Ungroup Asset(s)", self) @@ -979,12 +939,11 @@ def show_assets_container_context_menu( self.expand_container(currently_selected_container) def show_assets_container_box_context_menu( - self, position,currently_selected_container_box - ): + self, position, currently_selected_container_box + ): print("Assets Container Box Context menu activated") menu = QMenu() - assets_compress_container_box_action = \ - QAction("Compress Container Box", self) + assets_compress_container_box_action = QAction("Compress Container Box", self) menu.addAction(assets_compress_container_box_action) action = menu.exec(position) @@ -1026,25 +985,23 @@ def show_items_details(self): self.main_window.update_asset_relations_window(None) else: - self.main_window.item_details_window\ - .update_item_details_window(None) + self.main_window.item_details_window.update_item_details_window(None) self.main_window.update_properties_window(None) self.main_window.update_attack_steps_window(None) self.main_window.update_asset_relations_window(None) def calc_surrounding_rect_for_grouped_assets_in_container( - self, - contained_item_for_bounding_rect_calc - ): + self, contained_item_for_bounding_rect_calc + ): """Calculate the surrounding rect for assets in container""" if not contained_item_for_bounding_rect_calc: return QRectF() - min_x = float('inf') - max_x = float('-inf') - min_y = float('inf') - max_y = float('-inf') + min_x = float("inf") + max_x = float("-inf") + min_y = float("inf") + max_y = float("-inf") margin = 10.0 for item in contained_item_for_bounding_rect_calc: # Get the item's bounding rectangle in scene coordinates @@ -1065,6 +1022,6 @@ def calc_surrounding_rect_for_grouped_assets_in_container( return QRectF(QPointF(min_x, min_y), QPointF(max_x, max_y)) def update_connections(self, item: AssetItem): - if hasattr(item, 'connections'): + if hasattr(item, "connections"): for connection in item.connections: connection.update_path() diff --git a/mal_gui/model_view.py b/mal_gui/model_view.py index a0fdd4c..41382be 100644 --- a/mal_gui/model_view.py +++ b/mal_gui/model_view.py @@ -16,11 +16,11 @@ def __init__(self, scene, main_window): def zoomIn(self): """Overrides base""" - self.zoom(1.5) # Akash: This value need to discuss with Andrei + self.zoom(1.5) # Akash: This value need to discuss with Andrei def zoomOut(self): """Overrides base""" - self.zoom(1 / 1.5) # Akash: This value need to discuss with Andrei + self.zoom(1 / 1.5) # Akash: This value need to discuss with Andrei def wheelEvent(self, event): """Overrides base""" diff --git a/mal_gui/object_explorer/__init__.py b/mal_gui/object_explorer/__init__.py index 4d0d563..0e04e3f 100644 --- a/mal_gui/object_explorer/__init__.py +++ b/mal_gui/object_explorer/__init__.py @@ -4,8 +4,4 @@ from .attacker_item import AttackerItem from .editable_text_item import EditableTextItem -__all__ = [ - "AssetItem", - "AssetFactory", - "EditableTextItem" -] +__all__ = ["AssetItem", "AssetFactory", "EditableTextItem"] diff --git a/mal_gui/object_explorer/asset_factory.py b/mal_gui/object_explorer/asset_factory.py index 14d2394..553c224 100644 --- a/mal_gui/object_explorer/asset_factory.py +++ b/mal_gui/object_explorer/asset_factory.py @@ -8,14 +8,12 @@ from .attacker_item import AttackerItem if TYPE_CHECKING: - from maltoolbox.model import ModelAsset, AttackerAttachment + from maltoolbox.model import ModelAsset -AssetInfo = namedtuple( - 'AssetInfo', ['asset_type', 'asset_name', 'asset_image'] -) +AssetInfo = namedtuple("AssetInfo", ["asset_type", "asset_name", "asset_image"]) -class AssetFactory(): +class AssetFactory: def __init__(self, parent=None): self.asset_registry: dict[str, list[AssetInfo]] = {} @@ -31,13 +29,10 @@ def add_key_value_to_asset_registry(self, key, value): def register_asset(self, asset_name, image_path): self.add_key_value_to_asset_registry( - asset_name, - AssetInfo(asset_name, asset_name, image_path) + asset_name, AssetInfo(asset_name, asset_name, image_path) ) - def create_asset_item( - self, asset: ModelAsset, pos: QPointF - ): + def create_asset_item(self, asset: ModelAsset, pos: QPointF): asset_type = asset.lg_asset.name asset_info: AssetInfo = self.asset_registry[asset_type][0] requested_item = AssetItem(asset, asset_info.asset_image) @@ -48,19 +43,13 @@ def create_asset_item( requested_item.build() return requested_item - def create_attacker_item( - self, name: str, pos: QPointF, entry_points=None - ): - asset_type = 'Attacker' + def create_attacker_item(self, name: str, pos: QPointF, entry_points=None): + asset_type = "Attacker" asset_info: AssetInfo = self.asset_registry[asset_type][0] - requested_item = AttackerItem( - name, asset_info.asset_image, entry_points - ) + requested_item = AttackerItem(name, asset_info.asset_image, entry_points) requested_item.setPos(pos) - requested_item.type_text_item.setPlainText( - name or "Unnamed Attacker" - ) + requested_item.type_text_item.setPlainText(name or "Unnamed Attacker") requested_item.build() return requested_item diff --git a/mal_gui/object_explorer/asset_item.py b/mal_gui/object_explorer/asset_item.py index 65e8a78..a8466c8 100644 --- a/mal_gui/object_explorer/asset_item.py +++ b/mal_gui/object_explorer/asset_item.py @@ -1,22 +1,22 @@ from __future__ import annotations from typing import TYPE_CHECKING -from PySide6.QtCore import Qt from .item_base import ItemBase if TYPE_CHECKING: from maltoolbox.model import ModelAsset + class AssetItem(ItemBase): # Starting Sequence Id with normal start at 100 (randomly taken) def __init__( - self, - asset: ModelAsset, - image_path: str, - parent=None, - ): + self, + asset: ModelAsset, + image_path: str, + parent=None, + ): print("Create Asset item with parent", parent) self.asset = asset @@ -24,7 +24,6 @@ def __init__( super().__init__(asset.lg_asset.name, image_path, parent) - def update_name(self): super().update_name() self.asset.name = self.title @@ -50,8 +49,8 @@ def set_item_attribute_value(self, attr_name, new_value_str) -> None: def serialize(self): return { - 'title': self.title, - 'image_path': self.image_path, - 'type': 'asset', - 'object': self.asset + "title": self.title, + "image_path": self.image_path, + "type": "asset", + "object": self.asset, } diff --git a/mal_gui/object_explorer/attacker_item.py b/mal_gui/object_explorer/attacker_item.py index bf25e2c..42f713a 100644 --- a/mal_gui/object_explorer/attacker_item.py +++ b/mal_gui/object_explorer/attacker_item.py @@ -1,4 +1,3 @@ -from typing import Any from PySide6.QtGui import QColor from PySide6.QtCore import QTimer from PySide6.QtWidgets import QGraphicsItem @@ -16,17 +15,18 @@ policies.TTCSoftMinAttacker, ] + class AttackerItem(ItemBase): # Starting Sequence Id with normal start at 100 (randomly taken) def __init__( - self, - name: str, - image_path: str, - entry_points=None, - goals=None, - parent=None, - ): + self, + name: str, + image_path: str, + entry_points=None, + goals=None, + parent=None, + ): self.entry_points: list[str] = entry_points or [] self.policy = policies.PassiveAgent @@ -40,7 +40,7 @@ def __init__( self.timer.timeout.connect(self.update_status_color) self.timer.start(500) - super().__init__('Attacker', image_path, parent) + super().__init__("Attacker", image_path, parent) def update_type_text_item_position(self): super().update_type_text_item_position() @@ -102,8 +102,8 @@ def itemChange(self, change, value): def serialize(self): return { - 'title': self.title, - 'image_path': self.image_path, - 'type': 'asset', - 'object': self.entry_points + "title": self.title, + "image_path": self.image_path, + "type": "asset", + "object": self.entry_points, } diff --git a/mal_gui/object_explorer/editable_text_item.py b/mal_gui/object_explorer/editable_text_item.py index bd5c7cb..df5eead 100644 --- a/mal_gui/object_explorer/editable_text_item.py +++ b/mal_gui/object_explorer/editable_text_item.py @@ -2,6 +2,7 @@ from PySide6.QtGui import QFont, QTextCursor from PySide6.QtWidgets import QGraphicsTextItem + class EditableTextItem(QGraphicsTextItem): lostFocus = Signal() diff --git a/mal_gui/object_explorer/item_base.py b/mal_gui/object_explorer/item_base.py index 52b7097..7cca07a 100644 --- a/mal_gui/object_explorer/item_base.py +++ b/mal_gui/object_explorer/item_base.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING from abc import abstractmethod -from PySide6.QtCore import QRectF, Qt, QPointF, QSize, QSizeF, QTimer +from PySide6.QtCore import QRectF, Qt, QPointF, QSize, QSizeF from PySide6.QtGui import ( QPixmap, QFont, @@ -12,24 +12,23 @@ QPainterPath, QFontMetrics, QLinearGradient, - QImage + QImage, ) -from PySide6.QtWidgets import QGraphicsItem +from PySide6.QtWidgets import QGraphicsItem from .editable_text_item import EditableTextItem if TYPE_CHECKING: - from maltoolbox.model import ModelAsset from ..connection_item import IConnectionItem -class ItemBase(QGraphicsItem): +class ItemBase(QGraphicsItem): def __init__( - self, - title: str, - image_path: str, - parent=None, - ): + self, + title: str, + image_path: str, + parent=None, + ): super().__init__(parent) self.setZValue(1) # rect items are on top @@ -37,9 +36,7 @@ def __init__( self.title = title self.image_path = image_path - self.image = self.load_image_with_quality( - self.image_path, QSize(512, 512) - ) + self.image = self.load_image_with_quality(self.image_path, QSize(512, 512)) self.setFlags( QGraphicsItem.ItemIsSelectable @@ -59,8 +56,8 @@ def __init__( self.height = 70 self.size = QRectF(-self.width / 2, -self.height / 2, self.width, self.height) - self.asset_type_background_color = QColor(0, 200, 0) #Green - self.asset_name_background_color = QColor(20, 20, 20, 200) # Gray + self.asset_type_background_color = QColor(0, 200, 0) # Green + self.asset_name_background_color = QColor(20, 20, 20, 200) # Gray self.icon_path = None self.icon_visible = True @@ -72,7 +69,7 @@ def __init__( self.horizontal_margin = 15 # Horizontal margin self.vertical_margin = 15 # Vertical margin - self.status_color = QColor(0, 255, 0) + self.status_color = QColor(0, 255, 0) self.build() @@ -114,21 +111,32 @@ def paint(self, painter, option, widget=None): # resizedImageIcon = self.image.scaled(targetIconSize, Qt.KeepAspectRatio, Qt.SmoothTransformation) resizedImageIcon = self.image - # Calculate the position and size for the icon background - iconRect = QRectF(-self.width / 2 + 10, -self.height / 2 + 10, targetIconSize.width(), targetIconSize.height()) + iconRect = QRectF( + -self.width / 2 + 10, + -self.height / 2 + 10, + targetIconSize.width(), + targetIconSize.height(), + ) margin = 5 # Margin around the icon # Draw the background for the icon with additional margin backgroundRect = QRectF( iconRect.topLeft() - QPointF(margin, margin), - QSizeF(targetIconSize.width() + 2 * margin, targetIconSize.height() + 2 * margin) + QSizeF( + targetIconSize.width() + 2 * margin, + targetIconSize.height() + 2 * margin, + ), ) painter.setBrush(Qt.white) # Set the brush color to white - painter.drawRect(backgroundRect.toRect()) # Convert QRectF to QRect and draw the white background rectangle + painter.drawRect( + backgroundRect.toRect() + ) # Convert QRectF to QRect and draw the white background rectangle # Draw the resized icon on top of the white background - painter.drawPixmap(iconRect.toRect(), resizedImageIcon) # Convert QRectF to QRect and draw the resized icon + painter.drawPixmap( + iconRect.toRect(), resizedImageIcon + ) # Convert QRectF to QRect and draw the resized icon # Draw the highlight if selected if self.isSelected(): @@ -149,9 +157,7 @@ def itemChange(self, change, value): def mouseDoubleClickEvent(self, event): if event.button() == Qt.LeftButton: - self.type_text_item.setTextInteractionFlags( - Qt.TextEditorInteraction - ) + self.type_text_item.setTextInteractionFlags(Qt.TextEditorInteraction) self.type_text_item.setFocus() # Select all text when activated self.type_text_item.select_all_text() @@ -168,9 +174,8 @@ def mousePressEvent(self, event): """Overrides base method""" self.initial_position = self.pos() - if ( - self.type_text_item.hasFocus() - and not self.type_text_item.contains(event.pos()) + if self.type_text_item.hasFocus() and not self.type_text_item.contains( + event.pos() ): self.type_text_item.clearFocus() elif not self.type_text_item.contains(event.pos()): @@ -195,12 +200,7 @@ def build(self): # Draw the background of the node self.path = QPainterPath() self.path.addRoundedRect( - -fixed_width / 2, - -fixed_height / 2, - fixed_width, - fixed_height, - 6, - 6 + -fixed_width / 2, -fixed_height / 2, fixed_width, fixed_height, 6, 6 ) self.title_bg_path = QPainterPath() @@ -210,18 +210,13 @@ def build(self): fixed_width, title_font.pointSize() + 2 * self.vertical_margin, 6, - 6 + 6, ) # Draw status path self.status_path.setFillRule(Qt.WindingFill) self.status_path.addRoundedRect( - fixed_width / 2 - 12, - -fixed_height / 2 + 2, - 10, - 10, - 2, - 2 + fixed_width / 2 - 12, -fixed_height / 2 + 2, 10, 10, 2, 2 ) # Center title in the upper half @@ -232,7 +227,7 @@ def build(self): # Center vertically within its section -fixed_height / 2 + self.vertical_margin + title_font_metrics.ascent(), title_font, - self.title_text + self.title_text, ) # Set the font and default color for type_text_item @@ -243,13 +238,11 @@ def build(self): self.update_type_text_item_position() # Connect lostFocus signal to update position when text loses focus - self.type_text_item.lostFocus.connect( - self.update_type_text_item_position) + self.type_text_item.lostFocus.connect(self.update_type_text_item_position) # self.widget.move(-self.widget.size().width() / 2, # fixed_height / 2 - self.widget.size().height() + 5) - def add_connection(self, connection): self.connections.append(connection) @@ -257,7 +250,6 @@ def remove_connection(self, connection): if connection in self.connections: self.connections.remove(connection) - def update_type_text_item_position(self): # to update the position of the type_text_item so that it # remains centered within the lower half of the node @@ -268,10 +260,12 @@ def update_type_text_item_position(self): title_font_metrics = QFontMetrics(QFont("Arial", pointSize=12)) # Calculate the new position for type_text_item - type_text_item_pos_x = \ + type_text_item_pos_x = ( -type_font_metrics.horizontalAdvance(self.type_text_item.toPlainText()) / 2 - type_text_item_pos_y = \ + ) + type_text_item_pos_y = ( -fixed_height / 2 + title_font_metrics.height() + 2 * self.vertical_margin + ) # Update position self.type_text_item.setPos(type_text_item_pos_x, type_text_item_pos_y) @@ -285,8 +279,7 @@ def update_name(self): associated_scene = self.type_text_item.scene() if associated_scene: print("Asset Name Changed by user") - associated_scene.main_window\ - .update_childs_in_object_explorer_signal.emit() + associated_scene.main_window.update_childs_in_object_explorer_signal.emit() def setIcon(self, icon_path=None): """Overrides base method""" @@ -304,10 +297,7 @@ def load_image_with_quality(self, path, size): image = QImage(path) if not image.isNull(): return QPixmap.fromImage( - image.scaled( - size, Qt.KeepAspectRatio, - Qt.SmoothTransformation - ) + image.scaled(size, Qt.KeepAspectRatio, Qt.SmoothTransformation) ) return QPixmap() diff --git a/mal_gui/undo_redo_commands/__init__.py b/mal_gui/undo_redo_commands/__init__.py index 8fe097b..f708515 100644 --- a/mal_gui/undo_redo_commands/__init__.py +++ b/mal_gui/undo_redo_commands/__init__.py @@ -1,9 +1,7 @@ from .containerize_assets_command import ContainerizeAssetsCommand from .copy_command import CopyCommand -from .create_association_connection_command import \ - CreateAssociationConnectionCommand -from .create_entrypoint_connection_command import \ - CreateEntrypointConnectionCommand +from .create_association_connection_command import CreateAssociationConnectionCommand +from .create_entrypoint_connection_command import CreateEntrypointConnectionCommand from .cut_command import CutCommand from .delete_command import DeleteCommand from .delete_connection_command import DeleteConnectionCommand diff --git a/mal_gui/undo_redo_commands/containerize_assets_command.py b/mal_gui/undo_redo_commands/containerize_assets_command.py index 9cd21b4..64ff413 100644 --- a/mal_gui/undo_redo_commands/containerize_assets_command.py +++ b/mal_gui/undo_redo_commands/containerize_assets_command.py @@ -11,24 +11,27 @@ if TYPE_CHECKING: from ..connection_item import IConnectionItem + class ContainerizeAssetsCommand(QUndoCommand): def __init__(self, scene, items, parent=None): super().__init__(parent) self.scene = scene - self.items: list[AssetItem] = \ - [item for item in items if item.asset_type != 'Attacker'] + self.items: list[AssetItem] = [ + item for item in items if item.asset_type != "Attacker" + ] self.connections: list[IConnectionItem] = [] - self.centroid = QPointF(0,0) + self.centroid = QPointF(0, 0) - #Timer specific values for animation - self.animation_duration = 2000 # in milliseconds - self.animation_timer_interval = 20 # in milliseconds - self.num_steps_for_animation = \ + # Timer specific values for animation + self.animation_duration = 2000 # in milliseconds + self.animation_timer_interval = 20 # in milliseconds + self.num_steps_for_animation = ( self.animation_duration // self.animation_timer_interval + ) # Save connections of all items for item in self.items: - if hasattr(item, 'connections'): + if hasattr(item, "connections"): self.connections.extend(item.connections.copy()) self.new_assets_container = AssetsContainer( @@ -36,7 +39,7 @@ def __init__(self, scene, items, parent=None): "ContainerName", image_path("assetContainer.png"), image_path("assetContainerPlusSymbol.png"), - image_path("assetContainerMinusSymbol.png") + image_path("assetContainerMinusSymbol.png"), ) self.new_assets_container.build() @@ -49,22 +52,20 @@ def redo(self): center_y = sum(y_coords_list_of_assets) / len(y_coords_list_of_assets) self.centroid = QPointF(center_x, center_y) - #Display the container at centroid location + # Display the container at centroid location self.scene.addItem(self.new_assets_container) self.new_assets_container.setPos(self.centroid) for item in self.items: - if isinstance(item,AssetItem): - self.move_item_from_current_position_to_centroid( - item, self.centroid - ) + if isinstance(item, AssetItem): + self.move_item_from_current_position_to_centroid(item, self.centroid) def undo(self): """Undo containerization""" # Add items back to the scene for item_entry in self.new_assets_container.containerized_assets_list: - item = item_entry['item'] - original_position_of_item = self.centroid + item_entry['offset'] + item = item_entry["item"] + original_position_of_item = self.centroid + item_entry["offset"] self.scene.addItem(item) item.setPos(original_position_of_item) @@ -76,19 +77,21 @@ def undo(self): self.scene.removeItem(self.new_assets_container) - def update_item_position(self, item, start_pos, end_pos, is_redo): + def update_item_position(self, item, start_pos, end_pos, is_redo): if item.step_counter >= self.num_steps_for_animation: item.timer.timeout.disconnect() item.timer.stop() if is_redo: - #Store item and offset because items moving towards centroid + # Store item and offset because items moving towards centroid self.new_assets_container.containerized_assets_list.append( - {'item': item, 'offset': item.offsetFromCentroid}) + {"item": item, "offset": item.offsetFromCentroid} + ) item.setPos(self.centroid) self.update_connections(item) - self.new_assets_container.item_moved = \ + self.new_assets_container.item_moved = ( self.update_items_positions_relative_to_container + ) else: self.new_assets_container.containerized_assets_list.clear() return @@ -100,23 +103,23 @@ def update_item_position(self, item, start_pos, end_pos, is_redo): item.step_counter += 1 - def move_item_from_current_position_to_centroid(self,item,centroid): + def move_item_from_current_position_to_centroid(self, item, centroid): item.timer = QTimer() item.step_counter = 0 item.offsetFromCentroid = item.scenePos() - centroid item.timer.timeout.connect( lambda: self.update_item_position( - item, item.scenePos(),centroid,is_redo=True + item, item.scenePos(), centroid, is_redo=True ) ) item.timer.start(self.animation_timer_interval) def update_connections(self, item: AssetItem): - if hasattr(item, 'connections'): + if hasattr(item, "connections"): for connection in item.connections: connection.update_path() def update_items_positions_relative_to_container(self): for item_entry in self.new_assets_container.containerized_assets_list: - item = item_entry['item'] + item = item_entry["item"] item.setPos(self.new_assets_container.pos()) diff --git a/mal_gui/undo_redo_commands/copy_command.py b/mal_gui/undo_redo_commands/copy_command.py index 705bcf8..d869fed 100644 --- a/mal_gui/undo_redo_commands/copy_command.py +++ b/mal_gui/undo_redo_commands/copy_command.py @@ -6,8 +6,11 @@ from ..model_scene import ModelScene from ..object_explorer.asset_item import AssetItem + class CopyCommand(QUndoCommand): - def __init__(self, scene: ModelScene, items: list[AssetItem], clipboard, parent=None): + def __init__( + self, scene: ModelScene, items: list[AssetItem], clipboard, parent=None + ): super().__init__(parent) self.scene = scene self.items = items @@ -15,8 +18,9 @@ def __init__(self, scene: ModelScene, items: list[AssetItem], clipboard, parent= def redo(self): self.cut_item_flag = False - serialized_data = \ - self.scene.serialize_graphics_items(self.items, self.cut_item_flag) + serialized_data = self.scene.serialize_graphics_items( + self.items, self.cut_item_flag + ) self.clipboard.clear() self.clipboard.setText(serialized_data) diff --git a/mal_gui/undo_redo_commands/create_association_connection_command.py b/mal_gui/undo_redo_commands/create_association_connection_command.py index ddaa798..7039cd0 100644 --- a/mal_gui/undo_redo_commands/create_association_connection_command.py +++ b/mal_gui/undo_redo_commands/create_association_connection_command.py @@ -6,19 +6,19 @@ from ..model_scene import ModelScene from ..model_scene import AssetItem -class CreateAssociationConnectionCommand(QUndoCommand): +class CreateAssociationConnectionCommand(QUndoCommand): def __init__( self, scene: ModelScene, start_item: AssetItem, end_item: AssetItem, field_name, - parent=None + parent=None, ): super().__init__(parent) - self.scene = scene + self.scene = scene self.start_item = start_item self.end_item = end_item self.fieldname = field_name diff --git a/mal_gui/undo_redo_commands/create_entrypoint_connection_command.py b/mal_gui/undo_redo_commands/create_entrypoint_connection_command.py index 6cfd835..27c2cf5 100644 --- a/mal_gui/undo_redo_commands/create_entrypoint_connection_command.py +++ b/mal_gui/undo_redo_commands/create_entrypoint_connection_command.py @@ -6,6 +6,8 @@ if TYPE_CHECKING: from ..model_scene import ModelScene from ..object_explorer import AttackerItem, AssetItem + + class CreateEntrypointConnectionCommand(QUndoCommand): def __init__( self, @@ -13,7 +15,7 @@ def __init__( attacker_item: AttackerItem, asset_item: AssetItem, attack_step_name: str, - parent=None + parent=None, ): super().__init__(parent) self.scene = scene @@ -25,9 +27,7 @@ def __init__( def redo(self): """Create entrypoint for attacker""" self.connection = self.scene.add_entrypoint_connection( - self.attack_step_name, - self.attacker_item, - self.asset_item + self.attack_step_name, self.attacker_item, self.asset_item ) self.attacker_item.entry_points.append( self.asset_item.asset.name + ":" + self.attack_step_name diff --git a/mal_gui/undo_redo_commands/create_goal_connection_command.py b/mal_gui/undo_redo_commands/create_goal_connection_command.py index 34e4376..0e5b6a2 100644 --- a/mal_gui/undo_redo_commands/create_goal_connection_command.py +++ b/mal_gui/undo_redo_commands/create_goal_connection_command.py @@ -6,6 +6,8 @@ if TYPE_CHECKING: from ..model_scene import ModelScene from ..object_explorer import AttackerItem, AssetItem + + class CreateGoalConnectionCommand(QUndoCommand): def __init__( self, @@ -13,7 +15,7 @@ def __init__( attacker_item: AttackerItem, asset_item: AssetItem, attack_step_name: str, - parent=None + parent=None, ): super().__init__(parent) self.scene = scene @@ -25,9 +27,7 @@ def __init__( def redo(self): """Create entrypoint for attacker""" self.connection = self.scene.add_goal_connection( - self.attack_step_name, - self.attacker_item, - self.asset_item + self.attack_step_name, self.attacker_item, self.asset_item ) self.attacker_item.goals.append( self.asset_item.asset.name + ":" + self.attack_step_name diff --git a/mal_gui/undo_redo_commands/cut_command.py b/mal_gui/undo_redo_commands/cut_command.py index 80c6951..a30defb 100644 --- a/mal_gui/undo_redo_commands/cut_command.py +++ b/mal_gui/undo_redo_commands/cut_command.py @@ -7,14 +7,9 @@ from ..model_scene import ModelScene from ..connection_item import IConnectionItem + class CutCommand(QUndoCommand): - def __init__( - self, - scene: ModelScene, - items, - clipboard, - parent=None - ): + def __init__(self, scene: ModelScene, items, clipboard, parent=None): super().__init__(parent) self.scene = scene self.items = items @@ -23,14 +18,15 @@ def __init__( # Save connections of all items for item in self.items: - if hasattr(item, 'connections'): + if hasattr(item, "connections"): self.connections.extend(item.connections.copy()) def redo(self): """Perform cut command""" self.cut_item_flag = True serialized_data = self.scene.serialize_graphics_items( - self.items, self.cut_item_flag) + self.items, self.cut_item_flag + ) self.clipboard.clear() self.clipboard.setText(serialized_data) @@ -42,7 +38,7 @@ def redo(self): for item in self.items: self.scene.removeItem(item) - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() def undo(self): @@ -59,5 +55,5 @@ def undo(self): self.clipboard.clear() - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() diff --git a/mal_gui/undo_redo_commands/delete_command.py b/mal_gui/undo_redo_commands/delete_command.py index ef6ecd8..d88f9de 100644 --- a/mal_gui/undo_redo_commands/delete_command.py +++ b/mal_gui/undo_redo_commands/delete_command.py @@ -2,19 +2,19 @@ from typing import TYPE_CHECKING from PySide6.QtGui import QUndoCommand from ..object_explorer import AssetItem, AttackerItem -from ..connection_item import AssociationConnectionItem, EntrypointConnectionItem, GoalConnectionItem +from ..connection_item import ( + AssociationConnectionItem, + EntrypointConnectionItem, + GoalConnectionItem, +) if TYPE_CHECKING: from ..model_scene import ModelScene from ..connection_item import IConnectionItem + class DeleteCommand(QUndoCommand): - def __init__( - self, - scene: ModelScene, - items, - parent=None - ): + def __init__(self, scene: ModelScene, items, parent=None): super().__init__(parent) self.scene = scene self.items = items @@ -22,7 +22,7 @@ def __init__( # Save connections of all items for item in self.items: - if hasattr(item, 'connections'): + if hasattr(item, "connections"): self.connections.extend(item.connections.copy()) def redo(self): @@ -34,19 +34,26 @@ def redo(self): self.scene.removeItem(connection) if isinstance(connection, EntrypointConnectionItem): - step_full_name = connection.asset_item.asset.name + ":" + connection.attack_step_name + step_full_name = ( + connection.asset_item.asset.name + ":" + connection.attack_step_name + ) try: connection.attacker_item.entry_points.remove(step_full_name) except ValueError: - print(f"Entrypoint {step_full_name} not found in attacker {connection.attacker_item.name}") + print( + f"Entrypoint {step_full_name} not found in attacker {connection.attacker_item.name}" + ) if isinstance(connection, GoalConnectionItem): - step_full_name = connection.asset_item.asset.name + ":" + connection.attack_step_name + step_full_name = ( + connection.asset_item.asset.name + ":" + connection.attack_step_name + ) try: connection.attacker_item.goals.remove(step_full_name) except ValueError: - print(f"Entrypoint {step_full_name} not found in attacker {connection.attacker_item.name}") - + print( + f"Entrypoint {step_full_name} not found in attacker {connection.attacker_item.name}" + ) for item in self.items: if isinstance(item, AssetItem): @@ -55,7 +62,7 @@ def redo(self): if isinstance(item, AttackerItem): self.scene.remove_attacker(item) - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() def undo(self): @@ -70,12 +77,11 @@ def undo(self): # Restore connections for connection in self.connections: - if isinstance(connection, EntrypointConnectionItem): self.scene.add_entrypoint_connection( connection.attack_step_name, connection.attacker_item, - connection.asset_item + connection.asset_item, ) step_full_name = ( connection.asset_item.asset.name + ":" + connection.attack_step_name @@ -86,7 +92,7 @@ def undo(self): self.scene.add_goal_connection( connection.attack_step_name, connection.attacker_item, - connection.asset_item + connection.asset_item, ) step_full_name = ( connection.asset_item.asset.name + ":" + connection.attack_step_name @@ -97,12 +103,11 @@ def undo(self): self.scene.add_association_connection( connection.start_item, connection.end_item, - connection.right_fieldname + connection.right_fieldname, ) connection.start_item.asset.add_associated_assets( connection.right_fieldname, {connection.end_item.asset} ) - - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() diff --git a/mal_gui/undo_redo_commands/delete_connection_command.py b/mal_gui/undo_redo_commands/delete_connection_command.py index 65a723b..0ee5710 100644 --- a/mal_gui/undo_redo_commands/delete_connection_command.py +++ b/mal_gui/undo_redo_commands/delete_connection_command.py @@ -1,12 +1,17 @@ from __future__ import annotations from typing import TYPE_CHECKING from PySide6.QtGui import QUndoCommand -from ..connection_item import AssociationConnectionItem, EntrypointConnectionItem, GoalConnectionItem +from ..connection_item import ( + AssociationConnectionItem, + EntrypointConnectionItem, + GoalConnectionItem, +) if TYPE_CHECKING: from ..connection_item import IConnectionItem from ..model_scene import ModelScene + class DeleteConnectionCommand(QUndoCommand): def __init__(self, scene: ModelScene, item, parent=None): super().__init__(parent) @@ -34,19 +39,19 @@ def undo(self): self.connection = self.scene.add_association_connection( self.connection.start_item, self.connection.end_item, - self.connection.right_fieldname + self.connection.right_fieldname, ) elif isinstance(self.connection, EntrypointConnectionItem): self.connection = self.scene.add_entrypoint_connection( self.connection.attack_step_name, self.connection.attacker_item, - self.connection.asset_item + self.connection.asset_item, ) elif isinstance(self.connection, GoalConnectionItem): self.connection = self.scene.add_goal_connection( self.connection.attack_step_name, self.connection.attacker_item, - self.connection.asset_item + self.connection.asset_item, ) else: raise ValueError("Unknown connection type") diff --git a/mal_gui/undo_redo_commands/drag_drop_command.py b/mal_gui/undo_redo_commands/drag_drop_command.py index 9189239..f537ca0 100644 --- a/mal_gui/undo_redo_commands/drag_drop_command.py +++ b/mal_gui/undo_redo_commands/drag_drop_command.py @@ -7,13 +7,14 @@ if TYPE_CHECKING: from mal_gui.model_scene import ModelScene + class DragDropAssetCommand(QUndoCommand): def __init__( self, scene: ModelScene, asset_type: str, position: QPointF, - name = None, + name=None, parent=None, ): """ @@ -35,16 +36,14 @@ def redo(self): # If it is an 'actual' redo, we need to add the same asset again if self.item: # Create asset from previous deleted asset - self.item = self.scene.recreate_asset( - self.item, self.position - ) + self.item = self.scene.recreate_asset(self.item, self.position) else: # Create/add asset from scratch self.item = self.scene.create_asset( self.asset_type, self.position, self.name ) - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() def undo(self): @@ -77,13 +76,13 @@ def redo(self): self.item = self.scene.create_attacker( self.item.pos(), name=self.item.attacker.name, - attacker_id=self.item.attacker.id + attacker_id=self.item.attacker.id, ) else: # Create attacker from scratch - self.item = self.scene.create_attacker(self.position, 'Attacker') + self.item = self.scene.create_attacker(self.position, "Attacker") - #Update the Object Explorer when number of items change + # Update the Object Explorer when number of items change self.scene.main_window.update_childs_in_object_explorer_signal.emit() def undo(self): diff --git a/mal_gui/undo_redo_commands/move_command.py b/mal_gui/undo_redo_commands/move_command.py index fe66d2e..376e601 100644 --- a/mal_gui/undo_redo_commands/move_command.py +++ b/mal_gui/undo_redo_commands/move_command.py @@ -6,15 +6,16 @@ from ..object_explorer.asset_item import AssetItem from ..model_scene import ModelScene + class MoveCommand(QUndoCommand): def __init__( - self, - scene: ModelScene, - items: list, - start_positions, - end_positions, - parent=None - ): + self, + scene: ModelScene, + items: list, + start_positions, + end_positions, + parent=None, + ): super().__init__(parent) self.scene = scene self.items = items @@ -37,6 +38,6 @@ def undo(self): def update_connections(self, item: AssetItem): """Redraw connecting lines""" - if hasattr(item, 'connections'): + if hasattr(item, "connections"): for connection in item.connections: connection.update_path() diff --git a/mal_gui/undo_redo_commands/paste_command.py b/mal_gui/undo_redo_commands/paste_command.py index ec2d028..1d2d5ab 100644 --- a/mal_gui/undo_redo_commands/paste_command.py +++ b/mal_gui/undo_redo_commands/paste_command.py @@ -11,14 +11,9 @@ if TYPE_CHECKING: from ..model_scene import ModelScene + class PasteCommand(QUndoCommand): - def __init__( - self, - scene: ModelScene, - position, - clipboard, - parent=None - ): + def __init__(self, scene: ModelScene, position, clipboard, parent=None): super().__init__(parent) self.scene = scene @@ -37,25 +32,22 @@ def redo(self): print("\nPaste Redo is Called") serialized_data = self.clipboard.text() - deserialized_data = ( - self.scene.deserialize_graphics_items(serialized_data) - ) - print(json.dumps(deserialized_data, indent = 2)) + deserialized_data = self.scene.deserialize_graphics_items(serialized_data) + print(json.dumps(deserialized_data, indent=2)) # First pass: create items with new assetIds for data in deserialized_data: - - item_type = data['type'] - old_id = data['id'] - position_tuple = data['position'] + item_type = data["type"] + old_id = data["id"] + position_tuple = data["position"] position = QPointF(position_tuple[0], position_tuple[1]) if item_type == "attacker": new_item = self.scene.create_attacker(position) elif item_type == "asset": - asset_type = data['properties']['type'] - asset_name = data['properties']['name'] + asset_type = data["properties"]["type"] + asset_name = data["properties"]["name"] new_asset_id = None @@ -67,10 +59,7 @@ def redo(self): else: new_item = self.scene.create_asset( - asset_type, - position, - name=asset_name, - asset_id=new_asset_id + asset_type, position, name=asset_name, asset_id=new_asset_id ) else: @@ -80,10 +69,8 @@ def redo(self): # Adjust the position of all assetItems with offset values # Find the top-leftmost position among the items to be pasted - min_x = min( - item.pos().x() for item in self.original_id_to_item.values()) - min_y = min( - item.pos().y() for item in self.original_id_to_item.values()) + min_x = min(item.pos().x() for item in self.original_id_to_item.values()) + min_y = min(item.pos().y() for item in self.original_id_to_item.values()) top_left = QPointF(min_x, min_y) # Calculate the offset from the top-leftmost @@ -96,20 +83,18 @@ def redo(self): # Second pass: re-establish connections with new assetSequenceIds for data in deserialized_data: - item_type = data['type'] - old_id = data['id'] - position_tuple = data['position'] + item_type = data["type"] + old_id = data["id"] + position_tuple = data["position"] item = self.original_id_to_item[old_id] if isinstance(item, AssetItem): # Must be an asset - associated_assets = data['properties']['associated_assets'] + associated_assets = data["properties"]["associated_assets"] for fieldname, assets in associated_assets.items(): for asset_id in assets: right_item = self.original_id_to_item[asset_id] - item.asset.add_associated_assets( - fieldname, {right_item.asset} - ) + item.asset.add_associated_assets(fieldname, {right_item.asset}) con = self.scene.add_association_connection( item, right_item, fieldname ) @@ -117,21 +102,18 @@ def redo(self): elif isinstance(item, AttackerItem): # Add attacker entrypoints - for entrypoint in data['entrypoints']: - print(f'ENTRYPOINT: {entrypoint}') + for entrypoint in data["entrypoints"]: + print(f"ENTRYPOINT: {entrypoint}") old_start_id, old_end_id, label = entrypoint - new_attacker_item: AttackerItem = ( - self.original_id_to_item[old_start_id] - ) - new_asset_item: AssetItem = ( - self.original_id_to_item[old_end_id] - ) + new_attacker_item: AttackerItem = self.original_id_to_item[ + old_start_id + ] + new_asset_item: AssetItem = self.original_id_to_item[old_end_id] - new_connection = self.scene\ - .add_entrypoint_connection( - label, new_attacker_item, new_asset_item - ) + new_connection = self.scene.add_entrypoint_connection( + label, new_attacker_item, new_asset_item + ) self.pasted_entrypoints.append(new_connection) new_attacker_item.attacker.add_entry_point( diff --git a/pyproject.toml b/pyproject.toml index 903c2cf..c3f07fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,14 +9,14 @@ authors = [ {name="Sandor Berglund", email="sandorb@kth.se"}, ] readme = "README.md" -requires-python = ">=3.10,<=3.13" +requires-python = ">=3.10,<3.15" dependencies = [ "networkx==3.2.1", - "numpy==1.26.4", - "PySide6~=6.8.1", - "PySide6_Addons~=6.8.1", - "PySide6_Essentials~=6.8.1", - "shiboken6~=6.8.1", + "numpy==2.3.3", + "PySide6~=6.10.1", + "PySide6_Addons~=6.10.1", + "PySide6_Essentials~=6.10.1", + "shiboken6~=6.10.1", "mal-toolbox==2.*", "mal-simulator==2.*", "qt-material==2.14", diff --git a/tests/conftest.py b/tests/conftest.py index f9203ce..71cb09e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import pytest from PySide6.QtWidgets import QApplication + @pytest.fixture def lang_file_path(): return "tests/testdata/org.mal-lang.coreLang-1.0.0.mar" diff --git a/tests/test_main_window.py b/tests/test_main_window.py index d6137ba..f5f4829 100644 --- a/tests/test_main_window.py +++ b/tests/test_main_window.py @@ -1,7 +1,7 @@ import pytest -from PySide6.QtWidgets import QApplication, QMainWindow, QToolBar -from PySide6.QtCore import Qt, QPointF +from PySide6.QtWidgets import QMainWindow, QToolBar +from PySide6.QtCore import QPointF from maltoolbox.language import LanguageGraph from maltoolbox.model import Model @@ -22,6 +22,7 @@ def main_window(app, lang_file_path): # Initialization # ------------------------------------------------------------------- + def test_main_window_initialization(main_window): assert isinstance(main_window, QMainWindow) assert main_window.windowTitle() == "MAL GUI" @@ -47,9 +48,12 @@ def test_scene_is_model_scene(main_window): # Menu & Actions # ------------------------------------------------------------------- + def test_menu_bar_created(main_window): menu_bar = main_window.menuBar() - actions = [menu.title() for menu in menu_bar.findChildren(type(menu_bar.addMenu("tmp")))] + actions = [ + menu.title() for menu in menu_bar.findChildren(type(menu_bar.addMenu("tmp"))) + ] assert "&File" in actions assert "Edit" in actions @@ -70,6 +74,7 @@ def test_actions_exist(main_window): # Toolbar behavior # ------------------------------------------------------------------- + def test_zoom_actions(main_window): initial_zoom = main_window.view.zoom_factor @@ -92,6 +97,7 @@ def test_zoom_line_edit(main_window): # Scene reload / clearing # ------------------------------------------------------------------- + def test_clear_window(main_window): # Sanity: items exist initially assert main_window.scene is not None @@ -120,6 +126,7 @@ def test_load_scene_recreates_components(app, lang_file_path): # Object explorer update signal # ------------------------------------------------------------------- + def test_update_explorer_signal(main_window): # Should not raise main_window.update_childs_in_object_explorer_signal.emit() @@ -129,6 +136,7 @@ def test_update_explorer_signal(main_window): # Theme handling # ------------------------------------------------------------------- + def test_theme_selection(main_window): # First item is "None" main_window.theme_combo_box.setCurrentIndex(0) @@ -139,6 +147,7 @@ def test_theme_selection(main_window): # Model interaction (lightweight) # ------------------------------------------------------------------- + def test_add_asset_updates_scene(main_window): scene = main_window.scene @@ -152,6 +161,7 @@ def test_add_asset_updates_scene(main_window): # Quit behavior # ------------------------------------------------------------------- + def test_quit_app_calls_app_quit(monkeypatch, main_window): called = {"quit": False} diff --git a/tests/test_model_scene.py b/tests/test_model_scene.py index 7d55b7a..0f822ab 100644 --- a/tests/test_model_scene.py +++ b/tests/test_model_scene.py @@ -1,7 +1,6 @@ import pytest -from PySide6.QtWidgets import QApplication, QGraphicsScene, QGraphicsLineItem -from PySide6.QtCore import Qt, QPointF, QMimeData, QEvent -from PySide6.QtGui import QDropEvent, QMouseEvent +from PySide6.QtWidgets import QGraphicsScene +from PySide6.QtCore import QPointF from maltoolbox.language import LanguageGraph from maltoolbox.model import Model @@ -9,20 +8,23 @@ from mal_gui.model_scene import ModelScene from mal_gui.object_explorer import AssetItem, AttackerItem + @pytest.fixture def model_scene(app): - lang_file_path = 'tests/testdata/org.mal-lang.coreLang-1.0.0.mar' + lang_file_path = "tests/testdata/org.mal-lang.coreLang-1.0.0.mar" lang_graph = LanguageGraph.load_from_file(str(lang_file_path)) main_window = MainWindow(app, str(lang_file_path)) asset_factory = main_window.asset_factory model = Model("TestModel", lang_graph) return ModelScene(asset_factory, lang_graph, model, main_window) + def test_scene_initialization(model_scene): assert isinstance(model_scene, QGraphicsScene) assert model_scene.undo_stack.count() == 0 assert isinstance(model_scene.clipboard, type(model_scene.clipboard)) + def test_add_asset(model_scene): pos = QPointF(50, 50) asset_item = model_scene.create_asset("Application", pos, name="Application1") @@ -30,6 +32,7 @@ def test_add_asset(model_scene): assert asset_item.pos() == pos assert asset_item in model_scene._asset_id_to_item.values() + def test_add_attacker(model_scene): pos = QPointF(0, 0) attacker_item = model_scene.create_attacker(pos, "Attacker1") @@ -37,6 +40,7 @@ def test_add_attacker(model_scene): assert attacker_item.pos() == pos assert attacker_item in model_scene.attacker_items + # def test_connection_creation(model_scene): # a1 = model_scene.create_asset("Application", QPointF(0, 0)) # a2 = model_scene.create_asset("Application", QPointF(100, 0)) @@ -51,6 +55,7 @@ def test_add_attacker(model_scene): # assert model_scene.start_item is None # assert model_scene.end_item is None + def test_undo_redo(model_scene): pos = QPointF(0, 0) asset_item = model_scene.create_asset("Application", pos) @@ -60,9 +65,10 @@ def test_undo_redo(model_scene): model_scene.undo_stack.undo() assert len(model_scene.items()) == initial_count + def test_serialize_deserialize(model_scene): asset_item = model_scene.create_asset("Application", QPointF(0, 0)) serialized = model_scene.serialize_graphics_items([asset_item], cut_intended=False) deserialized = model_scene.deserialize_graphics_items(serialized) assert isinstance(deserialized, list) - assert deserialized[0]['title'] == asset_item.title + assert deserialized[0]["title"] == asset_item.title