Skip to content

Commit ed6b5b9

Browse files
mrkicklingmrkickling
committed
Partly fix serialization of assocs and entrypoints and allow pasting of them
1 parent 1e5f88f commit ed6b5b9

File tree

4 files changed

+126
-45
lines changed

4 files changed

+126
-45
lines changed

‎mal_gui/connection_dialog.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,8 @@ def ok_button_clicked(self):
146146
selected_association_text = selected_item.text()
147147
# QMessageBox.information(self, "Selected Item", f"You selected: {selected_association_text}")
148148

149-
(assoc, left_asset, right_asset) = self._str_to_assoc[selected_association_text]
149+
(assoc, left_asset, right_asset) = \
150+
self._str_to_assoc[selected_association_text]
150151
# TODO: Create association based on its full name instead in order
151152
# to avoid conflicts when multiple associations with the same name
152153
# exist.

‎mal_gui/connection_item.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,17 +59,24 @@ def __init__(
5959
and self.end_item.asset_type != 'Attacker':
6060

6161
self.association_details = selected_assoc_text.split("-->")
62+
# Get left field name
6263
assoc_left_field = self.association_details[0]
63-
assoc_middle_name = self.association_details[1]
64+
self.left_fieldname = assoc_left_field.split(".")[1]
65+
# Get assoc name
66+
self.assoc_name = self.association_details[1]
6467
assoc_right_field = self.association_details[2]
68+
# Get right field name
69+
self.right_fieldname = assoc_right_field.split(".")[1]
6570

6671
# Create labels with background color
6772
self.label_assoc_left_field = \
68-
self.create_label(assoc_left_field.split(".")[1])
73+
self.create_label(self.left_fieldname)
74+
6975
self.label_assoc_middle_name = \
70-
self.create_label(assoc_middle_name)
76+
self.create_label(self.assoc_name)
77+
7178
self.label_assoc_right_field = \
72-
self.create_label(assoc_right_field.split(".")[1])
79+
self.create_label(self.right_fieldname)
7380

7481
else:
7582
# Need to check who is attacker and get
@@ -206,6 +213,7 @@ def __init__(
206213

207214
self.attacker_item.add_connection(self)
208215
self.asset_item.add_connection(self)
216+
self.attack_step_name = attack_step_name
209217
self.label_entrypoint = self.create_label(attack_step_name)
210218

211219
def create_label(self, text):

‎mal_gui/model_scene.py

Lines changed: 77 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -589,8 +589,77 @@ def compress_container(self,currently_selected_container_box):
589589

590590
self.removeItem(currently_selected_container_box)
591591

592+
def serialize_associations(
593+
self,
594+
connections: list[IConnectionItem],
595+
selected_sequence_ids: list[int]
596+
):
597+
"""Serialize selected connections"""
598+
599+
serialized_associations = []
600+
for conn in connections:
601+
# Copy associations where both item are selected
602+
if not isinstance(conn, AssociationConnectionItem):
603+
continue
604+
605+
both_items_selected = (
606+
conn.start_item.asset_sequence_id\
607+
in selected_sequence_ids and
608+
conn.end_item.asset_sequence_id \
609+
in selected_sequence_ids
610+
)
611+
if not both_items_selected:
612+
continue
613+
# If association and selected, serialize it
614+
serialized_associations.append(
615+
(
616+
conn.start_item.asset_sequence_id,
617+
conn.end_item.asset_sequence_id,
618+
conn.assoc_name,
619+
conn.left_fieldname,
620+
conn.right_fieldname,
621+
"-->".join(conn.association_details)
622+
)
623+
)
624+
625+
return serialized_associations
626+
627+
def serialize_entrypoints(
628+
self,
629+
entrypoints: list[IConnectionItem],
630+
selected_sequence_ids: list[int]
631+
):
632+
"""Serialize selected attacker entrypoints"""
633+
634+
serialized_entrypoints = []
635+
for conn in entrypoints:
636+
# Copy entrypoints where both item are selected
637+
if not isinstance(conn, EntrypointConnectionItem):
638+
continue
639+
640+
# If entry points
641+
both_items_selected = (
642+
conn.asset_item.asset_sequence_id\
643+
in selected_sequence_ids and
644+
conn.attacker_item.asset_sequence_id \
645+
in selected_sequence_ids
646+
)
647+
if not both_items_selected:
648+
continue
649+
650+
serialized_entrypoints.append(
651+
(
652+
conn.attacker_item.asset_sequence_id,
653+
conn.asset_item.asset_sequence_id,
654+
conn.attack_step_name
655+
)
656+
)
657+
658+
return serialized_entrypoints
659+
592660
def serialize_graphics_items(self, items: list[AssetBase], cut_intended):
593-
objdetails = []
661+
"""Serialize all selected items"""
662+
serialized_items = []
594663

595664
# Set of selected item IDs
596665
selected_sequence_ids = {item.asset_sequence_id for item in items}
@@ -606,27 +675,23 @@ def serialize_graphics_items(self, items: list[AssetBase], cut_intended):
606675
'asset_name': asset_name,
607676
'asset_sequence_id': item.asset_sequence_id,
608677
'position': (item.pos().x(), item.pos().y()),
609-
'connections': [
610-
(conn.start_item.asset_sequence_id,
611-
conn.end_item.asset_sequence_id,
612-
'-->'.join(conn.association_details))
613-
for conn in item.connections
614-
if conn.start_item.asset_sequence_id in selected_sequence_ids
615-
and conn.end_item.asset_sequence_id in selected_sequence_ids
616-
],
617678
'asset_properties': []
618679
}
619680

681+
item_details['associations'] = self.serialize_associations(
682+
item.connections, selected_sequence_ids)
683+
item_details['entrypoints'] = self.serialize_entrypoints(
684+
item.connections, selected_sequence_ids)
685+
620686
if item.asset_type != "Attacker":
621687
item_details['asset_properties'] = [
622688
(str(key),str(value))
623689
for key, value in item.asset._properties.items()
624690
if key not in prop_keys_to_ignore
625691
]
692+
serialized_items.append(item_details)
626693

627-
objdetails.append(item_details)
628-
629-
serialized_data = pickle.dumps(objdetails)
694+
serialized_data = pickle.dumps(serialized_items)
630695
base64_serialized_data = \
631696
base64.b64encode(serialized_data).decode('utf-8')
632697
return base64_serialized_data

‎mal_gui/undo_redo_commands/paste_command.py

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
from maltoolbox.model import AttackerAttachment
99

10-
from ..connection_item import AssociationConnectionItem
10+
from ..connection_item import AssociationConnectionItem, EntrypointConnectionItem
1111

1212
if TYPE_CHECKING:
1313
from ..model_scene import ModelScene
@@ -26,14 +26,13 @@ def __init__(
2626
self.position = position
2727
self.pasted_items = []
2828
self.pasted_connections: list[AssociationConnectionItem] = []
29+
self.pasted_entrypoints: list[EntrypointConnectionItem] = []
2930
self.clipboard = clipboard
3031

3132
def redo(self):
3233
"""Perform paste command"""
3334
print("\nPaste Redo is Called")
3435
serialized_data = self.clipboard.text()
35-
print("\nSerializedData = " + str(len(serialized_data)))
36-
print("\nSerializedData = " + str(serialized_data))
3736
if serialized_data:
3837
deserialized_data = \
3938
self.scene.deserialize_graphics_items(serialized_data)
@@ -86,41 +85,49 @@ def redo(self):
8685

8786
# Second pass: re-establish connections with new assetSequenceIds
8887
for data in deserialized_data:
89-
for conn in data['connections']:
90-
print(f'CONN: {conn}')
88+
for entrypoint in data['entrypoints']:
89+
print(f'ENTRYPOINT: {entrypoint}')
90+
old_start_id, old_end_id, label = entrypoint
91+
new_attacker_item = new_item_map[old_start_id]
92+
new_asset_item = new_item_map[old_end_id]
93+
new_connection = self.scene\
94+
.add_entrypoint_connection(
95+
label, new_attacker_item, new_asset_item
96+
)
97+
self.pasted_entrypoints.append(new_connection)
98+
new_attacker_item.attackerAttachment.add_entry_point(
99+
new_asset_item.asset,
100+
label
101+
)
102+
103+
added_assoc_labels = set()
104+
for assoc in data['associations']:
105+
print(data)
106+
print(f'CONN: {assoc}')
91107
old_start_id = data['asset_sequence_id']
92-
old_end_id, old_label = conn[1], conn[2]
108+
(_, old_end_id, assoc_type,
109+
left_field_name, right_field_name, label) = assoc
110+
93111
new_start_item = new_item_map[old_start_id]
94112
new_end_item = new_item_map[old_end_id]
95113

96-
#Avoid Self reference
114+
if label in added_assoc_labels:
115+
continue
116+
117+
added_assoc_labels.add(label)
118+
119+
# Avoid Self reference
97120
if new_start_item != new_end_item:
98-
new_connection = AssociationConnectionItem(old_label,
99-
new_start_item,
100-
new_end_item,
101-
self.scene
102-
)
103-
self.scene.addItem(new_connection)
121+
new_connection = self.scene\
122+
.add_association_connection(
123+
label, new_start_item, new_end_item
124+
)
125+
104126
self.pasted_connections.append(new_connection)
105-
assoc_type = old_label.split('-->')[1]
106127
print(f'ASSOC TYPE {assoc_type}')
107128
association = getattr(self.scene.lcs.ns, assoc_type)()
108129
left_asset = new_start_item.asset
109130
right_asset = new_end_item.asset
110-
print(
111-
f'LEFT ASSET NAME:{left_asset.name}'
112-
f'ID:{left_asset.id}'
113-
f'TYPE:{left_asset.type}'
114-
)
115-
print(
116-
f'RIGHT ASSET NAME:{right_asset.name}'
117-
f'ID:{right_asset.id}'
118-
f'TYPE:{right_asset.type}'
119-
)
120-
left_field_name, right_field_name = \
121-
self.scene.model.get_association_field_names(
122-
association
123-
)
124131
setattr(association, left_field_name, [left_asset])
125132
setattr(association, right_field_name, [right_asset])
126133
self.scene.model.add_association(association)

0 commit comments

Comments
 (0)