feat: Support for GEXF

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'
 ]

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)

pyproject.toml

@@ -19,7 +19,8 @@ dependencies = [
   "docopt",
   "PyYAML",
   "py2neo",
-  "networkx"
+  "networkx",
+  "gexfpy"
 ]
 license = {text = "Apache Software License"}
 keywords = ["mal"]

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}")