diff --git a/maltoolbox/translators/__init__.py b/maltoolbox/translators/__init__.py index 2bf270dc..a9e2311c 100644 --- a/maltoolbox/translators/__init__.py +++ b/maltoolbox/translators/__init__.py @@ -1,8 +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, model_to_gexf __all__ = [ 'attack_graph_to_nx', 'model_to_nx', - 'load_model_from_older_version' + 'load_model_from_older_version', + '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 new file mode 100644 index 00000000..a0c23539 --- /dev/null +++ b/maltoolbox/translators/gexf.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from maltoolbox.attackgraph import AttackGraph +from maltoolbox.model import Model + +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) + + +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/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..2bd0fd8b --- /dev/null +++ b/tests/translators/test_gexf.py @@ -0,0 +1,64 @@ +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 +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}") + +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: + pytest.fail(f"Output is not valid XML: {e}") \ No newline at end of file