assemblyline-v4-service 4.5.1.dev232__py3-none-any.whl → 4.5.1.dev233__py3-none-any.whl

assemblyline_v4_service/VERSION

@@ -1 +1 @@
-4.5.1.dev232
+4.5.1.dev233

assemblyline_v4_service/common/ontology_helper.py

@@ -1,9 +1,10 @@
 import json
 import os
 from collections import defaultdict
-from typing import Dict, Optional
+from typing import Any, Dict, Optional, List, Tuple
 
 from assemblyline_v4_service.common import helper
+from assemblyline_v4_service.common.result import HEUR_LIST, ResultSection
 
 from assemblyline.common import forge
 from assemblyline.common.dict_utils import flatten, get_dict_fingerprint_hash, unflatten
@@ -18,9 +19,30 @@ ONTOLOGY_CLASS_TO_FIELD = {
     NetworkConnection: "netflow"
 }
 
+if not HEUR_LIST:
+    # Get heuristics of service if not already set
+    HEUR_LIST = helper.get_heuristics()
+
 Classification = forge.get_classification()
 
 
+# Cleanup invalid tagging from service results
+def validate_tags(tag_map: Dict[str, List[str]]) -> Dict[str, List[str]]:
+    tag_map, _ = construct_safe(Tagging, unflatten(tag_map))
+    tag_map = flatten(tag_map.as_primitives(strip_null=True))
+    return tag_map
+
+# Merge tags
+def merge_tags(tag_a: Dict[str, List[str]], tag_b: Dict[str, List[str]]) -> Dict[str, List[str]]:
+    if not tag_a:
+        return tag_b
+
+    elif not tag_b:
+        return tag_a
+
+    all_keys = list(tag_a.keys()) + list(tag_b.keys())
+    return {key: list(set(tag_a.get(key, []) + tag_b.get(key, []))) for key in all_keys}
+
 class OntologyHelper:
     def __init__(self, logger, service_name) -> None:
         self.log = logger
@@ -84,32 +106,12 @@ class OntologyHelper:
 
         ontology['results'].update(self.results)
 
-    def _attach_ontology(self, request, working_dir) -> Optional[str]:
-        # Get heuristics of service
-        heuristics = helper.get_heuristics()
-
-        def preprocess_result_for_dump(sections, current_max, heur_tag_map, tag_map):
+    def _preprocess_result_for_dump(self, sections: List[ResultSection], current_max: str,
+                                    heur_tag_map: Dict[str, Dict[str, Any]], tag_map: Dict[str, List[str]], score: int) -> Tuple[str, Dict[str, Dict[str, Any]], Dict[str, List[str]], int]:
             for section in sections:
                 # Determine max classification of the overall result
                 current_max = Classification.max_classification(section.classification, current_max)
 
-                # Cleanup invalid tagging from service results
-                def validate_tags(tag_map):
-                    tag_map, _ = construct_safe(Tagging, unflatten(tag_map))
-                    tag_map = flatten(tag_map.as_primitives(strip_null=True))
-                    return tag_map
-
-                # Merge tags
-                def merge_tags(tag_a, tag_b):
-                    if not tag_a:
-                        return tag_b
-
-                    elif not tag_b:
-                        return tag_a
-
-                    all_keys = list(tag_a.keys()) + list(tag_b.keys())
-                    return {key: list(set(tag_a.get(key, []) + tag_b.get(key, []))) for key in all_keys}
-
                 # Append tags raised by the service, if any
                 section_tags = validate_tags(section.tags)
                 if section_tags:
@@ -117,8 +119,8 @@ class OntologyHelper:
 
                 # Append tags associated to heuristics raised by the service, if any
                 if section.heuristic:
-                    heur = heuristics[section.heuristic.heur_id]
-                    key = f'{request.task.service_name.upper()}_{heur.heur_id}'
+                    heur = HEUR_LIST[section.heuristic.heur_id]
+                    key = f'{self.service.upper()}_{heur.heur_id}'
                     heur_tag_map[key].update({
                         "heur_id": key,
                         "name": heur.name,
@@ -126,24 +128,27 @@ class OntologyHelper:
                         "score": heur.score,
                         "times_raised": heur_tag_map[key]["times_raised"] + 1
                     })
+                    score += section.heuristic.score
 
                 # Recurse through subsections
                 if section.subsections:
-                    current_max, heur_tag_map, tag_map = preprocess_result_for_dump(
-                        section.subsections, current_max, heur_tag_map, tag_map)
+                    current_max, heur_tag_map, tag_map = self._preprocess_result_for_dump(
+                        section.subsections, current_max, heur_tag_map, tag_map, score)
 
-            return current_max, heur_tag_map, tag_map
+            return current_max, heur_tag_map, tag_map, score
 
+    def _attach_ontology(self, request, working_dir) -> Optional[str]:
         if not request.result or not request.result.sections:
             # No service results, therefore no ontological output
             return
 
-        max_result_classification, heur_tag_map, tag_map = preprocess_result_for_dump(
+        max_result_classification, heur_tag_map, tag_map, score = self._preprocess_result_for_dump(
             sections=request.result.sections,
             current_max=Classification.max_classification(request.task.min_classification,
                                                           request.task.service_default_result_classification),
             heur_tag_map=defaultdict(lambda: {"tags": dict(), "times_raised": int()}),
-            tag_map=defaultdict(list))
+            tag_map=defaultdict(list),
+            score=0)
 
         if not heur_tag_map and not tag_map and not self._result_parts:
             # No heuristics, tagging, or ontologies found, therefore informational results
@@ -168,7 +173,8 @@ class OntologyHelper:
             },
             "results": {
                 "tags": tag_map,
-                "heuristics": list(heur_tag_map.values())
+                "heuristics": list(heur_tag_map.values()),
+                "score": score
             }
         }
 

{assemblyline_v4_service-4.5.1.dev232.dist-info → assemblyline_v4_service-4.5.1.dev233.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-v4-service
-Version: 4.5.1.dev232
+Version: 4.5.1.dev233
 Summary: Assemblyline 4 - Service base
 Home-page: https://github.com/CybercentreCanada/assemblyline-v4-service/
 Author: CCCS Assemblyline development team

{assemblyline_v4_service-4.5.1.dev232.dist-info → assemblyline_v4_service-4.5.1.dev233.dist-info}/RECORD

@@ -1,4 +1,4 @@
-assemblyline_v4_service/VERSION,sha256=SKe4bEiWAVHRKVwsXso5Rp0XOe-qhtnReZqhrTTGYE0,13
+assemblyline_v4_service/VERSION,sha256=Z6ozwNHb47WpLFd4QhbnMbJ7bL5wRoGPPh3JUf5voM4,13
 assemblyline_v4_service/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 assemblyline_v4_service/healthz.py,sha256=sS1cFkDLw8hUPMpj7tbHXFv8ZmHcazrwZ0l6oQDwwkQ,1575
 assemblyline_v4_service/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -9,7 +9,7 @@ assemblyline_v4_service/common/api.py,sha256=Xzp8j4HCCfjPvNSGKiZl5ttH2_Itg47cjlH
 assemblyline_v4_service/common/base.py,sha256=4rnK_183qsSfkRkuVp_0wy-haW49umx4jJIC0OalGDM,14218
 assemblyline_v4_service/common/helper.py,sha256=xs9quuf-M1JOdKieBqOmWaOece0CtzXFhhe85xQYmuY,3289
 assemblyline_v4_service/common/ocr.py,sha256=3fV0PyY3oui_ucAM9dkolP0VRYKACKJuGY4M64DudIE,8841
-assemblyline_v4_service/common/ontology_helper.py,sha256=4YW_iqwX3PFMc9efknxzBHNxep64tEaVASYS8WFZTHw,7924
+assemblyline_v4_service/common/ontology_helper.py,sha256=kiYI0ecQJc1zxEHePWTX73mAWTEdv4v3wpSk6ONQTsE,8270
 assemblyline_v4_service/common/request.py,sha256=W7fqC2xQE3i5i2jlCDyUDp3ZqJQQqSshNW0mQfJMkFg,11792
 assemblyline_v4_service/common/result.py,sha256=9AqM6qCYiia_Bpyn_fBFhzNQMcqJbtFSiGjp57fXW2E,32713
 assemblyline_v4_service/common/task.py,sha256=dJsvRpW0x88CCF_LW6w87jQ_UKTVaOs2Gb117IDNiU8,14233
@@ -33,13 +33,13 @@ test/test_common/test_api.py,sha256=7wlo7wgB12T23zMLbwjJ3GIomLHqE_Qvs3xkibSsR1U,
 test/test_common/test_base.py,sha256=fuJSSlPxIDHq6HU1xbvaMFitw2z1spOZNHD2SJ4UUic,13346
 test/test_common/test_helper.py,sha256=sO6YAiBhKTqaxlpLhFYDuy2ZdbuF2cg07Ylzo83ZzQs,2575
 test/test_common/test_ocr.py,sha256=mt_PgElgwQKJmNrp2nRVx9NjfMedVk40I6IV317vATI,1753
-test/test_common/test_ontology_helper.py,sha256=KhHEBg_ecJyQbDw79NMT4FzUyA4C1Aak3HEQCwBfM2s,7914
+test/test_common/test_ontology_helper.py,sha256=a_bX1MPQnwV1QILGgrUSNsgMgCnl663fAB40cw3ElHI,10237
 test/test_common/test_request.py,sha256=Ceyds8BNO1O0f1kH1VEb84faJcaupvSjVKIrGdHexsc,11842
 test/test_common/test_result.py,sha256=6BiOKxEPrKBjOY44jv3TY-yiXm0qI1ok_CZBnjP9TM4,45447
 test/test_common/test_task.py,sha256=P44mNcSe-3tJgDk9ppN3KbM7oN4LBVIuhONG-Gveh74,19007
 test/test_common/test_utils.py,sha256=TbnBxqpS_ZC5ptXR9XJX3xtbItD0mTbtiBxxdyP8J5k,5904
-assemblyline_v4_service-4.5.1.dev232.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
-assemblyline_v4_service-4.5.1.dev232.dist-info/METADATA,sha256=XPH7OTtdAKk7QHLCNCoirljClEHVbN5LztIQsTp7o5o,9493
-assemblyline_v4_service-4.5.1.dev232.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-assemblyline_v4_service-4.5.1.dev232.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
-assemblyline_v4_service-4.5.1.dev232.dist-info/RECORD,,
+assemblyline_v4_service-4.5.1.dev233.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
+assemblyline_v4_service-4.5.1.dev233.dist-info/METADATA,sha256=q1LCbQQGOjZUK2eAqnpTpZefohsJQB6M-ysdVYjXQw4,9493
+assemblyline_v4_service-4.5.1.dev233.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+assemblyline_v4_service-4.5.1.dev233.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
+assemblyline_v4_service-4.5.1.dev233.dist-info/RECORD,,

test/test_common/test_ontology_helper.py

@@ -7,7 +7,7 @@ from test.test_common import setup_module
 setup_module()
 
 import pytest
-from assemblyline_v4_service.common.ontology_helper import *
+from assemblyline_v4_service.common.ontology_helper import OntologyHelper, validate_tags, merge_tags
 from assemblyline_v4_service.common.result import ResultSection
 
 from assemblyline.odm.models.ontology.ontology import ODM_VERSION
@@ -224,15 +224,75 @@ def test_attach_ontology(dummy_request_class):
         'odm_type': 'Assemblyline Result Ontology',
         'odm_version': ODM_VERSION,
         'results': {'heuristics': [],
-                    'tags': {'network.static.domain': ['blah.com']}},
+                    'tags': {'network.static.domain': ['blah.com']},
+                    'score': 0},
         'service': {'name': 'blah',
                     'tool_version': 'blah',
-                    'version': '4.0.0'},
+                    'version': '4.0.0'}
     }
 
-    # TODO
-    # Test nested methods preprocess_result_for_dump, validate_tags and merge_tags
+    # Assign a heuristic to the section and check to see if the score gets updated correctly
+    req.result.sections[0].set_heuristic(1)
+    req.result.sections.append(req.result.sections[0])
+    req.task.supplementary = []
+    assert oh._attach_ontology(req, working_dir) is None
+    assert oh.results == {}
+    assert req.task.supplementary == [
+        {
+            'classification': 'TLP:C',
+            'description': 'Result Ontology from blah',
+            'name': 'blah_blah.ontology',
+            'path': os.path.join(working_dir, "blah.ontology")
+        }
+    ]
 
+    with open(os.path.join(working_dir, "blah.ontology"), "r") as f:
+        file_contents = json.loads(f.read())
+    # The sum of the overall result should be 500 points since the heuristic was raised twice with a score of 250 per section
+    assert file_contents == {
+        'classification': 'TLP:C',
+        'file': {'md5': 'blah',
+                'names': ['blah'],
+                'sha1': 'blah',
+                'sha256': 'blah',
+                'size': 123,
+                'type': None},
+        'odm_type': 'Assemblyline Result Ontology',
+        'odm_version': ODM_VERSION,
+        'results': {'heuristics': [{'tags': {'network.static.domain': ['blah.com']}, 'times_raised': 2, 'heur_id': 'BLAH_1', 'name': 'blah', 'score': 250}],
+                    'tags': {'network.static.domain': ['blah.com']},
+                    'score': 500},
+        'service': {'name': 'blah',
+                    'tool_version': 'blah',
+                    'version': '4.0.0'}
+    }
+
+
+def test_validate_tags():
+    # Valid tag
+    assert validate_tags({'network.static.domain':['blah.com']})
+
+    # Invalid tag
+    assert not validate_tags({'blah':['blah.com']})
+
+def test_merge_tags():
+    # Merge with null should yield the tag map with value
+    assert merge_tags({'abc':['xyz']}, None) == {'abc':['xyz']}
+    assert merge_tags(None, {'abc':['xyz']}) == {'abc':['xyz']}
+
+    # Merge with mutually exclusive keys
+    assert merge_tags({'abc':['xyz']}, {'def':['xyz']}) == {
+        'abc':['xyz'],
+        'def':['xyz']
+    }
+
+    # Merge with non-mutually exclusive keys
+    merged_tags = merge_tags({'abc':['xyz'], 'def':['tuv']}, {'def':['xyz']})
+    merged_tags['def'] = sorted(merged_tags['def'])
+    merged_tags == {
+        'abc':['xyz'],
+        'def':['tuv', 'xyz']
+    }
 
 def test_reset():
     oh = OntologyHelper(None, "blah")