mal-toolbox 1.0.2__tar.gz → 1.0.3__tar.gz

{mal_toolbox-1.0.2/mal_toolbox.egg-info → mal_toolbox-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mal-toolbox
-Version: 1.0.2
+Version: 1.0.3
 Summary: A collection of tools used to create MAL models and attack graphs.
 Author-email: Andrei Buhaiu <buhaiu@kth.se>, Joakim Loxdal <loxdal@kth.se>, Nikolaos Kakouros <nkak@kth.se>, Jakob Nyberg <jaknyb@kth.se>, Giuseppe Nebbione <nebbione@kth.se>
 License: Apache Software License
@@ -18,6 +18,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: AUTHORS
+Requires-Dist: graphviz
 Requires-Dist: antlr4-tools
 Requires-Dist: antlr4-python3-runtime
 Requires-Dist: docopt

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3/mal_toolbox.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mal-toolbox
-Version: 1.0.2
+Version: 1.0.3
 Summary: A collection of tools used to create MAL models and attack graphs.
 Author-email: Andrei Buhaiu <buhaiu@kth.se>, Joakim Loxdal <loxdal@kth.se>, Nikolaos Kakouros <nkak@kth.se>, Jakob Nyberg <jaknyb@kth.se>, Giuseppe Nebbione <nebbione@kth.se>
 License: Apache Software License
@@ -18,6 +18,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: AUTHORS
+Requires-Dist: graphviz
 Requires-Dist: antlr4-tools
 Requires-Dist: antlr4-python3-runtime
 Requires-Dist: docopt

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/mal_toolbox.egg-info/SOURCES.txt

@@ -18,13 +18,16 @@ maltoolbox/attackgraph/__init__.py
 maltoolbox/attackgraph/attackgraph.py
 maltoolbox/attackgraph/node.py
 maltoolbox/attackgraph/analyzers/__init__.py
-maltoolbox/ingestors/__init__.py
 maltoolbox/language/__init__.py
 maltoolbox/language/languagegraph.py
 maltoolbox/language/compiler/__init__.py
 maltoolbox/language/compiler/mal_lexer.py
 maltoolbox/language/compiler/mal_parser.py
+maltoolbox/patternfinder/__init__.py
+maltoolbox/patternfinder/attackgraph_patterns.py
 maltoolbox/translators/__init__.py
 maltoolbox/translators/securicad.py
 maltoolbox/translators/updater.py
+maltoolbox/visualization/__init__.py
+maltoolbox/visualization/graphviz_utils.py
 tests/test_model.py

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/mal_toolbox.egg-info/requires.txt

@@ -1,3 +1,4 @@
+graphviz
 antlr4-tools
 antlr4-python3-runtime
 docopt

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/maltoolbox/__init__.py

@@ -1,5 +1,5 @@
 # -*- encoding: utf-8 -*-
-# MAL Toolbox v1.0.2
+# MAL Toolbox v1.0.3
 # Copyright 2025, Andrei Buhaiu.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,7 +21,7 @@ MAL-Toolbox Framework
 """
 
 __title__ = "maltoolbox"
-__version__ = "1.0.2"
+__version__ = "1.0.3"
 __authors__ = [
     "Andrei Buhaiu",
     "Giuseppe Nebbione",

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/maltoolbox/__main__.py

@@ -2,15 +2,20 @@
 Command-line interface for MAL toolbox operations
 
 Usage:
-    maltoolbox attack-graph generate [options] <model_file> <lang_file>
     maltoolbox compile <lang_file> <output_file>
+    maltoolbox generate-attack-graph [--graphviz] <model_file> <lang_file>
     maltoolbox upgrade-model <model_file> <lang_file> <output_file>
+    maltoolbox visualize-model <model_file> <lang_file>
 
 Arguments:
     <model_file>    Path to JSON instance model file.
     <lang_file>     Path to .mar or .mal file containing MAL spec.
     <output_file>   Path to write the result of the compilation (yml/json).
 
+Options:
+  -h --help         Show this screen.
+  -g --graphviz     Visualize with graphviz
+
 Notes:
     - <lang_file> can be either a .mar file (generated by the older MAL
       compiler) or a .mal file containing the DSL written in MAL.
@@ -21,17 +26,19 @@ import json
 import docopt
 
 from . import log_configs
-from .attackgraph import create_attack_graph
+from .attackgraph import create_attack_graph, AttackGraph
 from .language.compiler import MalCompiler
 from .language.languagegraph import LanguageGraph
 from .translators.updater import load_model_from_older_version
+from .visualization.graphviz_utils import render_model, render_attack_graph
+from .model import Model
 
 logger = logging.getLogger(__name__)
 
 def generate_attack_graph(
         model_file: str,
-        lang_file: str,
-    ) -> None:
+        lang_file: str
+    ) -> AttackGraph:
     """Create an attack graph
 
     Args:
@@ -43,6 +50,7 @@ def generate_attack_graph(
         attack_graph.save_to_file(
             log_configs['attackgraph_file']
         )
+    return attack_graph
 
 def compile(lang_file: str, output_file: str) -> None:
     """Compile language and dump into output file"""
@@ -64,10 +72,13 @@ def upgrade_model(model_file: str, lang_file: str, output_file: str):
 def main():
     args = docopt.docopt(__doc__)
 
-    if args['attack-graph'] and args['generate']:
-        generate_attack_graph(
+    if args['generate-attack-graph']:
+        attack_graph = generate_attack_graph(
             args['<model_file>'], args['<lang_file>']
         )
+        if args['--graphviz']:
+            render_attack_graph(attack_graph)
+
     elif args['compile']:
         compile(
             args['<lang_file>'], args['<output_file>']
@@ -76,6 +87,10 @@ def main():
         upgrade_model(
             args['<model_file>'], args['<lang_file>'], args['<output_file>']
         )
+    elif args['visualize-model']:
+        lang_graph = LanguageGraph.load_from_file(args['<lang_file>'])
+        model = Model.load_from_file(args['<model_file>'], lang_graph)
+        render_model(model)
 
 if __name__ == "__main__":
     main()

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/maltoolbox/attackgraph/attackgraph.py

@@ -89,7 +89,7 @@ def create_attack_graph(
 
 class AttackGraph():
     """Graph representation of attack steps"""
-    def __init__(self, lang_graph, model: Optional[Model] = None):
+    def __init__(self, lang_graph: LanguageGraph, model: Optional[Model] = None):
         self.nodes: dict[int, AttackGraphNode] = {}
         # Dictionaries used in optimization to get nodes by id or full name
         # faster
@@ -535,15 +535,15 @@ class AttackGraph():
             if not ag_node.model_asset:
                 raise AttackGraphException('Attack graph node is missing '
                     'asset link')
-            lang_graph_asset = self.lang_graph.assets[
-                ag_node.model_asset.type]
 
-            lang_graph_attack_step = lang_graph_asset.attack_steps[
-                ag_node.name]
+            lang_graph_asset = self.lang_graph.assets[ag_node.model_asset.type]
+            lang_graph_attack_step: Optional[LanguageGraphAttackStep] = (
+                lang_graph_asset.attack_steps[ag_node.name]
+            )
 
             while lang_graph_attack_step:
-                for child in lang_graph_attack_step.children.values():
-                    for target_attack_step, expr_chain in child:
+                for target_attack_step, expr_chains in lang_graph_attack_step.children.items():
+                    for expr_chain in expr_chains:
                         target_assets = self._follow_expr_chain(
                             self.model,
                             set([ag_node.model_asset]),

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/maltoolbox/language/languagegraph.py

@@ -485,8 +485,13 @@ class LanguageGraphAttackStep:
     asset: LanguageGraphAsset
     ttc: Optional[dict] = field(default_factory = dict)
     overrides: bool = False
-    children: dict = field(default_factory = dict)
-    parents: dict = field(default_factory = dict)
+
+    own_children: dict[LanguageGraphAttackStep, list[ExpressionsChain | None]] = (
+        field(default_factory = dict)
+    )
+    own_parents: dict[LanguageGraphAttackStep, list[ExpressionsChain | None]] = (
+        field(default_factory = dict)
+    )
     info: dict = field(default_factory = dict)
     inherits: Optional[LanguageGraphAttackStep] = None
     own_requires: list[ExpressionsChain] = field(default_factory=list)
@@ -497,6 +502,39 @@ class LanguageGraphAttackStep:
     def __hash__(self):
         return hash(self.full_name)
 
+    @property
+    def children(self) -> dict[
+        LanguageGraphAttackStep, list[ExpressionsChain | None]
+    ]:
+        """
+        Get all (both own and inherited) children of a LanguageGraphAttackStep
+        """
+
+        all_children = dict(self.own_children)
+
+        if self.overrides:
+            # Override overrides the children
+            return all_children
+
+        if not self.inherits:
+            return all_children
+
+        for child_step, chains in self.inherits.children.items():
+            if child_step in all_children:
+                all_children[child_step] += [
+                    chain for chain in chains
+                    if chain not in all_children[child_step]
+                ]
+            else:
+                all_children[child_step] = chains
+
+        return all_children
+
+    @property
+    def parents(self) -> None:
+        raise NotImplementedError(
+            "Can not fetch parents of a LanguageGraphAttackStep"
+        )
 
     @property
     def full_name(self) -> str:
@@ -514,8 +552,8 @@ class LanguageGraphAttackStep:
             'type': self.type,
             'asset': self.asset.name,
             'ttc': self.ttc,
-            'children': {},
-            'parents': {},
+            'own_children': {},
+            'own_parents': {},
             'info': self.info,
             'overrides': self.overrides,
             'inherits': self.inherits.full_name if self.inherits else None,
@@ -524,25 +562,22 @@ class LanguageGraphAttackStep:
             self.detectors.items()},
         }
 
-        for child in self.children:
-            node_dict['children'][child] = []
-            for (_, expr_chain) in self.children[child]:
-                if expr_chain:
-                    node_dict['children'][child].append(
-                        expr_chain.to_dict())
+        for child, expr_chains in self.own_children.items():
+            node_dict['own_children'][child.full_name] = []
+            for chain in expr_chains:
+                if chain:
+                    node_dict['own_children'][child.full_name].append(chain.to_dict())
                 else:
-                    node_dict['children'][child].append(None)
-
-        for parent in self.parents:
-            node_dict['parents'][parent] = []
-            for (_, expr_chain) in self.parents[parent]:
-                if expr_chain:
-                    node_dict['parents'][parent].append(
-                        expr_chain.to_dict())
+                    node_dict['own_children'][child.full_name].append(None)
+        for parent, expr_chains in self.own_children.items():
+            node_dict['own_parents'][parent.full_name] = []
+            for chain in expr_chains:
+                if chain:
+                    node_dict['own_parents'][parent.full_name].append(chain.to_dict())
                 else:
-                    node_dict['parents'][parent].append(None)
+                    node_dict['own_parents'][parent.full_name].append(None)
 
-        if hasattr(self, 'own_requires'):
+        if self.own_requires:
             node_dict['requires'] = []
             for requirement in self.own_requires:
                 node_dict['requires'].append(requirement.to_dict())
@@ -965,8 +1000,8 @@ class LanguageGraph():
                     asset = asset,
                     ttc = get_ttc_distribution(attack_step_dict),
                     overrides = attack_step_dict['overrides'],
-                    children = {},
-                    parents = {},
+                    own_children = {},
+                    own_parents = {},
                     info = attack_step_dict['info'],
                     tags = list(attack_step_dict['tags'])
                 )
@@ -995,12 +1030,11 @@ class LanguageGraph():
             asset = lang_graph.assets[asset_dict['name']]
             for attack_step_dict in asset_dict['attack_steps'].values():
                 attack_step = asset.attack_steps[attack_step_dict['name']]
-                for child_target in attack_step_dict['children'].items():
+                for child_target in attack_step_dict['own_children'].items():
                     target_full_attack_step_name = child_target[0]
                     expr_chains = child_target[1]
                     target_asset_name, target_attack_step_name = \
-                        disaggregate_attack_step_full_name(
-                            target_full_attack_step_name)
+                        disaggregate_attack_step_full_name(target_full_attack_step_name)
                     target_asset = lang_graph.assets[target_asset_name]
                     target_attack_step = target_asset.attack_steps[
                         target_attack_step_name]
@@ -1009,33 +1043,30 @@ class LanguageGraph():
                             expr_chain_dict,
                             lang_graph
                         )
-                        if target_attack_step.full_name in attack_step.children:
-                            attack_step.children[target_attack_step.full_name].\
-                                append((target_attack_step, expr_chain))
+
+                        if target_attack_step in attack_step.own_children:
+                            attack_step.own_children[target_attack_step].append(expr_chain)
                         else:
-                            attack_step.children[target_attack_step.full_name] = \
-                                [(target_attack_step, expr_chain)]
+                            attack_step.own_children[target_attack_step] = [expr_chain]
 
-                for parent_target in attack_step_dict['parents'].items():
-                    target_full_attack_step_name = parent_target[0]
-                    expr_chains = parent_target[1]
+                for (target_step_full_name, expr_chains) in attack_step_dict['own_parents'].items():
                     target_asset_name, target_attack_step_name = \
                         disaggregate_attack_step_full_name(
-                            target_full_attack_step_name)
+                            target_step_full_name
+                        )
                     target_asset = lang_graph.assets[target_asset_name]
-                    target_attack_step = target_asset.attack_steps[
-                        target_attack_step_name]
+                    target_attack_step = target_asset.attack_steps[target_attack_step_name]
                     for expr_chain_dict in expr_chains:
                         expr_chain = ExpressionsChain._from_dict(
-                            expr_chain_dict,
-                            lang_graph
+                            expr_chain_dict, lang_graph
                         )
-                        if target_attack_step.full_name in attack_step.parents:
-                            attack_step.parents[target_attack_step.full_name].\
-                                append((target_attack_step, expr_chain))
+
+                        if target_attack_step in attack_step.own_parents:
+                            attack_step.own_parents[target_attack_step].append(
+                                expr_chain
+                            )
                         else:
-                            attack_step.parents[target_attack_step.full_name] = \
-                                [(target_attack_step, expr_chain)]
+                            attack_step.own_parents[target_attack_step] = [expr_chain]
 
                 # Recreate the requirements of exist and notExist attack steps
                 if attack_step.type == 'exist' or \
@@ -1625,6 +1656,7 @@ class LanguageGraph():
         """
         Link assets based on inheritance and associations.
         """
+
         for asset_dict in lang_spec['assets']:
             asset = assets[asset_dict['name']]
             if asset_dict['superAsset']:
@@ -1689,8 +1721,8 @@ class LanguageGraph():
                         attack_step_attribs['reaches']['overrides']
                         if attack_step_attribs['reaches'] else False
                     ),
-                    children = {},
-                    parents = {},
+                    own_children = {},
+                    own_parents = {},
                     info = attack_step_attribs['meta'],
                     tags = list(attack_step_attribs['tags'])
                 )
@@ -1732,8 +1764,8 @@ class LanguageGraph():
                                 asset = asset,
                                 ttc = attack_step.ttc,
                                 overrides = False,
-                                children = {},
-                                parents = {},
+                                own_children = {},
+                                own_parents = {},
                                 info = attack_step.info,
                                 tags = list(attack_step.tags)
                             )
@@ -1799,16 +1831,14 @@ class LanguageGraph():
                         target_attack_step_name]
 
                     # Link to the children target attack steps
-                    attack_step.children.setdefault(target_attack_step.full_name, [])
-                    attack_step.children[target_attack_step.full_name].append(
-                        (target_attack_step, expr_chain)
-                    )
+                    attack_step.own_children.setdefault(target_attack_step, [])
+                    attack_step.own_children[target_attack_step].append(expr_chain)
 
                     # Reverse the children associations chains to get the
                     # parents associations chain.
-                    target_attack_step.parents.setdefault(attack_step.full_name, [])
-                    target_attack_step.parents[attack_step.full_name].append(
-                        (attack_step, self.reverse_expr_chain(expr_chain, None))
+                    target_attack_step.own_parents.setdefault(attack_step, [])
+                    target_attack_step.own_parents[attack_step].append(
+                        self.reverse_expr_chain(expr_chain, None)
                     )
 
                 # Evaluate the requirements of exist and notExist attack steps

mal_toolbox-1.0.3/maltoolbox/patternfinder/attackgraph_patterns.py

@@ -0,0 +1,134 @@
+"""Utilities for finding patterns in the AttackGraph"""
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable
+from maltoolbox.attackgraph import AttackGraph, AttackGraphNode
+
+class SearchPattern:
+    """A pattern consists of conditions, the conditions are used
+    to find all matching sequences of nodes in an AttackGraph."""
+    conditions: list[SearchCondition]
+
+    def __init__(self, conditions):
+        self.conditions = conditions
+
+    def find_matches(self, graph: AttackGraph):
+        """Search through a graph for a pattern using
+        its conditions, and return sequences of nodes
+        that match all the conditions in the pattern
+
+        Args:
+        graph   - The AttackGraph to search in
+    
+        Return: list[list[AttackGraphNode]] matching paths of Nodes
+        """
+
+        # Find the starting nodes which match the first condition
+        condition = self.conditions[0]
+        matching_paths = []
+        for node in graph.nodes.values():
+            if condition.matches(node):
+                matching_paths.extend(
+                    find_matches_recursively(node, self.conditions)
+                )
+        return matching_paths
+
+@dataclass
+class SearchCondition:
+    """A condition that has to be true for a node to match"""
+
+    # Predefined search conditions
+    ANY = lambda _: True
+
+    # `matches` should be a lambda that takes node as input and returns bool
+    # If lamdba returns True for a node, the node matches
+    # If the lamdba returns False for a node, the node does not match
+    matches: Callable[[AttackGraphNode], bool]
+    greedy: bool = False
+
+    # It is possible to require/allow a Condition to repeat
+    min_repeated: int = 1
+    max_repeated: int = 1
+
+    def can_match_again(self, num_matches):
+        """Returns true if condition can be used again"""
+        return num_matches < self.max_repeated
+
+    def must_match_again(self, num_matches):
+        """Returns true if condition must match again to be fulfilled"""
+        return num_matches < self.min_repeated
+
+
+def find_matches_recursively(
+        node: AttackGraphNode,
+        condition_list: list[SearchCondition],
+        current_path: list[AttackGraphNode] | None = None,
+        matching_paths: set[tuple[AttackGraphNode,...]] | None = None,
+        condition_match_count: int = 0
+    ):
+    """Find all paths of nodes that match the list of conditions.
+    When a sequence of conditions is fulfilled for a path of nodes,
+    add the path of nodes to the returned `matching_paths`
+    The function runs recursively down all paths of children nodes.
+
+    Args:
+    node                  - node to check if current `condition` matches for
+    condition_list        - first condition in list will attempt match `node`
+    current_path          - list of matched nodes so far (recursively built)
+    matching_paths        - set of matched paths so far (recursively built)
+    condition_match_count - number of matches on current condition so far
+
+    Return: set of tuples (paths) of AttackGraphNodes that match the condition
+    """
+
+    # Init path lists if None, or copy/init into new lists for each iteration
+    current_path = [] if current_path is None else list(current_path)
+    matching_paths = set() if matching_paths is None else matching_paths
+
+    curr_cond, *next_conds = condition_list
+
+    if node in current_path:
+        # Stop the chain, infinite loop
+        return matching_paths
+
+    if next_conds and not curr_cond.must_match_again(condition_match_count):
+        # Try next condition for current node if there are more and
+        # current condition is already fulfilled.
+        matching_paths = find_matches_recursively(
+            node,
+            next_conds,
+            current_path=current_path,
+            matching_paths=matching_paths
+        )
+
+    if curr_cond.matches(node):
+        # Current node matches, add to current_path and increment match_count
+        current_path.append(node)
+        condition_match_count += 1
+
+        if next_conds:
+            # If there are more conditions, try next one for all children
+            for child in node.children:
+                matching_paths = find_matches_recursively(
+                    child,
+                    next_conds,
+                    current_path=current_path,
+                    matching_paths=matching_paths,
+                )
+        if curr_cond.can_match_again(condition_match_count):
+            # If we can match current condition again, try for all children
+            for child in node.children:
+                matching_paths = find_matches_recursively(
+                    child,
+                    [curr_cond] + next_conds,
+                    current_path=current_path,
+                    matching_paths=matching_paths,
+                    condition_match_count=condition_match_count
+                )
+
+        if not next_conds:
+            # Congrats - matched a full unique search pattern!
+            matching_paths.add(tuple(current_path)) # tuple is hashable
+
+    return matching_paths

mal_toolbox-1.0.3/maltoolbox/visualization/graphviz_utils.py

@@ -0,0 +1,102 @@
+import random
+import graphviz
+
+from ..model import Model
+from ..attackgraph import AttackGraph
+
+graphviz_bright_colors = [
+    "aliceblue", "antiquewhite", "antiquewhite1", "antiquewhite2", "azure", "azure1", "azure2",
+    "beige", "bisque", "bisque1", "bisque2", "blanchedalmond",
+    "cornsilk", "cornsilk1", "cornsilk2",
+    "floralwhite", "gainsboro", "ghostwhite", "gold", "gold1", "gold2",
+    "honeydew", "honeydew1", "honeydew2",
+    "ivory", "ivory1", "ivory2",
+    "khaki", "khaki1", "khaki2",
+    "lavender", "lavenderblush", "lavenderblush1", "lavenderblush2",
+    "lemonchiffon", "lemonchiffon1", "lemonchiffon2",
+    "lightblue", "lightblue1", "lightblue2",
+    "lightcyan", "lightcyan1", "lightcyan2",
+    "lightgoldenrod", "lightgoldenrod1", "lightgoldenrod2",
+    "lightgoldenrodyellow", "lightgray",
+    "lightpink", "lightpink1", "lightpink2",
+    "lightsalmon", "lightsalmon1", "lightsalmon2",
+    "lightskyblue", "lightskyblue1", "lightskyblue2",
+    "lightslategray", "lightsteelblue", "lightsteelblue1", "lightsteelblue2",
+    "lightyellow", "lightyellow1", "lightyellow2",
+    "linen", "mintcream", "mistyrose", "mistyrose1", "mistyrose2",
+    "moccasin", "navajowhite", "navajowhite1", "navajowhite2",
+    "oldlace", "papayawhip", "peachpuff", "peachpuff1", "peachpuff2",
+    "pink", "pink1", "pink2", "plum", "plum1", "plum2", "powderblue",
+    "seashell", "seashell1", "seashell2",
+    "snow", "snow1", "snow2",
+    "thistle", "thistle1", "thistle2",
+    "wheat", "wheat1", "wheat2",
+    "white", "whitesmoke", "yellow", "yellow1", "yellow2",
+]
+
+
+def render_model(model: Model):
+    """Render a model in graphviz, create pdf and open it"""
+    dot = graphviz.Digraph(model.name)
+
+    # Create nodes
+    asset_type_colors: dict[str, str] = {}
+    for asset in model.assets.values():
+        bg_color = asset_type_colors.get(asset.lg_asset.name)
+        if not bg_color:
+            bg_color = random.choice(graphviz_bright_colors)
+            asset_type_colors[asset.lg_asset.name] = bg_color
+        dot.node(
+            str(asset.id), asset.name, style="filled", fillcolor=bg_color
+        )
+
+    # Create edges
+    for from_asset in model.assets.values():
+
+        for fieldname, to_assets in from_asset.associated_assets.items():
+            for to_asset in to_assets:
+                dot.edge(
+                    str(from_asset.id), str(to_asset.id), label=fieldname
+                )
+    dot.render(directory='.', view=True)
+
+def render_attack_graph(attack_graph: AttackGraph):
+    """Render attack graph graphviz, create pdf and open it"""
+    assert attack_graph.model, "Attack graph needs a model"
+    dot = graphviz.Graph(attack_graph.model.name)
+    dot.graph_attr['nodesep'] = '3.0' # Node separation
+    dot.graph_attr['ratio'] = 'compress'
+
+    # Create nodes
+    asset_colors: dict[str, str] = {}
+    for node in attack_graph.nodes.values():
+        assert node.model_asset, "Node needs model"
+        bg_color = asset_colors.get(node.model_asset.name)
+        if not bg_color:
+            bg_color = random.choice(graphviz_bright_colors)
+            asset_colors[node.model_asset.name] = bg_color
+        path_color = 'white'
+        match node.type:
+            case 'defense':
+                path_color = 'blue'
+            case 'or':
+                path_color = 'red'
+            case 'and':
+                path_color = 'red'
+            case 'exist':
+                path_color = 'grey'
+            case 'notExist':
+                path_color = 'grey'
+            case t:
+                raise ValueError(f'Type {t} not supported')
+
+        dot.node(
+            str(node.id), node.full_name, style="filled", color=path_color, fillcolor=bg_color
+        )
+
+    # Create edges
+    for parent in attack_graph.nodes.values():
+        for child in parent.children:
+            dot.edge(str(parent.id), str(child.id))
+
+    dot.render(directory='.', view=True)

{mal_toolbox-1.0.2 → mal_toolbox-1.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mal-toolbox"
-version = "1.0.2"
+version = "1.0.3"
 authors = [
   { name="Andrei Buhaiu", email="buhaiu@kth.se" },
   { name="Joakim Loxdal", email="loxdal@kth.se" },
@@ -12,6 +12,7 @@ description = "A collection of tools used to create MAL models and attack graphs
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+  "graphviz",
   "antlr4-tools",
   "antlr4-python3-runtime",
   "docopt",