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

maltoolbox/attackgraph/node.py

@@ -5,14 +5,17 @@ MAL-Toolbox Attack Graph Node Dataclass
 from __future__ import annotations
 import copy
 from dataclasses import field, dataclass
+from functools import cached_property
 from typing import Any, Optional
 
 from . import Attacker
+from ..language import LanguageGraphAttackStep
 
 @dataclass
 class AttackGraphNode:
     """Node part of AttackGraph"""
     type: str
+    lang_graph_attack_step: LanguageGraphAttackStep
     name: str
     ttc: Optional[dict] = None
     id: Optional[int] = None
@@ -24,18 +27,19 @@ class AttackGraphNode:
     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: [])
+    tags: set[str] = field(default_factory=set)
     attributes: Optional[dict] = None
 
     # Optional extra metadata for AttackGraphNode
     extras: dict = field(default_factory=dict)
 
+
     def to_dict(self) -> dict:
         """Convert node to dictionary"""
         node_dict: dict = {
             'id': self.id,
             'type': self.type,
+            'lang_graph_attack_step': self.lang_graph_attack_step.full_name,
             'name': self.name,
             'ttc': self.ttc,
             'children': {},
@@ -58,18 +62,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())
 
+
     def __deepcopy__(self, memo) -> AttackGraphNode:
         """Deep copy an attackgraph node
 
@@ -86,6 +90,7 @@ class AttackGraphNode:
 
         copied_node = AttackGraphNode(
             self.type,
+            self.lang_graph_attack_step,
             self.name,
             self.ttc,
             self.id,
@@ -97,8 +102,7 @@ class AttackGraphNode:
             self.is_viable,
             self.is_necessary,
             [],
-            self.mitre_info,
-            [],
+            set(),
             {},
             {}
         )
@@ -112,6 +116,7 @@ class AttackGraphNode:
 
         return copied_node
 
+
     def is_compromised(self) -> bool:
         """
         Return True if any attackers have compromised this node.
@@ -119,6 +124,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 +136,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 +146,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 +158,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 +169,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 +179,7 @@ class AttackGraphNode:
             'suppress' not in self.tags and \
             self.defense_status != 1.0
 
+
     @property
     def full_name(self) -> str:
         """
@@ -181,3 +192,8 @@ class AttackGraphNode:
         else:
             full_name = str(self.id) + ':' + self.name
         return full_name
+
+
+    @cached_property
+    def info(self) -> dict[str, str]:
+        return self.lang_graph_attack_step.info

maltoolbox/file_utils.py

@@ -24,15 +24,19 @@ def save_dict_to_yaml_file(filename: str, serialized_object: dict) -> None:
     data            - dict to output as yaml
     """
 
+    class NoAliasSafeDumper(yaml.SafeDumper):
+        def ignore_aliases(self, data):
+            return True
+
     # Handle Literal values from jsonschema_objects
     yaml.add_multi_representer(
         LiteralValue,
         lambda dumper, data: dumper.represent_data(data._value),
-        yaml.SafeDumper
+        NoAliasSafeDumper
     )
 
     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/language/__init__.py

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

maltoolbox/language/classes_factory.py

@@ -28,10 +28,10 @@ class LanguageClassesFactory:
         """
         Generate JSON Schema for asset types in the language specification.
         """
-        for asset in self.lang_graph.assets:
+        for asset_name, asset in self.lang_graph.assets.items():
             logger.debug('Creating %s asset JSON schema entry.', asset.name)
             asset_json_entry = {
-                'title': asset.name,
+                'title': 'Asset_' + asset.name,
                 'type': 'object',
                 'properties': {},
             }
@@ -41,29 +41,30 @@ class LanguageClassesFactory:
             asset_json_entry['properties']['type'] = \
                 {
                     'type' : 'string',
-                    'default': asset.name
+                    'default': asset_name
                 }
-            if asset.super_assets:
+            if asset.own_super_asset:
                 asset_json_entry['allOf'] = [
-                    {'$ref': '#/definitions/LanguageAsset/definitions/' + superasset.name}
-                    for superasset in asset.super_assets
+                    {'$ref': '#/definitions/LanguageAsset/definitions/Asset_'\
+                        + asset.own_super_asset.name}
                 ]
-            for defense in filter(lambda step: step.type == 'defense', asset.attack_steps):
-                if defense.ttc and defense.ttc['name'] == 'Enabled':
-                    default_defense_value = 1.0
-                else:
-                    default_defense_value = 0.0
-                asset_json_entry['properties'][defense.name] = \
-                    {
-                        'type' : 'number',
-                        'minimum' : 0,
-                        'maximum' : 1,
-                        'default': default_defense_value
-                    }
+            for step_name, step in asset.attack_steps.items():
+                if step.type == 'defense':
+                    if step.ttc and step.ttc['name'] == 'Enabled':
+                        default_defense_value = 1.0
+                    else:
+                        default_defense_value = 0.0
+                    asset_json_entry['properties'][step_name] = \
+                        {
+                            'type' : 'number',
+                            'minimum' : 0,
+                            'maximum' : 1,
+                            'default': default_defense_value
+                        }
             self.json_schema['definitions']['LanguageAsset']['definitions']\
-                [asset.name] = asset_json_entry
+                ['Asset_' + asset_name] = asset_json_entry
             self.json_schema['definitions']['LanguageAsset']['oneOf'].append(
-                {'$ref': '#/definitions/LanguageAsset/definitions/' + asset.name}
+                {'$ref': '#/definitions/LanguageAsset/definitions/Asset_' + asset_name}
             )
 
     def _generate_associations(self) -> None:
@@ -72,44 +73,21 @@ class LanguageClassesFactory:
         """
         def create_association_entry(assoc: SchemaGeneratedClass):
             logger.debug('Creating %s association JSON schema entry.', assoc.name)
-            assoc_json_entry = {'title': assoc.name, 'type': 'object', 'properties': {}}
+            assoc_json_entry = {
+                'title': 'Association_' + assoc.full_name,
+                'type': 'object',
+                'properties': {}
+            }
+            assoc_json_entry['properties']['type'] = \
+                {
+                    'type' : 'string',
+                    'default': assoc.name
+                }
 
             create_association_field(assoc, assoc_json_entry, 'left')
             create_association_field(assoc, assoc_json_entry, 'right')
             return assoc_json_entry
 
-        def create_association_with_subentries(assoc: SchemaGeneratedClass):
-            if (assoc.name not in self.json_schema['definitions']\
-                ['LanguageAssociation']['definitions']):
-                logger.info(
-                    'Multiple associations with name %s exist. '
-                    'Creating subentries for each one.', assoc.name
-                )
-                self.json_schema['definitions']['LanguageAssociation']\
-                ['definitions'][assoc.name] =\
-                    {
-                        'title': assoc.name,
-                        'type': 'object',
-                        'oneOf': [],
-                        'definitions': {}
-                    }
-                self.json_schema['definitions']['LanguageAssociation']['oneOf'].\
-                    append({'$ref': '#/definitions/LanguageAssociation/definitions/'
-                    + assoc.name})
-
-            assoc_json_subentry = create_association_entry(assoc)
-            subentry_name = assoc.name + '_' + assoc.left_field.asset.name + '_' \
-                + assoc.right_field.asset.name
-
-            logger.info('Creating %s subentry association.', subentry_name)
-            assoc_json_subentry['title'] = subentry_name
-            self.json_schema['definitions']['LanguageAssociation']\
-                ['definitions'][assoc.name]['definitions'][subentry_name] = assoc_json_subentry
-            self.json_schema['definitions']['LanguageAssociation']\
-                ['definitions'][assoc.name]['oneOf'].append(
-                    {'$ref': '#/definitions/LanguageAssociation/definitions/' \
-                    + assoc.name + '/definitions/' + subentry_name})
-
         def create_association_field(
                 assoc: SchemaGeneratedClass,
                 assoc_json_entry: dict,
@@ -122,7 +100,7 @@ class LanguageClassesFactory:
                     'items' :
                         {
                         '$ref':
-                        '#/definitions/LanguageAsset/definitions/' +
+                        '#/definitions/LanguageAsset/definitions/Asset_' +
                             field.asset.name
                         }
                 }
@@ -130,22 +108,17 @@ class LanguageClassesFactory:
                 assoc_json_entry['properties'][field.fieldname]\
                     ['maxItems'] = field.maximum
 
-
-        for assoc in self.lang_graph.associations:
-            count = len(list(filter(lambda temp_assoc: temp_assoc.name ==
-                assoc.name, self.lang_graph.associations)))
-            if count > 1:
-                # If there are multiple associations with the same name we
-                # will need to create separate entries for each using their
-                # fieldnames.
-                create_association_with_subentries(assoc)
-            else:
-                assoc_json_entry = create_association_entry(assoc)
-                self.json_schema['definitions']['LanguageAssociation']\
-                    ['definitions'][assoc.name] = assoc_json_entry
-                self.json_schema['definitions']['LanguageAssociation']['oneOf'].\
-                    append({'$ref': '#/definitions/LanguageAssociation/' +
-                    'definitions/' + assoc.name})
+        for asset_name, asset in self.lang_graph.assets.items():
+            for assoc_name, assoc in asset.associations.items():
+                if assoc_name not in self.json_schema['definitions']\
+                        ['LanguageAssociation']['definitions']:
+                    assoc_json_entry = create_association_entry(assoc)
+                    self.json_schema['definitions']['LanguageAssociation']\
+                        ['definitions']['Association_' + assoc_name] = \
+                            assoc_json_entry
+                    self.json_schema['definitions']['LanguageAssociation']['oneOf'].\
+                        append({'$ref': '#/definitions/LanguageAssociation/' +
+                        'definitions/Association_' + assoc_name})
 
     def _create_classes(self) -> None:
         """
@@ -241,3 +214,46 @@ class LanguageClassesFactory:
                 return full_name
         else:
             return assoc_name
+
+    def get_asset_class(self,
+            asset_name: str
+        ) -> Optional[SchemaGeneratedClass]:
+        class_name = 'Asset_' + asset_name
+        if hasattr(self.ns, class_name):
+            class_obj = getattr(self.ns, class_name)
+            class_obj.__hash__ = lambda self: hash(self.name)
+            return class_obj
+        else:
+            logger.warning('Could not find Asset "%s" in classes factory.' %
+                asset_name)
+            return None
+
+    def get_association_class(self,
+            assoc_name: str
+        ) -> Optional[SchemaGeneratedClass]:
+        class_name = 'Association_' + assoc_name
+        if hasattr(self.ns, class_name):
+            return getattr(self.ns, class_name)
+        else:
+            logger.warning('Could not find Association "%s" in classes factory.' %
+                assoc_name)
+            return None
+
+    def get_association_class_by_fieldnames(self,
+            assoc_name: str,
+            fieldname1: str,
+            fieldname2: str
+        ) -> Optional[SchemaGeneratedClass]:
+        class_name = 'Association_%s_%s_%s' % (assoc_name,
+            fieldname1, fieldname2)
+        class_name_alt = 'Association_%s_%s_%s' % (assoc_name,
+            fieldname2, fieldname1)
+
+        if hasattr(self.ns, class_name):
+            return getattr(self.ns, class_name)
+        elif hasattr(self.ns, class_name_alt):
+            return getattr(self.ns, class_name_alt)
+        else:
+            logger.warning('Could not find Association "%s" or "%s" in '
+                'classes factory.' % (class_name, class_name_alt))
+            return None