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

maltoolbox/language/languagegraph.py

@@ -32,6 +32,36 @@ def disaggregate_attack_step_full_name(
     return attack_step_full_name.split(':')
 
 
+@dataclass
+class Detector:
+    context: Context
+    name: Optional[str]
+    type: Optional[str]
+    tprate: Optional[dict]
+
+    def to_dict(self) -> dict:
+        return {
+            "context": self.context.to_dict(),
+            "name": self.name,
+            "type": self.type,
+            "tprate": self.tprate,
+        }
+
+
+class Context(dict):
+    def __init__(self, context) -> None:
+        super().__init__(context)
+        self._context_dict = context
+        for label, asset in context.items():
+            setattr(self, label, asset)
+
+    def to_dict(self) -> dict:
+        return {label: asset.name for label, asset in self.items()}
+
+    def __str__(self) -> str:
+        return str({label: asset.name for label, asset in self._context_dict.items()})
+
+
 @dataclass
 class LanguageGraphAsset:
     name: str
@@ -60,8 +90,8 @@ class LanguageGraphAsset:
             'is_abstract': self.is_abstract
         }
 
-        for assoc in self.own_associations.values():
-            node_dict['associations'][assoc.full_name] = assoc.to_dict()
+        for fieldname, assoc in self.own_associations.items():
+            node_dict['associations'][fieldname] = assoc.to_dict()
         for attack_step in self.attack_steps.values():
             node_dict['attack_steps'][attack_step.name] = \
                 attack_step.to_dict()
@@ -75,7 +105,7 @@ class LanguageGraphAsset:
 
 
     def __repr__(self) -> str:
-        return str(self.to_dict())
+        return f'LanguageGraphAsset(name: "{self.name}")'
 
 
     def __hash__(self):
@@ -101,6 +131,7 @@ class LanguageGraphAsset:
             current_asset = current_asset.own_super_asset
         return False
 
+
     @cached_property
     def sub_assets(self) -> set[LanguageGraphAsset]:
         """
@@ -244,7 +275,9 @@ class LanguageGraphAssociation:
 
 
     def __repr__(self) -> str:
-        return str(self.to_dict())
+        return (f'LanguageGraphAssociation(name: "{self.name}", '
+            f'left_field: {self.left_field}, '
+            f'right_field: {self.right_field})')
 
 
     @property
@@ -262,6 +295,15 @@ class LanguageGraphAssociation:
         return full_name
 
 
+    def get_field(self, fieldname: str) -> LanguageGraphAssociationField:
+        """
+        Return the field that matches the `fieldname` given as parameter.
+        """
+        if self.right_field.fieldname == fieldname:
+            return self.right_field
+        return self.left_field
+
+
     def contains_fieldname(self, fieldname: str) -> bool:
         """
         Check if the association contains the field name given as a parameter.
@@ -317,35 +359,6 @@ class LanguageGraphAssociation:
         raise LanguageGraphAssociationError(msg % (fieldname, self.name))
 
 
-    def get_opposite_asset(
-            self, asset: LanguageGraphAsset
-        ) -> Optional[LanguageGraphAsset]:
-        """
-        Return the opposite asset if the association matches the asset given
-        as a parameter. A match can either be an explicit one or if the asset
-        given subassets either of the two assets that are part of the
-        association.
-
-        Arguments:
-        asset       - the asset to look for
-        Return the other asset if the parameter matched either of the
-        two. None, otherwise.
-        """
-        #TODO Should check to see which one is the tightest fit for
-        #     associations between assets on different levels of the same
-        #     inheritance chain.
-        if asset.is_subasset_of(self.left_field.asset):
-            return self.right_field.asset
-        if asset.is_subasset_of(self.right_field.asset):
-            return self.left_field.asset
-
-        logger.warning(
-            'Requested asset "%s" from association %s'
-            'which did not contain it!', asset.name, self.name
-        )
-        return None
-
-
 @dataclass
 class LanguageGraphAttackStep:
     name: str
@@ -359,6 +372,7 @@ class LanguageGraphAttackStep:
     inherits: Optional[LanguageGraphAttackStep] = None
     tags: set = field(default_factory = set)
     _attributes: Optional[dict] = None
+    detectors: dict = field(default_factory = lambda: {})
 
 
     def __hash__(self):
@@ -386,7 +400,9 @@ class LanguageGraphAttackStep:
             'info': self.info,
             'overrides': self.overrides,
             'inherits': self.inherits.full_name if self.inherits else None,
-            'tags': list(self.tags)
+            'tags': list(self.tags),
+            'detectors': {label: detector.to_dict() for label, detector in
+            self.detectors.items()},
         }
 
         for child in self.children:
@@ -484,7 +500,7 @@ class ExpressionsChain:
                     )
 
                 return {
-                    self.association.full_name:
+                    self.association.name:
                     {
                         'fieldname': self.fieldname,
                         'asset type': asset_type
@@ -527,7 +543,7 @@ class ExpressionsChain:
             serialized_expr_chain: dict,
             lang_graph: LanguageGraph,
         ) -> Optional[ExpressionsChain]:
-        """Create LanguageGraph from dict
+        """Create ExpressionsChain from dict
         Args:
         serialized_expr_chain   - expressions chain in dict format
         lang_graph              - the LanguageGraph that contains the assets,
@@ -567,10 +583,26 @@ class ExpressionsChain:
                 assoc_name = list(serialized_expr_chain.keys())[0]
                 target_asset = lang_graph.assets[\
                     serialized_expr_chain[assoc_name]['asset type']]
+                fieldname = serialized_expr_chain[assoc_name]['fieldname']
+
+                association = None
+                for assoc in target_asset.associations.values():
+                    if assoc.contains_fieldname(fieldname) and \
+                            assoc.name == assoc_name:
+                        association = assoc
+                        break
+
+                if association is None:
+                    msg = 'Failed to find association "%s" with '\
+                        'fieldname "%s"'
+                    logger.error(msg % (assoc_name, fieldname))
+                    raise LanguageGraphException(msg % (assoc_name,
+                        fieldname))
+
                 new_expr_chain = ExpressionsChain(
                     type = 'field',
-                    association = target_asset.associations[assoc_name],
-                    fieldname = serialized_expr_chain[assoc_name]['fieldname']
+                    association = association,
+                    fieldname = fieldname
                 )
                 return new_expr_chain
 
@@ -629,6 +661,11 @@ class LanguageGraph():
             self._generate_graph()
 
 
+    def __repr__(self) -> str:
+        return (f'LanguageGraph(id: "{self.metadata.get("id", "N/A")}", '
+            f'version: "{self.metadata.get("version", "N/A")}")')
+
+
     @classmethod
     def from_mal_spec(cls, mal_spec_file: str) -> LanguageGraph:
         """
@@ -663,12 +700,18 @@ class LanguageGraph():
             'Serializing %s assets.', len(self.assets.items())
         )
 
-        serialized_graph = {}
+        serialized_graph = {'metadata': self.metadata}
         for asset in self.assets.values():
             serialized_graph[asset.name] = asset.to_dict()
 
         return serialized_graph
 
+    def _link_association_to_assets(cls,
+            assoc: LanguageGraphAssociation,
+            left_asset: LanguageGraphAsset,
+            right_asset: LanguageGraphAsset):
+        left_asset.own_associations[assoc.right_field.fieldname] = assoc
+        right_asset.own_associations[assoc.left_field.fieldname] = assoc
 
     def save_to_file(self, filename: str) -> None:
         """Save to json/yml depending on extension"""
@@ -684,6 +727,8 @@ class LanguageGraph():
 
         logger.debug('Create language graph from dictionary.')
         lang_graph = LanguageGraph()
+        lang_graph.metadata = serialized_graph.pop('metadata')
+
         # Recreate all of the assets
         for asset_dict in serialized_graph.values():
             logger.debug(
@@ -766,11 +811,8 @@ class LanguageGraph():
                 )
 
                 # Add the association to the left and right asset
-                associated_assets = [left_asset, right_asset]
-                while associated_assets != []:
-                    asset = associated_assets.pop()
-                    if assoc_node.full_name not in asset.own_associations:
-                        asset.own_associations[assoc_node.full_name] = assoc_node
+                lang_graph._link_association_to_assets(assoc_node,
+                    left_asset, right_asset)
 
         # Recreate the variables
         for asset_dict in serialized_graph.values():
@@ -1360,11 +1402,8 @@ class LanguageGraph():
                 )
 
                 # Add the association to the left and right asset
-                associated_assets = [left_asset, right_asset]
-                while associated_assets != []:
-                    asset = associated_assets.pop()
-                    if assoc_node.full_name not in asset.own_associations:
-                        asset.own_associations[assoc_node.full_name] = assoc_node
+                self._link_association_to_assets(assoc_node,
+                    left_asset, right_asset)
 
         # Set the variables
         for asset in self.assets.values():
@@ -1407,6 +1446,20 @@ class LanguageGraph():
                 asset.attack_steps[attack_step_attribs['name']] = \
                     attack_step_node
 
+                for detector in attack_step_attribs.get("detectors",
+                                                        {}).values():
+                    attack_step_node.detectors[detector["name"]] = Detector(
+                        context=Context(
+                            {
+                                label: self.assets[asset]
+                                for label, asset in detector["context"].items()
+                            }
+                        ),
+                        name=detector.get("name"),
+                        type=detector.get("type"),
+                        tprate=detector.get("tprate"),
+                    )
+
         # Create the inherited attack steps
         assets = list(self.assets.values())
         while len(assets) > 0:
@@ -1540,8 +1593,7 @@ class LanguageGraph():
                                 None,
                                 step_expression
                             )
-                        attack_step.own_requires.append(
-                            self.reverse_expr_chain(result_expr_chain, None))
+                        attack_step.own_requires.append(result_expr_chain)
 
     def _get_attacks_for_asset_type(self, asset_type: str) -> dict:
         """
@@ -1681,51 +1733,3 @@ class LanguageGraph():
         self._generate_graph()
 
 
-    def get_association_by_fields_and_assets(
-            self,
-            first_field: str,
-            second_field: str,
-            first_asset_name: str,
-            second_asset_name: str
-        ) -> Optional[LanguageGraphAssociation]:
-        """
-        Get an association based on its field names and asset types
-
-        Arguments:
-        first_field         - a string containing the first field
-        second_field        - a string containing the second field
-        first_asset_name    - a string representing the first asset type
-        second_asset_name   - a string representing the second asset type
-
-        Return:
-        The association matching the fieldnames and asset types.
-        None if there is no match.
-        """
-        first_asset = self.assets[first_asset_name]
-        second_asset = self.assets[second_asset_name]
-
-        for assoc in first_asset.associations.values():
-            logger.debug(
-                'Compare ("%s", "%s", "%s", "%s") to '
-                '("%s", "%s", "%s", "%s").',
-                first_asset_name, first_field,
-                second_asset_name, second_field,
-                assoc.left_field.asset.name, assoc.left_field.fieldname,
-                assoc.right_field.asset.name, assoc.right_field.fieldname
-            )
-
-            # If the asset and fields match either way we accept it as a
-            # match.
-            if assoc.left_field.fieldname == first_field and \
-                assoc.right_field.fieldname == second_field and \
-                first_asset.is_subasset_of(assoc.left_field.asset) and \
-                second_asset.is_subasset_of(assoc.right_field.asset):
-                return assoc
-
-            if assoc.left_field.fieldname == second_field and \
-                assoc.right_field.fieldname == first_field and \
-                second_asset.is_subasset_of(assoc.left_field.asset) and \
-                first_asset.is_subasset_of(assoc.right_field.asset):
-                return assoc
-
-        return None