From 653cafdfa67623c69f0c85a986cbc228b07a1608 Mon Sep 17 00:00:00 2001 From: sandorstormen Date: Mon, 20 Oct 2025 15:45:08 +0200 Subject: [PATCH 1/3] Add Gephi GEXF exporter for attack graphs --- maltoolbox/translators/gephi.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 maltoolbox/translators/gephi.py diff --git a/maltoolbox/translators/gephi.py b/maltoolbox/translators/gephi.py new file mode 100644 index 00000000..9cddf5a3 --- /dev/null +++ b/maltoolbox/translators/gephi.py @@ -0,0 +1,18 @@ + +from __future__ import annotations + +import networkx as nx + +from maltoolbox.attackgraph import AttackGraph +from .networkx import attack_graph_to_nx + + +def attack_graph_to_gexf(attack_graph: AttackGraph, output_path: str) -> None: + """Export an attack graph to Gephi-compatible GEXF format + + Arguments: + attack_graph - the attack graph to export + output_path - file path where the .gexf file will be written + """ + G: nx.DiGraph = attack_graph_to_nx(attack_graph) + nx.write_gexf(G, output_path) From ebc8f34f826997eff84d9def00805d98bec87aed Mon Sep 17 00:00:00 2001 From: sandorstormen Date: Mon, 20 Oct 2025 16:54:12 +0200 Subject: [PATCH 2/3] add: gexf format conversion for attack graphs --- maltoolbox/translators/__init__.py | 4 +- maltoolbox/translators/gephi.py | 18 ----- maltoolbox/translators/gexf.py | 117 +++++++++++++++++++++++++++++ pyproject.toml | 3 +- tests/translators/test_gexf.py | 35 +++++++++ 5 files changed, 157 insertions(+), 20 deletions(-) delete mode 100644 maltoolbox/translators/gephi.py create mode 100644 maltoolbox/translators/gexf.py create mode 100644 tests/translators/test_gexf.py diff --git a/maltoolbox/translators/__init__.py b/maltoolbox/translators/__init__.py index 2bf270dc..086ebf18 100644 --- a/maltoolbox/translators/__init__.py +++ b/maltoolbox/translators/__init__.py @@ -1,8 +1,10 @@ from .networkx import attack_graph_to_nx, model_to_nx from .updater import load_model_from_older_version +from .gexf import attack_graph_to_gexf __all__ = [ 'attack_graph_to_nx', 'model_to_nx', - 'load_model_from_older_version' + 'load_model_from_older_version', + 'attack_graph_to_gexf' ] \ No newline at end of file diff --git a/maltoolbox/translators/gephi.py b/maltoolbox/translators/gephi.py deleted file mode 100644 index 9cddf5a3..00000000 --- a/maltoolbox/translators/gephi.py +++ /dev/null @@ -1,18 +0,0 @@ - -from __future__ import annotations - -import networkx as nx - -from maltoolbox.attackgraph import AttackGraph -from .networkx import attack_graph_to_nx - - -def attack_graph_to_gexf(attack_graph: AttackGraph, output_path: str) -> None: - """Export an attack graph to Gephi-compatible GEXF format - - Arguments: - attack_graph - the attack graph to export - output_path - file path where the .gexf file will be written - """ - G: nx.DiGraph = attack_graph_to_nx(attack_graph) - nx.write_gexf(G, output_path) diff --git a/maltoolbox/translators/gexf.py b/maltoolbox/translators/gexf.py new file mode 100644 index 00000000..ccb1701c --- /dev/null +++ b/maltoolbox/translators/gexf.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from maltoolbox.attackgraph import AttackGraph + +from gexfpy import stringify +from gexfpy import ( + Gexf, + Graph, + Nodes, + Edges, + Node, + Edge, + Color, + Attvalues, + Attvalue, + Thickness, + NodeShapeContent, + NodeShapeType, + EdgeShapeContent, + EdgeShapeType, + IdtypeType, + ModeType, + DefaultedgetypeType, + Attribute, + AttrtypeType, + Attributes, +) + + +def attack_graph_to_gexf( + attack_graph: AttackGraph, + node_type_shape_map: dict[str, NodeShapeContent] = { + "or": NodeShapeContent(value=NodeShapeType.DISC), + "and": NodeShapeContent(value=NodeShapeType.DIAMOND), + "exist": NodeShapeContent(value=NodeShapeType.TRIANGLE), + "notExist": NodeShapeContent(value=NodeShapeType.TRIANGLE), + "defense": NodeShapeContent(value=NodeShapeType.SQUARE), + }, + color_map: dict[str, Color] | dict[int, Color] = {}, + edge_thickness: Thickness = Thickness(value=3.0), + edge_shape: EdgeShapeContent = EdgeShapeContent(value=EdgeShapeType.SOLID), +) -> Gexf: + """Export an attack graph to GEXF format""" + + node_list: list[Node] = [] + edge_list: list[Edge] = [] + for node_id, node in attack_graph.nodes.items(): + node_list.append( + Node( + id=node_id, + label=node.full_name, + attvalues=Attvalues( + [ + Attvalue( + for_value=(key if isinstance(key, int) else str(key)), + value=str(value), + ) + for key, value in node.to_dict().items() + if key not in {"children", "parents"} + ] + ), + color=[color_map.get(node_id, Color(r=0, g=0, b=0, a=0.5))], + shape=[ + node_type_shape_map.get( + node.type, NodeShapeContent(value=NodeShapeType.DISC) + ) + ], + ) + ) + + for child in node.children: + edge_list.append( + Edge( + source=node_id, + target=child.id, + thickness=[edge_thickness], + shape=[edge_shape], + ) + ) + + graph = Graph( + attributes=Attributes( + attribute=[ + Attribute( + default=[attack_graph.model.name], + id=0, + title="model_name", + type=AttrtypeType.STRING, + ), + Attribute( + default=[attack_graph.model.maltoolbox_version], + id=1, + title="maltoolbox_version", + type=AttrtypeType.STRING, + ), + Attribute( + default=[attack_graph.lang_graph.metadata["version"]], + id=2, + title="lang_version", + type=AttrtypeType.STRING, + ), + Attribute( + default=[attack_graph.lang_graph.metadata["id"]], + id=3, + title="lang_id", + type=AttrtypeType.STRING, + ), + ] + ), + nodes=Nodes(node=node_list, count=len(node_list)), + edges=Edges(edge=edge_list, count=len(edge_list)), + defaultedgetype=DefaultedgetypeType.DIRECTED, + idtype=IdtypeType.INTEGER, + mode=ModeType.STATIC, + ) + + return Gexf(graph=graph) diff --git a/pyproject.toml b/pyproject.toml index 35e366f3..053410dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,8 @@ dependencies = [ "docopt", "PyYAML", "py2neo", - "networkx" + "networkx", + "gexfpy" ] license = {text = "Apache Software License"} keywords = ["mal"] diff --git a/tests/translators/test_gexf.py b/tests/translators/test_gexf.py new file mode 100644 index 00000000..0a81758e --- /dev/null +++ b/tests/translators/test_gexf.py @@ -0,0 +1,35 @@ +from maltoolbox.translators import attack_graph_to_gexf +from maltoolbox.attackgraph import AttackGraph +import pytest +from gexfpy import stringify +import xml.etree.ElementTree as ET + +def test_attack_graph_to_gexf(example_attackgraph: AttackGraph): + """Test conversion of attack graph to GEXF format""" + + def number_of_edges(attack_graph: AttackGraph) -> int: + edges = set() + for node in attack_graph.nodes.values(): + for child in node.children: + edges.add((node.id, child.id)) + for parent in node.parents: + edges.add((parent.id, node.id)) + return len(edges) + + + gexf = attack_graph_to_gexf(example_attackgraph) + + assert gexf.graph.nodes.count == len(example_attackgraph.nodes) + assert gexf.graph.edges.count == number_of_edges(example_attackgraph) + + for node in example_attackgraph.nodes.values(): + assert gexf.graph.nodes.node[node.id].id == node.id, f"Node id mismatch" + assert gexf.graph.nodes.node[node.id].label == node.full_name, f"Node label mismatch" + node_dict = node.to_dict() + for attvalue in gexf.graph.nodes.node[node.id].attvalues.attvalue: + assert str(node_dict[attvalue.for_value]) == attvalue.value, f"Attribute {attvalue.for_value} value mismatch" + + try: + ET.fromstring(stringify(gexf)) + except ET.ParseError as e: + pytest.fail(f"Output is not valid XML: {e}") \ No newline at end of file From af0e1e06899068400105bae27c13f5b5120ae212 Mon Sep 17 00:00:00 2001 From: sandorstormen Date: Mon, 20 Oct 2025 17:17:46 +0200 Subject: [PATCH 3/3] feat: add GEXF export functionality for models --- maltoolbox/translators/__init__.py | 5 +- maltoolbox/translators/gexf.py | 94 ++++++++++++++++++++++++++++++ tests/translators/test_gexf.py | 31 +++++++++- 3 files changed, 127 insertions(+), 3 deletions(-) diff --git a/maltoolbox/translators/__init__.py b/maltoolbox/translators/__init__.py index 086ebf18..a9e2311c 100644 --- a/maltoolbox/translators/__init__.py +++ b/maltoolbox/translators/__init__.py @@ -1,10 +1,11 @@ from .networkx import attack_graph_to_nx, model_to_nx from .updater import load_model_from_older_version -from .gexf import attack_graph_to_gexf +from .gexf import attack_graph_to_gexf, model_to_gexf __all__ = [ 'attack_graph_to_nx', 'model_to_nx', 'load_model_from_older_version', - 'attack_graph_to_gexf' + 'attack_graph_to_gexf', + 'model_to_gexf' ] \ No newline at end of file diff --git a/maltoolbox/translators/gexf.py b/maltoolbox/translators/gexf.py index ccb1701c..a0c23539 100644 --- a/maltoolbox/translators/gexf.py +++ b/maltoolbox/translators/gexf.py @@ -1,6 +1,7 @@ from __future__ import annotations from maltoolbox.attackgraph import AttackGraph +from maltoolbox.model import Model from gexfpy import stringify from gexfpy import ( @@ -115,3 +116,96 @@ def attack_graph_to_gexf( ) return Gexf(graph=graph) + + +def model_to_gexf( + model: Model, + color_map: dict[str, Color] | dict[int, Color] = {}, + edge_thickness: Thickness = Thickness(value=8.0), + edge_shape: EdgeShapeContent = EdgeShapeContent(value=EdgeShapeType.SOLID), +) -> Gexf: + """Export a model to GEXF format""" + + node_list: list[Node] = [] + edge_list: list[Edge] = [] + + # Create nodes for each asset + for asset_id, asset in model.assets.items(): + node_list.append( + Node( + id=asset_id, + label=asset.name, + attvalues=Attvalues( + [ + Attvalue( + for_value="type", + value=asset.type, + ), + Attvalue( + for_value="name", + value=asset.name, + ), + ] + + [ + Attvalue( + for_value=defense_name, + value=str(defense_value), + ) + for defense_name, defense_value in asset.defenses.items() + ] + ), + color=[color_map.get(asset_id, Color(r=0, g=0, b=0, a=0.5))], + shape=[NodeShapeContent(value=NodeShapeType.SQUARE)], + ) + ) + + # Create edges for associations + for fieldname, associated_assets in asset.associated_assets.items(): + for associated_asset in associated_assets: + edge_list.append( + Edge( + source=asset_id, + target=associated_asset.id, + thickness=[edge_thickness], + shape=[edge_shape], + label=fieldname, + ) + ) + + graph = Graph( + attributes=Attributes( + attribute=[ + Attribute( + default=[model.name], + id=0, + title="model_name", + type=AttrtypeType.STRING, + ), + Attribute( + default=[model.maltoolbox_version], + id=1, + title="maltoolbox_version", + type=AttrtypeType.STRING, + ), + Attribute( + default=[model.lang_graph.metadata["version"]], + id=2, + title="lang_version", + type=AttrtypeType.STRING, + ), + Attribute( + default=[model.lang_graph.metadata["id"]], + id=3, + title="lang_id", + type=AttrtypeType.STRING, + ), + ] + ), + nodes=Nodes(node=node_list, count=len(node_list)), + edges=Edges(edge=edge_list, count=len(edge_list)), + defaultedgetype=DefaultedgetypeType.UNDIRECTED, + idtype=IdtypeType.INTEGER, + mode=ModeType.STATIC, + ) + + return Gexf(graph=graph) diff --git a/tests/translators/test_gexf.py b/tests/translators/test_gexf.py index 0a81758e..2bd0fd8b 100644 --- a/tests/translators/test_gexf.py +++ b/tests/translators/test_gexf.py @@ -1,4 +1,5 @@ -from maltoolbox.translators import attack_graph_to_gexf +from maltoolbox.model import Model +from maltoolbox.translators import attack_graph_to_gexf, model_to_gexf from maltoolbox.attackgraph import AttackGraph import pytest from gexfpy import stringify @@ -29,6 +30,34 @@ def number_of_edges(attack_graph: AttackGraph) -> int: for attvalue in gexf.graph.nodes.node[node.id].attvalues.attvalue: assert str(node_dict[attvalue.for_value]) == attvalue.value, f"Attribute {attvalue.for_value} value mismatch" + try: + ET.fromstring(stringify(gexf)) + except ET.ParseError as e: + pytest.fail(f"Output is not valid XML: {e}") + +def test_model_to_gexf(example_model: Model): + """Test conversion of model to GEXF format""" + + def number_of_edges(model: Model) -> int: + edges = set() + for asset in model.assets.values(): + for _fieldname, associated_assets in asset.associated_assets.items(): + for associated_asset in associated_assets: + edges.add((asset.id, associated_asset.id)) + return len(edges) + + gexf = model_to_gexf(example_model) + + assert gexf.graph.nodes.count == len(example_model.assets) + assert gexf.graph.edges.count == number_of_edges(example_model) + + for asset in example_model.assets.values(): + assert gexf.graph.nodes.node[asset.id].id == asset.id, f"Node id mismatch" + assert gexf.graph.nodes.node[asset.id].label == asset.name, f"Node label mismatch" + asset_dict = asset._to_dict()[asset.id] + for attvalue in gexf.graph.nodes.node[asset.id].attvalues.attvalue: + assert str(asset_dict[attvalue.for_value]) == attvalue.value, f"Attribute {attvalue.for_value} value mismatch" + try: ET.fromstring(stringify(gexf)) except ET.ParseError as e: