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

maltoolbox/ingestors/neo4j.py

@@ -1,23 +1,24 @@
 """
 MAL-Toolbox Neo4j Ingestor Module
 """
+# mypy: ignore-errors
 
 import logging
 
 from py2neo import Graph, Node, Relationship, Subgraph
 
-from maltoolbox.model import model
-from maltoolbox.language import specification
+from ..model import AttackerAttachment, Model
+from ..language import LanguageGraph, LanguageClassesFactory
 
 logger = logging.getLogger(__name__)
 
 def ingest_attack_graph(graph,
-        uri,
-        username,
-        password,
-        dbname,
-        delete=False
-) -> None:
+        uri: str,
+        username: str,
+        password: str,
+        dbname: str,
+        delete: bool = False
+    ) -> None:
     """
     Ingest an attack graph into a neo4j database
 
@@ -31,7 +32,6 @@ def ingest_attack_graph(graph,
                            before ingesting the new attack graph
     """
 
-
     g = Graph(uri=uri, user=username, password=password, name=dbname)
     if delete:
         g.delete_all()
@@ -43,7 +43,7 @@ def ingest_attack_graph(graph,
         nodes[node.id] = Node(
             node_dict['asset'] if 'asset' in node_dict else node_dict['id'],
             name = node_dict['name'],
-            full_name = node_dict['id'],
+            full_name = node.full_name,
             type = node_dict['type'],
             ttc = str(node_dict['ttc']),
             is_necessary = str(node.is_necessary),
@@ -65,12 +65,12 @@ def ingest_attack_graph(graph,
 
 
 def ingest_model(model,
-        uri,
-        username,
-        password,
-        dbname,
-        delete=False
-) -> None:
+        uri: str,
+        username: str,
+        password: str,
+        dbname: str,
+        delete: bool = False
+    ) -> None:
     """
     Ingest an instance model graph into a Neo4J database
 
@@ -91,16 +91,14 @@ def ingest_model(model,
     rels = []
 
     for asset in model.assets:
-        nodeid = asset.name
 
-        nodes[str(asset.id)] = Node(str(asset.metaconcept),
+        nodes[str(asset.id)] = Node(str(asset.type),
                 name=str(asset.name),
                 asset_id=str(asset.id),
-                metaconcept=str(asset.metaconcept))
+                type=str(asset.type))
 
     for assoc in model.associations:
-        firstElementName = list(vars(assoc)['_properties'])[0]
-        secondElementName = list(vars(assoc)['_properties'])[1]
+        firstElementName, secondElementName = assoc._properties.keys()
         firstElements = getattr(assoc, firstElementName)
         secondElements = getattr(assoc, secondElementName)
         for first_asset in firstElements:
@@ -120,46 +118,47 @@ def ingest_model(model,
     g.commit(tx)
 
 
-def get_model(uri,
-        username,
-        password,
-        dbname,
-        lang_spec,
-        lang_classes_factory
-) -> model.Model:
+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.Model('Neo4j imported model', lang_spec,
-        lang_classes_factory)
+    instance_model = Model('Neo4j imported model', lang_classes_factory)
     # Get all assets
-    assets_results = g.run('MATCH (a) WHERE a.metaconcept IS NOT NULL RETURN DISTINCT a').data()
+    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' \
-            + str(asset_data))
-        if asset_data['metaconcept'] == 'Attacker':
+        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 = model.Attacker()
+            attacker = AttackerAttachment()
             attacker.entry_points = []
             instance_model.add_attacker(attacker, attacker_id = attacker_id)
             continue
 
-        if not hasattr(lang_classes_factory.ns,
-            asset_data['metaconcept']):
-            logger.error(f'Failed to find {asset_data["metaconcept"]} '
-                'asset in language specification!')
-            return None
+        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['metaconcept'])(name = asset_data['name'])
+            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.metaconcept IS NOT NULL RETURN DISTINCT a, r1, r2, b').data()
+    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]
@@ -167,15 +166,14 @@ def get_model(uri,
         left_asset = dict(assoc['a'])
         right_asset = dict(assoc['b'])
 
-        logger.debug(f'Load association '
-            f'(\"{left_field}\",'
-            f'\"{right_field}\",'
-            f'\"{left_asset["metaconcept"]}\",'
-            f'\"{right_asset["metaconcept"]}\") '
-            f'from Neo4j instance.')
+        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
@@ -186,59 +184,72 @@ def get_model(uri,
             target_id = right_id
             target_prop = left_field
 
-        if attacker_id:
+        if attacker_id is not None:
             attacker = instance_model.get_attacker_by_id(attacker_id)
             if not attacker:
-                logger.error(f'Failed to find attacker with id {attacker_id} '
-                    'in model!')
-                return None
+                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:
-                logger.error(f'Failed to find asset with id {target_id} '
-                    'in model!')
-                return None
+                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 not left_asset:
-            logger.error(f'Failed to find asset with id {left_id} '
-                'in model!')
-            return None
+        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 not right_asset:
-            logger.error(f'Failed to find asset with id {right_id} '
-                'in model!')
-            return None
+        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_name = specification.get_association_by_fields_and_assets(
-            lang_spec,
+        assoc = lang_graph.get_association_by_fields_and_assets(
             left_field,
             right_field,
-            left_asset.metaconcept,
-            right_asset.metaconcept)
-        logger.debug(f'Found \"{assoc_name}\" association.')
-
-        if not assoc_name:
-            logger.error(f'Failed to find '
-                f'(\"{left_asset.metaconcept}\",'
-                f'\"{right_asset.metaconcept}\",'
-                f'\"{left_field}\",'
-                f'\"{right_field}\") '
-                'association in language specification!')
+            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
 
-        if not hasattr(lang_classes_factory.ns,
-            assoc_name):
-            logger.error(f'Failed to find {assoc_name} '
-                'association in language specification!')
-            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])
-        instance_model.add_association(assoc)
+        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

@@ -0,0 +1,4 @@
+"""Contains tools to process MAL languages"""
+
+from .languagegraph import LanguageGraph
+from .classes_factory import LanguageClassesFactory

maltoolbox/language/classes_factory.py

@@ -1,43 +1,59 @@
 """
 MAL-Toolbox Language Classes Factory Module
+Uses python_jsonschema_objects to generate python classes from a MAL language
 """
-
-import python_jsonschema_objects as pjs
+from __future__ import annotations
 import json
 import logging
+from typing import TYPE_CHECKING
+
+import python_jsonschema_objects as pjs
+
+if TYPE_CHECKING:
+    from typing import Literal, Optional, TypeAlias
+    from maltoolbox.language import LanguageGraph
+    from python_jsonschema_objects.classbuilder import ProtocolBase
+
+    SchemaGeneratedClass: TypeAlias = ProtocolBase
 
 logger = logging.getLogger(__name__)
 
 class LanguageClassesFactory:
-    def __init__(self, langspec):
-        self.langspec = langspec
-        self.json_schema = {}
+    def __init__(self, lang_graph: LanguageGraph):
+        self.lang_graph: LanguageGraph = lang_graph
+        self.json_schema: dict = {}
+        self._create_classes()
 
-    def generate_assets(self):
+    def _generate_assets(self) -> None:
         """
-        Generate JSON Schema for the assets in the language specification.
+        Generate JSON Schema for asset types in the language specification.
         """
-        for asset in self.langspec['assets']:
-            logger.debug(f'Creating {asset["name"]} asset JSON '\
-            'schema entry.')
-            asset_json_entry = {'title': asset['name'], 'type': 'object',
-                'properties': {}}
-            asset_json_entry['properties']['metaconcept'] = \
+        for asset in self.lang_graph.assets:
+            logger.debug('Creating %s asset JSON schema entry.', asset.name)
+            asset_json_entry = {
+                'title': asset.name,
+                'type': 'object',
+                'properties': {},
+            }
+            asset_json_entry['properties']['id'] = {
+                'type' : 'integer',
+            }
+            asset_json_entry['properties']['type'] = \
                 {
                     'type' : 'string',
-                    'default': asset['name']
+                    'default': asset.name
                 }
-            if asset['superAsset']:
-                asset_json_entry['allOf'] = [{
-                    '$ref': '#/definitions/LanguageAsset/definitions/' + asset['superAsset']
-                }]
-            for defense in filter(lambda item: item['type'] == 'defense',
-                asset['attackSteps']):
-                if defense['ttc'] and defense['ttc']['name'] == 'Enabled':
+            if asset.super_assets:
+                asset_json_entry['allOf'] = [
+                    {'$ref': '#/definitions/LanguageAsset/definitions/' + superasset.name}
+                    for superasset in asset.super_assets
+                ]
+            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']] = \
+                asset_json_entry['properties'][defense.name] = \
                     {
                         'type' : 'number',
                         'minimum' : 0,
@@ -45,89 +61,93 @@ class LanguageClassesFactory:
                         'default': default_defense_value
                     }
             self.json_schema['definitions']['LanguageAsset']['definitions']\
-                [asset['name']] = asset_json_entry
-            self.json_schema['definitions']['LanguageAsset']['oneOf'].\
-                append({'$ref': '#/definitions/LanguageAsset/definitions/'
-                + asset['name']})
+                [asset.name] = asset_json_entry
+            self.json_schema['definitions']['LanguageAsset']['oneOf'].append(
+                {'$ref': '#/definitions/LanguageAsset/definitions/' + asset.name}
+            )
 
-    def generate_associations(self):
+    def _generate_associations(self) -> None:
         """
-        Generate JSON Schema for the associations in the language specification.
+        Generate JSON Schema for association types in the language specification.
         """
-        def create_association_entry(assoc):
-            logger.debug(f'Creating {assoc["name"]} association JSON '\
-            'schema entry.')
-            assoc_json_entry = {'title': assoc['name'], 'type': 'object',
-                'properties': {}}
+        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': {}}
 
             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):
-            if (assoc['name'] not in self.json_schema['definitions']\
+        def create_association_with_subentries(assoc: SchemaGeneratedClass):
+            if (assoc.name not in self.json_schema['definitions']\
                 ['LanguageAssociation']['definitions']):
-                logger.info('Multiple associations with the same '\
-                    f'name, {assoc["name"]}, exist. '\
-                    'Creating subentries for each one.')
+                logger.info(
+                    'Multiple associations with name %s exist. '
+                    'Creating subentries for each one.', assoc.name
+                )
                 self.json_schema['definitions']['LanguageAssociation']\
-                ['definitions'][assoc['name']] =\
+                ['definitions'][assoc.name] =\
                     {
-                        'title': assoc['name'],
+                        'title': assoc.name,
                         'type': 'object',
                         'oneOf': [],
                         'definitions': {}
                     }
                 self.json_schema['definitions']['LanguageAssociation']['oneOf'].\
                     append({'$ref': '#/definitions/LanguageAssociation/definitions/'
-                    + assoc['name']})
+                    + assoc.name})
 
             assoc_json_subentry = create_association_entry(assoc)
-            subentry_name = assoc['name'] + '_' + assoc['leftAsset'] + '_' \
-                + assoc['rightAsset']
+            subentry_name = assoc.name + '_' + assoc.left_field.asset.name + '_' \
+                + assoc.right_field.asset.name
 
-            logger.info(f'Creating {subentry_name} subentry association.')
+            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
+                ['definitions'][assoc.name]['definitions'][subentry_name] = assoc_json_subentry
             self.json_schema['definitions']['LanguageAssociation']\
-                ['definitions'][assoc['name']]['oneOf'].append(
+                ['definitions'][assoc.name]['oneOf'].append(
                     {'$ref': '#/definitions/LanguageAssociation/definitions/' \
-                    + assoc['name'] + '/definitions/' + subentry_name})
-
-        def create_association_field(assoc, assoc_json_entry, position):
-            assoc_json_entry['properties'][assoc[position + 'Field']] = \
+                    + assoc.name + '/definitions/' + subentry_name})
+
+        def create_association_field(
+                assoc: SchemaGeneratedClass,
+                assoc_json_entry: dict,
+                position: Literal['left', 'right']
+            ) -> None:
+            field = getattr(assoc, position + "_field")
+            assoc_json_entry['properties'][field.fieldname] = \
                 {
                     'type' : 'array',
                     'items' :
                         {
                         '$ref':
                         '#/definitions/LanguageAsset/definitions/' +
-                            assoc[position + 'Asset']
+                            field.asset.name
                         }
                 }
-            if assoc[position + 'Multiplicity']['max']:
-                assoc_json_entry['properties'][assoc[position + 'Field']]\
-                    ['maxItems'] = assoc[position + 'Multiplicity']['max']
+            if field.maximum:
+                assoc_json_entry['properties'][field.fieldname]\
+                    ['maxItems'] = field.maximum
 
 
-        for assoc in self.langspec['associations']:
-            count = len(list(filter(lambda temp_assoc: temp_assoc['name'] ==
-                assoc['name'], self.langspec['associations'])))
+        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.
-                assoc_json_entry = create_association_with_subentries(assoc)
+                create_association_with_subentries(assoc)
             else:
                 assoc_json_entry = create_association_entry(assoc)
                 self.json_schema['definitions']['LanguageAssociation']\
-                    ['definitions'][assoc['name']] = assoc_json_entry
+                    ['definitions'][assoc.name] = assoc_json_entry
                 self.json_schema['definitions']['LanguageAssociation']['oneOf'].\
                     append({'$ref': '#/definitions/LanguageAssociation/' +
-                    'definitions/' + assoc['name']})
+                    'definitions/' + assoc.name})
 
-    def create_classes(self):
+    def _create_classes(self) -> None:
         """
         Create classes based on the language specification.
         """
@@ -135,6 +155,8 @@ class LanguageClassesFactory:
         # First, we have to translate the language specification into a JSON
         # schema. Initialize the overall JSON schema structure.
         self.json_schema = {
+            '$schema': 'http://json-schema.org/draft-04/schema#',
+            'id': f"urn:mal:{__name__.replace('.', ':')}",
             'title': 'LanguageObject',
             'type': 'object',
             'oneOf':[
@@ -153,10 +175,69 @@ class LanguageClassesFactory:
             'oneOf': [],
             'definitions': {}}
 
-        self.generate_assets()
-        self.generate_associations()
-        logger.debug(json.dumps(self.json_schema, indent = 2))
+        self._generate_assets()
+        self._generate_associations()
+
+        if logger.isEnabledFor(logging.DEBUG):
+            # Avoid running json.dumps when not in debug
+            logger.debug(json.dumps(self.json_schema, indent = 2))
 
         # Once we have the JSON schema we create the actual classes.
         builder = pjs.ObjectBuilder(self.json_schema)
         self.ns = builder.build_classes(standardize_names=False)
+
+    def get_association_by_signature(
+        self,
+        assoc_name: str,
+        left_asset: str,
+        right_asset: str
+    ) -> Optional[str]:
+        """
+        Get association name based on its signature. This is primarily
+        relevant for getting the exact association full name when multiple
+        associations with the same name exist.
+
+        Arguments:
+        assoc_name          - the association name
+        left_asset          - the name of the left asset type
+        right_asset         - the name of the right asset type
+
+        Return: The matching association name if a match is found.
+        None if there is no match.
+        """
+        lang_assocs_entries = self.json_schema['definitions']\
+                ['LanguageAssociation']['definitions']
+        if not assoc_name in lang_assocs_entries:
+            raise LookupError(
+                'Failed to find "%s" association in the language json '
+                'schema.' % assoc_name
+            )
+        assoc_entry = lang_assocs_entries[assoc_name]
+        # If the association has a oneOf property it should always have more
+        # than just one alternative, but check just in case
+        if 'definitions' in assoc_entry and \
+                len(assoc_entry['definitions']) > 1:
+            full_name = '%s_%s_%s' % (
+                assoc_name,
+                left_asset,
+                right_asset
+            )
+            full_name_flipped = '%s_%s_%s' % (
+                assoc_name,
+                right_asset,
+                left_asset
+            )
+            if not full_name in assoc_entry['definitions']:
+                if not full_name_flipped in assoc_entry['definitions']:
+                    raise LookupError(
+                        'Failed to find "%s" or "%s" association in the '
+                        'language json schema.'
+                        % (full_name,
+                        full_name_flipped)
+                    )
+                else:
+                    return full_name_flipped
+            else:
+                return full_name
+        else:
+            return assoc_name

maltoolbox/language/{lexer_parser/__main__.py → compiler/__init__.py}

@@ -1,21 +1,21 @@
 #!/usr/bin/env python
+# mypy: ignore-errors
+
+import os
+from typing import Optional
 
 from antlr4 import FileStream, CommonTokenStream
 from .mal_lexer import malLexer
 from .mal_parser import malParser
 from .mal_visitor import malVisitor
 
-import sys
-import os
-import json
-
 
 class MalCompiler:
     def __init__(self):
         self.path = None
         self.current_file = None
 
-    def compile(self, malfile: str = None):
+    def compile(self, malfile: Optional[str] = None):
         if not self.path:
             self.path = os.path.dirname(malfile)
 
@@ -30,10 +30,3 @@ class MalCompiler:
         tree = parser.mal()
 
         return malVisitor(compiler=self).visit(tree)
-
-
-if __name__ == "__main__":
-    compiler = MalCompiler()
-
-    with open("new_langspec.json", "w") as f:
-        json.dump(compiler.compile(sys.argv[1]), f, indent=2)

maltoolbox/language/{lexer_parser → compiler}/mal_lexer.py

@@ -1,6 +1,6 @@
+# mypy: ignore-errors
 # Generated from mal.g4 by ANTLR 4.13.1
 from antlr4 import *
-from io import StringIO
 import sys
 if sys.version_info[1] > 5:
     from typing import TextIO

maltoolbox/language/{lexer_parser → compiler}/mal_parser.py

@@ -1,7 +1,7 @@
+# mypy: ignore-errors
 # Generated from mal.g4 by ANTLR 4.13.1
 # encoding: utf-8
 from antlr4 import *
-from io import StringIO
 import sys
 if sys.version_info[1] > 5:
 	from typing import TextIO