Merge pull request #4 from mal-lang/upgrade-maltoolbox-0.3.0

Upgrade maltoolbox 0.3.0

Author: mrkickling

mal_gui/connection_dialog.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING
+from PySide6.QtCore import Qt
 from PySide6.QtWidgets import (
     QDialog,
     QVBoxLayout,
@@ -13,7 +14,9 @@
 )
 
 if TYPE_CHECKING:
-    from .object_explorer.asset_base import AssetBase
+    from .object_explorer import AssetItem, AttackerItem
+    from maltoolbox.language import LanguageGraph, LanguageGraphAsset
+    from maltoolbox.model import Model, ModelAsset
 
 class ConnectionDialog(QDialog):
     def filter_items(self, text):
@@ -26,16 +29,15 @@ def ok_button_clicked(self):
 class AssociationConnectionDialog(ConnectionDialog):
     def __init__(
             self,
-            start_item: AssetBase,
-            end_item: AssetBase,
-            lang_graph,
-            lcs,model,
+            start_item: AssetItem,
+            end_item: AssetItem,
+            lang_graph: LanguageGraph,
+            model: Model,
             parent=None
         ):
         super().__init__(parent)
 
-        self.lang_graph = lang_graph
-        self.lcs = lcs
+        self.lang_graph: LanguageGraph = lang_graph
         self.model = model
 
         self.setWindowTitle("Select Association Type")
@@ -46,75 +48,38 @@ def __init__(
 
         self.association_list_widget = QListWidget()
 
-        start_asset = start_item.asset
-        end_asset = end_item.asset
-        self.start_asset_type = start_asset.type
-        self.end_asset_type = end_asset.type
-        self.start_asset_name = start_asset.name
-        self.end_asset_name = end_asset.name
+        self.start_asset: ModelAsset = start_item.asset
+        self.end_asset: ModelAsset = end_item.asset
+        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...")
         self.filter_edit.textChanged.connect(self.filter_items)
         self.layout.addWidget(self.filter_edit)
-        lang_graph_start_asset = next(
-                (asset for asset in self.lang_graph.assets
-                 if asset.name == start_asset.type), None
-            )
-        if lang_graph_start_asset is None:
-            raise LookupError(f'Failed to find asset "{start_asset.type}" '
-                'in language graph.')
-        lang_graph_end_asset = next(
-                (asset for asset in self.lang_graph.assets
-                 if asset.name == end_asset.type), None
-            )
-        if lang_graph_end_asset is None:
-            raise LookupError(f'Failed to find asset "{end_asset.type}" '
-                'in language graph.')
-
-        self._str_to_assoc = {}
-        for assoc in lang_graph_start_asset.associations:
-            asset_pairs = []
-            opposite_asset = assoc.get_opposite_asset(lang_graph_start_asset)
-            # Check if the other side of the association matches the other end
-            # and if the exact association does not already exist in the
-            # model.
-            if lang_graph_end_asset.is_subasset_of(opposite_asset):
-                print("IDENTIFIED MATCH  ++++++++++++")
-                if lang_graph_start_asset.is_subasset_of(assoc.left_field.asset):
-                    asset_pairs.append((start_asset, end_asset))
-                else:
-                    asset_pairs.append((end_asset, start_asset))
-            if lang_graph_start_asset.is_subasset_of(opposite_asset):
-                # The association could be applied either way, add the
-                # reverse association as well.
-                other_asset = assoc.get_opposite_asset(opposite_asset)
-                # Check if the other side of the association matches the other end
-                # and if the exact association does not already exist in the
-                # model.
-                if lang_graph_end_asset.is_subasset_of(other_asset):
-                    print("REVERSE ASSOC  ++++++++++++")
-                    # We need to create the reverse association as well
-                    asset_pairs.append((end_asset, start_asset))
-            for (left_asset, right_asset) in asset_pairs:
-                if not self.model.association_exists_between_assets(
-                        assoc.name,
-                        left_asset,
-                        right_asset):
-                    formatted_assoc_str = left_asset.name + "." + \
-                        assoc.left_field.fieldname + "-->" + \
-                        assoc.name + "-->" + \
-                        right_asset.name + "." + \
-                        assoc.right_field.fieldname
-                    self._str_to_assoc[formatted_assoc_str] = (
-                        assoc,
-                        left_asset,
-                        right_asset
-                    )
-                    self.association_list_widget.addItem(QListWidgetItem(formatted_assoc_str))
+
+        possible_assocs = self.start_asset.lg_asset.associations
+        for fieldname, association in possible_assocs.items():
+            field = association.get_field(fieldname)
+
+            # 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.setData(
+                    Qt.UserRole,
+                    {
+                        'from': self.start_asset,
+                        'to': self.end_asset,
+                        'fieldname': fieldname
+                    }
+                )
+                self.association_list_widget.addItem(assoc_list_item)
+
         self.layout.addWidget(self.association_list_widget)
 
         button_layout = QHBoxLayout()
@@ -142,65 +107,53 @@ def filter_items(self, text):
 
     def ok_button_clicked(self):
         selected_item = self.association_list_widget.currentItem()
+
         if selected_item:
-            selected_association_text = selected_item.text()
-            # QMessageBox.information(self, "Selected Item", f"You selected: {selected_association_text}")
-
-            (assoc, left_asset, right_asset) = \
-                self._str_to_assoc[selected_association_text]
-            # TODO: Create association based on its full name instead in order
-            # to avoid conflicts when multiple associations with the same name
-            # exist.
-            association = getattr(self.lcs.ns, assoc.name)()
-            print(
-                f'N:{assoc.name} LF:{assoc.left_field.fieldname} '
-                f'LA:{left_asset.name} RF:{assoc.right_field.fieldname} '
-                f'RA:{right_asset.name}'
-            )
-            setattr(association, assoc.left_field.fieldname, [left_asset])
-            setattr(association, assoc.right_field.fieldname, [right_asset])
-            selected_item.association = association
-            # self.model.add_association(association)
+            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')
+
+
         self.accept()
 
 class EntrypointConnectionDialog(ConnectionDialog):
     def __init__(
             self,
-            attacker_item,
-            asset_item,
-            lang_graph,
-            lcs,
+            attacker_item: AttackerItem,
+            asset_item: AssetItem,
+            lang_graph: LanguageGraph,
             model,
             parent=None
         ):
         super().__init__(parent)
 
         self.lang_graph = lang_graph
-        self.lcs = lcs
         self.model = model
 
         self.setWindowTitle("Select Entry Point")
         self.setMinimumWidth(300)
 
-        print(f'Attacker ITEM TYPE {attacker_item.asset_type}')
-        print(f'Asset ITEM TYPE {asset_item.asset_type}')
-
         self.attack_step_list_widget = QListWidget()
-        attacker = attacker_item.attackerAttachment
+        attacker = attacker_item.attacker
 
         if asset_item.asset is not None:
-            asset_type = \
-                self.lang_graph.get_asset_by_name(asset_item.asset.type)
+            asset_type = self.lang_graph.assets[asset_item.asset.type]
 
             # Find asset attack steps already part of attacker entry points
             entry_point_tuple = attacker.get_entry_point_tuple(
-                asset_item.asset)
+                asset_item.asset
+            )
+
             if entry_point_tuple is not None:
                 entry_point_attack_steps = entry_point_tuple[1]
             else:
                 entry_point_attack_steps = []
 
-            for attack_step in asset_type.attack_steps:
+            for attack_step in asset_type.attack_steps.values():
                 if attack_step.type not in ['or', 'and']:
                     continue
 

mal_gui/connection_item.py

@@ -10,13 +10,14 @@
 )
 
 if TYPE_CHECKING:
+    from maltoolbox.language import LanguageGraphAssociation
     from .model_scene import ModelScene
-    from .object_explorer.asset_base import AssetBase
+    from .object_explorer import AssetItem, AttackerItem
 
 class IConnectionItem(QGraphicsLineItem):
     """'interface' for Connection Item"""
-    start_item: AssetBase
-    end_item: AssetBase
+    start_item: AssetItem
+    end_item: AssetItem
     association_details: list[str]
 
     def create_label(self, text):
@@ -35,9 +36,9 @@ def restore_labels(self):
 class AssociationConnectionItem(IConnectionItem):
     def __init__(
         self,
-        selected_assoc_text: str,
-        start_item: AssetBase,
-        end_item: AssetBase,
+        fieldname: str,
+        start_item: AssetItem,
+        end_item: AssetItem,
         scene: ModelScene,
         parent = None
     ):
@@ -55,18 +56,22 @@ def __init__(
         self.start_item.add_connection(self)
         self.end_item.add_connection(self)
 
-        if self.start_item.asset_type != 'Attacker'\
-            and self.end_item.asset_type != 'Attacker':
+        if (
+            self.start_item.asset_type != 'Attacker'
+            and self.end_item.asset_type != 'Attacker'
+        ):
+            # Fetch the association and the fieldnames
+            lg_assoc: LanguageGraphAssociation = (
+                start_item.asset.lg_asset.associations[fieldname]
+            )
+            opposite_fieldname = lg_assoc.get_opposite_fieldname(fieldname)
 
-            self.association_details = selected_assoc_text.split("-->")
             # Get left field name
-            assoc_left_field = self.association_details[0]
-            self.left_fieldname = assoc_left_field.split(".")[1]
+            self.left_fieldname = opposite_fieldname
             # Get assoc name
-            self.assoc_name = self.association_details[1]
-            assoc_right_field = self.association_details[2]
+            self.assoc_name = lg_assoc.name
             # Get right field name
-            self.right_fieldname = assoc_right_field.split(".")[1]
+            self.right_fieldname = fieldname
 
             # Create labels with background color
             self.label_assoc_left_field = \
@@ -196,8 +201,8 @@ class EntrypointConnectionItem(IConnectionItem):
     def __init__(
         self,
         attack_step_name,
-        attacker_item: AssetBase,
-        asset_item: AssetBase,
+        attacker_item: AttackerItem,
+        asset_item: AssetItem,
         scene: ModelScene,
         parent = None
     ):

mal_gui/docked_windows/draggable_tree_view.py

@@ -12,7 +12,7 @@
 )
 
 if TYPE_CHECKING:
-    from ..object_explorer.asset_base import AssetBase
+    from ..object_explorer.asset_item import AssetItem
 
 class DraggableTreeView(QTreeWidget):
     def __init__(
@@ -134,7 +134,7 @@ def hide_unhide_asset_item(self, eye_button, item):
             self.eye_visibility = Visibility.HIDE
 
             #First Hide the connections associtaed with the asset item
-            asset_item: AssetBase = item.assetItemReference
+            asset_item: AssetItem = item.assetItemReference
 
             if hasattr(asset_item, 'connections'):
                 connections = asset_item.connections

mal_gui/docked_windows/item_details_window.py

@@ -1,17 +1,19 @@
 from PySide6.QtWidgets import QTreeWidget,QTreeWidgetItem
 
+from ..object_explorer import ItemBase
+
 class ItemDetailsWindow(QTreeWidget):
     def __init__(self, parent=None):
         super(ItemDetailsWindow, self).__init__(parent)
         self.setHeaderLabel(None)
         self.setColumnCount(2)
         self.setHeaderLabels(["Attribute","Value"])
 
-    def update_item_details_window(self, asset_item):
+    def update_item_details_window(self, item_object: ItemBase):
         self.clear()
-        if asset_item is not None:
+        if item_object is not None:
             # item has a method that returns a dict
-            asset_details = asset_item.get_item_attribute_vakues()
+            asset_details = item_object.get_item_attribute_values()
             for (key, value) in asset_details.items():
                 print(f"Attribute:{key} Value:{str(value)}")
                 item = QTreeWidgetItem([key, str(value)])

mal_gui/docked_windows/style_configuration.py

@@ -12,7 +12,7 @@
 from PySide6.QtCore import Signal
 from PySide6.QtGui import QColor
 
-from ..object_explorer.asset_base import AssetBase
+from ..object_explorer.asset_item import AssetItem
 
 class Visibility(Enum):
     HIDE = 1
@@ -199,7 +199,7 @@ def accept(self):
         super().accept()
         # self.update_color_callback(self.get_color_1(), self.get_color_2())
         for item in self.scene.items():
-            if isinstance(item, (AssetBase)):
+            if isinstance(item, (AssetItem)):
                 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()