mal-toolbox 0.1.12__py3-none-any.whl → 0.3.0__py3-none-any.whl

maltoolbox/attackgraph/node.py

@@ -1,55 +1,68 @@
 """
-MAL-Toolbox Attack Graph Node Dataclass
+MAL-Toolbox Attack Graph Node
 """
 
 from __future__ import annotations
 import copy
-from dataclasses import field, dataclass
-from typing import Any, Optional
+from functools import cached_property
+from typing import TYPE_CHECKING
 
-from . import Attacker
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from . import Attacker
+    from ..language import LanguageGraphAttackStep, Detector
+    from ..model import ModelAsset
 
-@dataclass
 class AttackGraphNode:
     """Node part of AttackGraph"""
-    type: str
-    name: str
-    ttc: Optional[dict] = None
-    id: Optional[int] = None
-    asset: Optional[Any] = None
-    children: list[AttackGraphNode] = field(default_factory=list)
-    parents: list[AttackGraphNode] = field(default_factory=list)
-    defense_status: Optional[float] = None
-    existence_status: Optional[bool] = None
-    is_viable: bool = True
-    is_necessary: bool = True
-    compromised_by: list[Attacker] = field(default_factory=list)
-    mitre_info: Optional[str] = None
-    tags: list[str] = field(default_factory=lambda: [])
-    attributes: Optional[dict] = None
-
-    # Optional extra metadata for AttackGraphNode
-    extras: dict = field(default_factory=dict)
+
+    def __init__(
+        self,
+        node_id: int,
+        lg_attack_step: LanguageGraphAttackStep,
+        model_asset: Optional[ModelAsset] = None,
+        defense_status: Optional[float] = None,
+        existence_status: Optional[bool] = None
+    ):
+        self.lg_attack_step = lg_attack_step
+        self.name = lg_attack_step.name
+        self.type = lg_attack_step.type
+        self.ttc = lg_attack_step.ttc
+        self.tags = lg_attack_step.tags
+        self.detectors = lg_attack_step.detectors
+
+        self.id = node_id
+        self.model_asset = model_asset
+        self.defense_status = defense_status
+        self.existence_status = existence_status
+
+        self.children: set[AttackGraphNode] = set()
+        self.parents: set[AttackGraphNode] = set()
+        self.is_viable: bool = True
+        self.is_necessary: bool = True
+        self.compromised_by: set[Attacker] = set()
+        self.extras: dict = {}
 
     def to_dict(self) -> dict:
         """Convert node to dictionary"""
         node_dict: dict = {
             'id': self.id,
             'type': self.type,
+            'lang_graph_attack_step': self.lg_attack_step.full_name,
             'name': self.name,
             'ttc': self.ttc,
-            'children': {},
-            'parents': {},
+            'children': {child.id: child.full_name for child in
+                self.children},
+            'parents': {parent.id: parent.full_name for parent in
+                self.parents},
             'compromised_by': [attacker.name for attacker in \
                 self.compromised_by]
         }
 
-        for child in self.children:
-            node_dict['children'][child.id] = child.full_name
-        for parent in self.parents:
-            node_dict['parents'][parent.id] = parent.full_name
-        if self.asset is not None:
-            node_dict['asset'] = str(self.asset.name)
+        for detector in self.detectors.values():
+            node_dict.setdefault('detectors', {})[detector.name] = detector.to_dict()
+        if self.model_asset is not None:
+            node_dict['asset'] = str(self.model_asset.name)
         if self.defense_status is not None:
             node_dict['defense_status'] = str(self.defense_status)
         if self.existence_status is not None:
@@ -58,17 +71,18 @@ class AttackGraphNode:
             node_dict['is_viable'] = str(self.is_viable)
         if self.is_necessary is not None:
             node_dict['is_necessary'] = str(self.is_necessary)
-        if self.mitre_info is not None:
-            node_dict['mitre_info'] = str(self.mitre_info)
         if self.tags:
-            node_dict['tags'] = str(self.tags)
+            node_dict['tags'] = list(self.tags)
         if self.extras:
             node_dict['extras'] = self.extras
 
         return node_dict
 
+
     def __repr__(self) -> str:
-        return str(self.to_dict())
+        return (f'AttackGraphNode(name: "{self.full_name}", id: {self.id}, '
+            f'type: {self.type})')
+
 
     def __deepcopy__(self, memo) -> AttackGraphNode:
         """Deep copy an attackgraph node
@@ -85,33 +99,26 @@ class AttackGraphNode:
             return memo[id(self)]
 
         copied_node = AttackGraphNode(
-            self.type,
-            self.name,
-            self.ttc,
-            self.id,
-            self.asset,
-            [],
-            [],
-            self.defense_status,
-            self.existence_status,
-            self.is_viable,
-            self.is_necessary,
-            [],
-            self.mitre_info,
-            [],
-            {},
-            {}
+            node_id = self.id,
+            model_asset = self.model_asset,
+            lg_attack_step = self.lg_attack_step
         )
 
         copied_node.tags = copy.deepcopy(self.tags, memo)
-        copied_node.attributes = copy.deepcopy(self.attributes, memo)
         copied_node.extras = copy.deepcopy(self.extras, memo)
+        copied_node.ttc = copy.deepcopy(self.ttc, memo)
+
+        copied_node.defense_status = self.defense_status
+        copied_node.existence_status = self.existence_status
+        copied_node.is_viable = self.is_viable
+        copied_node.is_necessary = self.is_necessary
 
         # Remember that self was already copied
         memo[id(self)] = copied_node
 
         return copied_node
 
+
     def is_compromised(self) -> bool:
         """
         Return True if any attackers have compromised this node.
@@ -119,6 +126,7 @@ class AttackGraphNode:
         """
         return len(self.compromised_by) > 0
 
+
     def is_compromised_by(self, attacker: Attacker) -> bool:
         """
         Return True if the attacker given as an argument has compromised this
@@ -130,6 +138,7 @@ class AttackGraphNode:
         """
         return attacker in self.compromised_by
 
+
     def compromise(self, attacker: Attacker) -> None:
         """
         Have the attacker given as a parameter compromise this node.
@@ -139,6 +148,7 @@ class AttackGraphNode:
         """
         attacker.compromise(self)
 
+
     def undo_compromise(self, attacker: Attacker) -> None:
         """
         Remove the attacker given as a parameter from the list of attackers
@@ -150,6 +160,7 @@ class AttackGraphNode:
         """
         attacker.undo_compromise(self)
 
+
     def is_enabled_defense(self) -> bool:
         """
         Return True if this node is a defense node and it is enabled and not
@@ -160,6 +171,7 @@ class AttackGraphNode:
             'suppress' not in self.tags and \
             self.defense_status == 1.0
 
+
     def is_available_defense(self) -> bool:
         """
         Return True if this node is a defense node and it is not fully enabled
@@ -169,6 +181,7 @@ class AttackGraphNode:
             'suppress' not in self.tags and \
             self.defense_status != 1.0
 
+
     @property
     def full_name(self) -> str:
         """
@@ -176,8 +189,13 @@ class AttackGraphNode:
         asset name to which the attack step belongs and attack step name
         itself.
         """
-        if self.asset:
-            full_name = self.asset.name + ':' + self.name
+        if self.model_asset:
+            full_name = self.model_asset.name + ':' + self.name
         else:
             full_name = str(self.id) + ':' + self.name
         return full_name
+
+
+    @cached_property
+    def info(self) -> dict[str, str]:
+        return self.lg_attack_step.info

maltoolbox/attackgraph/query.py

@@ -186,7 +186,8 @@ def get_defense_surface(graph: AttackGraph) -> list[AttackGraphNode]:
     graph       - the attack graph
     """
     logger.debug('Get the defense surface.')
-    return [node for node in graph.nodes if node.is_available_defense()]
+    return [node for node in graph.nodes.values()
+        if node.is_available_defense()]
 
 def get_enabled_defenses(graph: AttackGraph) -> list[AttackGraphNode]:
     """
@@ -197,4 +198,5 @@ def get_enabled_defenses(graph: AttackGraph) -> list[AttackGraphNode]:
     graph       - the attack graph
     """
     logger.debug('Get the enabled defenses.')
-    return [node for node in graph.nodes if node.is_enabled_defense()]
+    return [node for node in graph.nodes.values()
+        if node.is_enabled_defense()]

maltoolbox/file_utils.py

@@ -2,7 +2,6 @@
 
 import json
 import yaml
-from python_jsonschema_objects.literals import LiteralValue
 
 def save_dict_to_json_file(filename: str, serialized_object: dict) -> None:
     """Save serialized object to a json file.
@@ -24,15 +23,12 @@ def save_dict_to_yaml_file(filename: str, serialized_object: dict) -> None:
     data            - dict to output as yaml
     """
 
-    # Handle Literal values from jsonschema_objects
-    yaml.add_multi_representer(
-        LiteralValue,
-        lambda dumper, data: dumper.represent_data(data._value),
-        yaml.SafeDumper
-    )
+    class NoAliasSafeDumper(yaml.SafeDumper):
+        def ignore_aliases(self, data):
+            return True
 
     with open(filename, 'w', encoding='utf-8') as f:
-        yaml.dump(serialized_object, f, Dumper=yaml.SafeDumper)
+        yaml.dump(serialized_object, f, Dumper=NoAliasSafeDumper)
 
 
 def load_dict_from_yaml_file(filename: str) -> dict:

maltoolbox/ingestors/neo4j.py

@@ -7,9 +7,6 @@ import logging
 
 from py2neo import Graph, Node, Relationship, Subgraph
 
-from ..model import AttackerAttachment, Model
-from ..language import LanguageGraph, LanguageClassesFactory
-
 logger = logging.getLogger(__name__)
 
 def ingest_attack_graph(graph,
@@ -90,26 +87,24 @@ def ingest_model(model,
     nodes = {}
     rels = []
 
-    for asset in model.assets:
-
-        nodes[str(asset.id)] = Node(str(asset.type),
-                name=str(asset.name),
-                asset_id=str(asset.id),
-                type=str(asset.type))
-
-    for assoc in model.associations:
-        firstElementName, secondElementName = assoc._properties.keys()
-        firstElements = getattr(assoc, firstElementName)
-        secondElements = getattr(assoc, secondElementName)
-        for first_asset in firstElements:
-            for second_asset in secondElements:
-                rels.append(Relationship(nodes[str(first_asset.id)],
-                    str(firstElementName),
-                    nodes[str(second_asset.id)]))
-                rels.append(Relationship(nodes[str(second_asset.id)],
-                    str(secondElementName),
-                    nodes[str(first_asset.id)]))
+    for asset in model.assets.values():
+        nodes[str(asset.id)] = Node(
+            str(asset.type),
+            name=str(asset.name),
+            asset_id=str(asset.id),
+            type=str(asset.type)
+        )
 
+    for asset in model.assets.values():
+        for fieldname, other_assets in asset.associated_assets.items():
+            for other_asset in other_assets:
+                rels.append(
+                    Relationship(
+                        nodes[str(asset.id)],
+                        str(fieldname),
+                        nodes[str(other_asset.id)]
+                    )
+                )
 
     subgraph = Subgraph(list(nodes.values()), rels)
 
@@ -118,138 +113,132 @@ def ingest_model(model,
     g.commit(tx)
 
 
-def get_model(
-        uri: str,
-        username: str,
-        password: str,
-        dbname: str,
-        lang_graph: LanguageGraph,
-        lang_classes_factory: LanguageClassesFactory
-    ) -> Model:
-    """Load a model from Neo4j"""
-
-    g = Graph(uri=uri, user=username, password=password, name=dbname)
-
-    instance_model = Model('Neo4j imported model', lang_classes_factory)
-    # Get all assets
-    assets_results = g.run('MATCH (a) WHERE a.type IS NOT NULL RETURN DISTINCT a').data()
-    for asset in assets_results:
-        asset_data = dict(asset['a'])
-        logger.debug(
-            'Loading asset from Neo4j instance:\n%s', str(asset_data)
-        )
-        if asset_data['type'] == 'Attacker':
-            attacker_id = int(asset_data['asset_id'])
-            attacker = AttackerAttachment()
-            attacker.entry_points = []
-            instance_model.add_attacker(attacker, attacker_id = attacker_id)
-            continue
-
-        if not hasattr(lang_classes_factory.ns, asset_data['type']):
-            msg = 'Failed to find %s asset in language specification!'
-            logger.error(msg, asset_data["type"])
-            raise LookupError(msg % asset_data["type"])
-
-        asset_obj = getattr(lang_classes_factory.ns,
-            asset_data['type'])(name = asset_data['name'])
-        asset_id = int(asset_data['asset_id'])
-
-        #TODO Process defense values when they are included in Neo4j
-        instance_model.add_asset(asset_obj, asset_id)
-
-    # Get all relationships
-    assocs_results = g.run('MATCH (a)-[r1]->(b),(a)<-[r2]-(b) WHERE a.type IS NOT NULL RETURN DISTINCT a, r1, r2, b').data()
-
-    for assoc in assocs_results:
-        left_field = list(assoc['r1'].types())[0]
-        right_field = list(assoc['r2'].types())[0]
-        left_asset = dict(assoc['a'])
-        right_asset = dict(assoc['b'])
-
-        logger.debug(
-            'Load association ("%s", "%s", "%s", "%s") from Neo4j instance.',
-            left_field, right_field, left_asset["type"], right_asset["type"]
-        )
-
-        left_id = int(left_asset['asset_id'])
-        right_id = int(right_asset['asset_id'])
-
-        attacker_id = None
-        if left_field == 'firstSteps':
-            attacker_id = right_id
-            target_id = left_id
-            target_prop = right_field
-        elif right_field == 'firstSteps':
-            attacker_id = left_id
-            target_id = right_id
-            target_prop = left_field
-
-        if attacker_id is not None:
-            attacker = instance_model.get_attacker_by_id(attacker_id)
-            if not attacker:
-                msg = 'Failed to find attacker with id %s in model!'
-                logger.error(msg, attacker_id)
-                raise LookupError(msg % attacker_id)
-            target_asset = instance_model.get_asset_by_id(target_id)
-            if not target_asset:
-                msg = 'Failed to find asset with id %d in model!'
-                logger.error(msg, target_id)
-                raise LookupError(msg % target_id)
-            attacker.entry_points.append((target_asset,
-                [target_prop]))
-            continue
-
-        left_asset = instance_model.get_asset_by_id(left_id)
-        if left_asset is None:
-            msg = 'Failed to find asset with id %d in model!'
-            logger.error(msg, left_id)
-            raise LookupError(msg % left_id)
-        right_asset = instance_model.get_asset_by_id(right_id)
-        if right_asset is None:
-            msg = 'Failed to find asset with id %d in model!'
-            logger.error(msg, right_id)
-            raise LookupError(msg % right_id)
-
-        assoc = lang_graph.get_association_by_fields_and_assets(
-            left_field,
-            right_field,
-            left_asset.type,
-            right_asset.type)
-
-        if not assoc:
-            logger.error(
-                'Failed to find ("%s", "%s", "%s", "%s")'
-                'association in language specification!',
-                left_asset.type, right_asset.type,
-                left_field, right_field
-            )
-            return None
-
-        logger.debug('Found "%s" association.', assoc.name)
-
-        assoc_name = lang_classes_factory.get_association_by_signature(
-            assoc.name,
-            left_asset.type,
-            right_asset.type
-        )
-
-        if not assoc_name:
-            msg = 'Failed to find \"%s\" association in language specification!'
-            logger.error(msg, assoc.name)
-            raise LookupError(msg % assoc.name)
-
-        assoc = getattr(lang_classes_factory.ns, assoc_name)()
-        setattr(assoc, left_field, [left_asset])
-        setattr(assoc, right_field, [right_asset])
-        if not (instance_model.association_exists_between_assets(
-            assoc_name,
-            left_asset,
-            right_asset
-        ) or instance_model.association_exists_between_assets(
-            assoc_name,
-            right_asset,
-            left_asset
-        )):
-            instance_model.add_association(assoc)
-
-    return instance_model
+# def get_model(
+#         uri: str,
+#         username: str,
+#         password: str,
+#         dbname: str,
+#         lang_graph: LanguageGraph,
+#     ) -> Model:
+#     """Load a model from Neo4j"""
+
+#     g = Graph(uri=uri, user=username, password=password, name=dbname)
+
+#     instance_model = Model('Neo4j imported model', lang_graph)
+#     # Get all assets
+#     assets_results = g.run('MATCH (a) WHERE a.type IS NOT NULL RETURN DISTINCT a').data()
+#     for asset in assets_results:
+#         asset_data = dict(asset['a'])
+#         logger.debug(
+#             'Loading asset from Neo4j instance:\n%s', str(asset_data)
+#         )
+#         if asset_data['type'] == 'Attacker':
+#             attacker_id = int(asset_data['asset_id'])
+#             attacker = AttackerAttachment()
+#             attacker.entry_points = []
+#             instance_model.add_attacker(attacker, attacker_id = attacker_id)
+#             continue
+
+#         asset_id = int(asset_data['asset_id'])
+
+#         #TODO Process defense values when they are included in Neo4j
+#         instance_model.add_asset(asset_data['type'], asset_id=asset_id)
+
+#     # Get all relationships
+#     assocs_results = g.run(
+#         'MATCH (a)-[r1]->(b),(a)<-[r2]-(b) WHERE a.type IS NOT NULL RETURN DISTINCT a, r1, r2, b'
+#     ).data()
+
+#     for assoc in assocs_results:
+#         left_field = list(assoc['r1'].types())[0]
+#         right_field = list(assoc['r2'].types())[0]
+#         left_asset = dict(assoc['a'])
+#         right_asset = dict(assoc['b'])
+
+#         logger.debug(
+#             'Load association ("%s", "%s", "%s", "%s") from Neo4j instance.',
+#             left_field, right_field, left_asset["type"], right_asset["type"]
+#         )
+
+#         left_id = int(left_asset['asset_id'])
+#         right_id = int(right_asset['asset_id'])
+
+#         attacker_id = None
+#         if left_field == 'firstSteps':
+#             attacker_id = right_id
+#             target_id = left_id
+#             target_prop = right_field
+#         elif right_field == 'firstSteps':
+#             attacker_id = left_id
+#             target_id = right_id
+#             target_prop = left_field
+
+#         if attacker_id is not None:
+#             attacker = instance_model.get_attacker_by_id(attacker_id)
+#             if not attacker:
+#                 msg = 'Failed to find attacker with id %s in model!'
+#                 logger.error(msg, attacker_id)
+#                 raise LookupError(msg % attacker_id)
+#             target_asset = instance_model.get_asset_by_id(target_id)
+#             if not target_asset:
+#                 msg = 'Failed to find asset with id %d in model!'
+#                 logger.error(msg, target_id)
+#                 raise LookupError(msg % target_id)
+#             attacker.entry_points.append((target_asset,
+#                 [target_prop]))
+#             continue
+
+#         left_asset = instance_model.get_asset_by_id(left_id)
+#         if left_asset is None:
+#             msg = 'Failed to find asset with id %d in model!'
+#             logger.error(msg, left_id)
+#             raise LookupError(msg % left_id)
+#         right_asset = instance_model.get_asset_by_id(right_id)
+#         if right_asset is None:
+#             msg = 'Failed to find asset with id %d in model!'
+#             logger.error(msg, right_id)
+#             raise LookupError(msg % right_id)
+
+#         assoc = lang_graph.get_association_by_fields_and_assets(
+#             left_field,
+#             right_field,
+#             left_asset.type,
+#             right_asset.type)
+
+#         if not assoc:
+#             logger.error(
+#                 'Failed to find ("%s", "%s", "%s", "%s")'
+#                 'association in language specification!',
+#                 left_asset.type, right_asset.type,
+#                 left_field, right_field
+#             )
+#             return None
+
+#         logger.debug('Found "%s" association.', assoc.name)
+
+#         assoc_name = lang_classes_factory.get_association_by_signature(
+#             assoc.name,
+#             left_asset.type,
+#             right_asset.type
+#         )
+
+#         if not assoc_name:
+#             msg = 'Failed to find \"%s\" association in language specification!'
+#             logger.error(msg, assoc.name)
+#             raise LookupError(msg % assoc.name)
+
+#         assoc = getattr(lang_classes_factory.ns, assoc_name)()
+#         setattr(assoc, left_field, [left_asset])
+#         setattr(assoc, right_field, [right_asset])
+#         if not (instance_model.association_exists_between_assets(
+#             assoc_name,
+#             left_asset,
+#             right_asset
+#         ) or instance_model.association_exists_between_assets(
+#             assoc_name,
+#             right_asset,
+#             left_asset
+#         )):
+#             instance_model.add_association(assoc)
+
+#     return instance_model

maltoolbox/language/__init__.py

@@ -1,4 +1,12 @@
 """Contains tools to process MAL languages"""
 
-from .languagegraph import LanguageGraph
-from .classes_factory import LanguageClassesFactory
+from .languagegraph import (
+    Context,
+    Detector,
+    ExpressionsChain,
+    LanguageGraph,
+    LanguageGraphAsset,
+    LanguageGraphAssociation,
+    LanguageGraphAttackStep,
+    disaggregate_attack_step_full_name,
+)