assemblyline-v4-service 4.5.0.11__py3-none-any.whl → 4.5.0.13__py3-none-any.whl

assemblyline_v4_service/VERSION

@@ -1 +1 @@
-4.5.0.11
+4.5.0.13

assemblyline_v4_service/common/base.py

@@ -19,6 +19,7 @@ from assemblyline.odm.messages.task import Task as ServiceTask
 from assemblyline_v4_service.common import helper
 from assemblyline_v4_service.common.api import PrivilegedServiceAPI, ServiceAPI
 from assemblyline_v4_service.common.ontology_helper import OntologyHelper
+from assemblyline_v4_service.common.ocr import update_ocr_config
 from assemblyline_v4_service.common.request import ServiceRequest
 from assemblyline_v4_service.common.task import Task
 
@@ -84,6 +85,10 @@ class ServiceBase:
         self.rules_hash: str = None
         self.signatures_meta: dict = {}
 
+        # OCR-related
+        if self.config.get('ocr'):
+            update_ocr_config(self.config['ocr'])
+
     @property
     def api_interface(self):
         return self.get_api_interface()

assemblyline_v4_service/common/ocr.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from typing import Dict, List, TextIO
+from typing import Any, Dict, List, TextIO, Union
 
 from assemblyline_v4_service.common.helper import get_service_manifest
 from assemblyline_v4_service.common.utils import PASSWORD_WORDS
 
-# TODO: Would prefer this mapping to be dynamic from trusted sources (ie. import from library), but will copy-paste for now
+# The terms related to each indicator category
 OCR_INDICATORS_TERMS: dict[str, list[str]] = {
     "ransomware": [
         # https://github.com/cuckoosandbox/community/blob/master/modules/signatures/windows/ransomware_message.py
@@ -154,6 +154,48 @@ OCR_INDICATORS_TERMS: dict[str, list[str]] = {
 # The minimum number of indicator hits to avoid FP detections
 OCR_INDICATORS_THRESHOLD: Dict[str, int] = {"ransomware": 2, "macros": 2, "banned": 1, "password": 1}
 
+def update_ocr_config(ocr_config: Dict[str, Union[List[str], Dict[str, Any]]] = None):
+    global OCR_INDICATORS_TERMS
+    global OCR_INDICATORS_THRESHOLD
+    if not ocr_config:
+        try:
+            # Retrieve service-configured OCR settings on module load (primary used in testing)
+            ocr_config: Dict = get_service_manifest().get("config", {}).get("ocr", {})
+        except Exception:
+            # No configuration updates provided
+            return
+
+    indicators = set(list(OCR_INDICATORS_TERMS.keys()) + list(ocr_config.keys()))
+    # Iterate over the different indicators and include lines of detection in response
+    for indicator in indicators:
+        indicator_config = ocr_config.get(indicator)
+        terms = OCR_INDICATORS_TERMS.get(indicator, [])
+        hit_threshold = OCR_INDICATORS_THRESHOLD.get(indicator, 1)
+        # Backwards compatibility: Check how the OCR configuration is formatted
+        if not indicator_config:
+            # Empty block/no override provided by service
+            pass
+        elif isinstance(indicator_config, list):
+            # Legacy support (before configurable indicator thresholds)
+            terms = indicator_config
+        elif isinstance(indicator_config, dict):
+            # Either you're exclusively overwriting the terms list or you're selectively including/excluding terms
+            if indicator_config.get("terms"):
+                # Overwrite terms list with service configuration
+                terms = indicator_config["terms"]
+            else:
+                included_terms = set(indicator_config.get("include", []))
+                excluded_terms = set(indicator_config.get("exclude", []))
+                # Compute the new terms list for indicator type
+                terms = list(set(terms).union(included_terms) - excluded_terms)
+
+            # Set the indicator hit threshold
+            hit_threshold = indicator_config.get("threshold", 1)
+
+        # Overwrite key-value in respective constants
+        OCR_INDICATORS_TERMS[indicator] = terms
+        OCR_INDICATORS_THRESHOLD[indicator] = hit_threshold
+
 
 def ocr_detections(image_path: str, ocr_io: TextIO = None) -> Dict[str, List[str]]:
     try:
@@ -185,57 +227,25 @@ def ocr_detections(image_path: str, ocr_io: TextIO = None) -> Dict[str, List[str
 
 
 def detections(ocr_output: str) -> Dict[str, List[str]]:
-    try:
-        # Retrieve service-configured OCR settings on module load
-        ocr_config: Dict = get_service_manifest().get("config", {}).get("ocr", {})
-    except Exception:
-        # Service manifest not found
-        ocr_config = {}
-
-    indicators = set(list(OCR_INDICATORS_TERMS.keys()) + list(ocr_config.keys()))
     detection_output: Dict[str, List[str]] = {}
-    # Iterate over the different indicators and include lines of detection in response
-    for indicator in indicators:
-        indicator_config = ocr_config.get(indicator)
-        terms = OCR_INDICATORS_TERMS.get(indicator, [])
-        hit_threshold = OCR_INDICATORS_THRESHOLD.get(indicator, 1)
-        # Backwards compatibility: Check how the OCR configuration is formatted
-        if not indicator_config:
-            # Empty block/no override provided by service
-            pass
-        elif isinstance(indicator_config, list):
-            # Legacy support (before configurable indicator thresholds)
-            terms = indicator_config
-        elif isinstance(indicator_config, dict):
-            # Either you're exclusively overwriting the terms list or you're selectively including/excluding terms
-            if indicator_config.get("terms"):
-                # Overwrite terms list with service configuration
-                terms = indicator_config["terms"]
-            else:
-                included_terms = set(indicator_config.get("include", []))
-                excluded_terms = set(indicator_config.get("exclude", []))
-                # Compute the new terms list for indicator type
-                terms = list(set(terms).union(included_terms) - excluded_terms)
-
-            # Set the indicator hit threshold
-            hit_threshold = indicator_config.get("threshold", 1)
-
+    for indicator, terms in OCR_INDICATORS_TERMS.items():
+        hit_threshold = OCR_INDICATORS_THRESHOLD[indicator]
         # Perform a pre-check to see if the terms even exist in the OCR text
         if not any([t.lower() in ocr_output.lower() for t in terms]):
             continue
 
         # Keep a track of the hits and the lines corresponding with term hits
-        indicator_hits: int = 0
+        indicator_hits: set = set()
         list_of_strings: List[str] = []
         for line in ocr_output.split("\n"):
             for t in terms:
                 term_count = line.lower().count(t.lower())
                 if term_count:
-                    indicator_hits += term_count
+                    indicator_hits.add(t)
                     if line not in list_of_strings:
                         list_of_strings.append(line)
 
-        if list_of_strings and indicator_hits >= hit_threshold:
+        if list_of_strings and len(indicator_hits) >= hit_threshold:
             # If we were to find hits and those hits are above the required threshold, then add them to output
             detection_output[indicator] = list_of_strings
     return detection_output

{assemblyline_v4_service-4.5.0.11.dist-info → assemblyline_v4_service-4.5.0.13.dist-info}/METADATA

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

{assemblyline_v4_service-4.5.0.11.dist-info → assemblyline_v4_service-4.5.0.13.dist-info}/RECORD

@@ -1,4 +1,4 @@
-assemblyline_v4_service/VERSION,sha256=1_9OL3qv21zjqWhS4QZ0yZBiBfz5-VmY3lDxRoxFMV4,9
+assemblyline_v4_service/VERSION,sha256=VuBzD9mQzFaZYAMCWM9b0vxAyjHIbKMlmPzW3unLZFw,9
 assemblyline_v4_service/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 assemblyline_v4_service/healthz.py,sha256=3QGBg0EZuXC6UN411HFwpLNEop9UvS9feFhvBUTP-k4,1576
 assemblyline_v4_service/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -6,9 +6,9 @@ assemblyline_v4_service/run_privileged_service.py,sha256=qd4DmHo5G_Tpv8tb0A96qNf
 assemblyline_v4_service/run_service.py,sha256=NiFX52NfsbBQY6E3nrjoB3e2XxIlwFcQpYYY-rADIk4,5996
 assemblyline_v4_service/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 assemblyline_v4_service/common/api.py,sha256=Xzp8j4HCCfjPvNSGKiZl5ttH2_Itg47cjlH0NXNtth0,6849
-assemblyline_v4_service/common/base.py,sha256=mKkkzbxVL_wVMy_VieU9mlHYLqZXndga_4dWWbnEnx8,14045
+assemblyline_v4_service/common/base.py,sha256=1wxzg_enESdqxloiAWCL0fUNzcKg0kAVfNNkGAfHX0U,14218
 assemblyline_v4_service/common/helper.py,sha256=xs9quuf-M1JOdKieBqOmWaOece0CtzXFhhe85xQYmuY,3289
-assemblyline_v4_service/common/ocr.py,sha256=A8OnjpEor-S3OUC_jZzJ-Er3KKAsMdQLEXTtMS81Xbk,8397
+assemblyline_v4_service/common/ocr.py,sha256=3fV0PyY3oui_ucAM9dkolP0VRYKACKJuGY4M64DudIE,8841
 assemblyline_v4_service/common/ontology_helper.py,sha256=QpwerYoS5hXjWzpx3Pmwv6j2330PQVYqxYGamjcpW3I,7890
 assemblyline_v4_service/common/request.py,sha256=NxtWxp8-ttC72i3Vnchc3fZTRKqPQAoMC4KAVwEj8-4,11714
 assemblyline_v4_service/common/result.py,sha256=9AqM6qCYiia_Bpyn_fBFhzNQMcqJbtFSiGjp57fXW2E,32713
@@ -28,18 +28,18 @@ test/conftest.py,sha256=W3SieQpZsZpGEmtLqY4aIlxREDSsHceyCrFcFsWUM0U,1851
 test/test_healthz.py,sha256=DkeLUlrb7rGx3nZ04aADU9HXXu5mZTf_DBwT0xhzIv4,7
 test/test_run_privileged_service.py,sha256=DkeLUlrb7rGx3nZ04aADU9HXXu5mZTf_DBwT0xhzIv4,7
 test/test_run_service.py,sha256=DkeLUlrb7rGx3nZ04aADU9HXXu5mZTf_DBwT0xhzIv4,7
-test/test_common/__init__.py,sha256=v3__rzWUBW0Smc2rFwfFR9WehQHM_uBTIFCdC_We3tA,1319
+test/test_common/__init__.py,sha256=RkOm3vnVp5L947mD1jTo4bdOgLTZJ24_NX-kqfMn5a8,1259
 test/test_common/test_api.py,sha256=7wlo7wgB12T23zMLbwjJ3GIomLHqE_Qvs3xkibSsR1U,4902
-test/test_common/test_base.py,sha256=d61an3lRaes_cx0AOPMNFMVWxTOjilrcGpuVP0dqTvM,13198
+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=s5kL0vKjLmbyXuCg-9v6V6wQwwFRu0w2hYpjG0_BXy4,1874
-test/test_common/test_ontology_helper.py,sha256=TuvTeP9BTRqklOlsLu_yMdN9wdPWlVAENx2JqUe9a-A,7856
-test/test_common/test_request.py,sha256=K-gaHe20AUztfU3pLdWYv5bc6088GBBKP-eb833cis0,11729
-test/test_common/test_result.py,sha256=Wm0Cs5kZRzlZr0jL-l8OTsYAvkoN2eaB3NkeXzvyssI,42208
+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_request.py,sha256=VzWw-NZKTRIXKwCACIT2SSHE70HMZ3IeBk1AJiT-cjU,11472
+test/test_common/test_result.py,sha256=b96bCfyW0ukdTcCsl01jS_l5YhfzXFVYs_VPOwz7IEU,41982
 test/test_common/test_task.py,sha256=RWsGEN0L-xKoRAsGv2x4JZcSAySS9aBF4dwL9t7tepo,18668
 test/test_common/test_utils.py,sha256=TbnBxqpS_ZC5ptXR9XJX3xtbItD0mTbtiBxxdyP8J5k,5904
-assemblyline_v4_service-4.5.0.11.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
-assemblyline_v4_service-4.5.0.11.dist-info/METADATA,sha256=DLdPO96aAuV6e-qQibM_zwEqYA5WY5mbUqOJfaj6Tz8,9495
-assemblyline_v4_service-4.5.0.11.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-assemblyline_v4_service-4.5.0.11.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
-assemblyline_v4_service-4.5.0.11.dist-info/RECORD,,
+assemblyline_v4_service-4.5.0.13.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
+assemblyline_v4_service-4.5.0.13.dist-info/METADATA,sha256=ww3G0mD7HZ9MsPJxRQGyMQm6-AaGDOfNVZrutQ63-fg,9495
+assemblyline_v4_service-4.5.0.13.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+assemblyline_v4_service-4.5.0.13.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
+assemblyline_v4_service-4.5.0.13.dist-info/RECORD,,

test/test_common/__init__.py

@@ -3,8 +3,7 @@ import subprocess
 
 from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
 
-SERVICE_CONFIG_NAME = "service_manifest.yml"
-TEMP_SERVICE_CONFIG_PATH = os.path.join("/tmp", SERVICE_CONFIG_NAME)
+TEMP_SERVICE_CONFIG_PATH ="/tmp/service_manifest.yml"
 
 ret = subprocess.run("dpkg -l | grep ^ii | awk '{print $2}' | grep -i 'tesseract'", capture_output=True, shell=True)
 TESSERACT_LIST = list(filter(None, ret.stdout.decode().split('\n')))

test/test_common/test_base.py

@@ -3,6 +3,12 @@ import os
 import time
 from logging import Logger
 
+
+from test.test_common import setup_module
+
+# Ensure service manifest is instantiated before loading assemblyline_v4_service module
+setup_module()
+
 import pytest
 import requests_mock
 from assemblyline_v4_service.common.base import *

test/test_common/test_ocr.py

@@ -1,16 +1,14 @@
 import os
-from test.test_common import TESSERACT_LIST, setup_module
+from test.test_common import TESSERACT_LIST
 
 import pytest
-from assemblyline_v4_service.common.ocr import *
 
+from assemblyline_v4_service.common.ocr import ocr_detections, detections, update_ocr_config
 
 @pytest.mark.skipif(len(TESSERACT_LIST) < 1, reason="Requires tesseract-ocr apt package")
 def test_ocr_detections():
-    if os.getcwd().endswith("/test"):
-        file_path = os.path.join(os.getcwd(), "test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
-    else:
-        file_path = os.path.join(os.getcwd(), "test/test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
+    update_ocr_config()
+    file_path = os.path.join(os.path.dirname(__file__), "b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
     assert ocr_detections(file_path) == {
         'ransomware': [
             "YOUR FILES HAVE BEEN ENCRYPTED AND YOU WON'T BE ABLE TO "
@@ -30,7 +28,8 @@ def test_ocr_detections():
 
 
 def test_detections():
-    setup_module()
+    update_ocr_config()
+
     # No detection
     assert detections("blah") == {}
 

test/test_common/test_ontology_helper.py

@@ -3,6 +3,9 @@ import logging
 import os
 import tempfile
 
+from test.test_common import setup_module
+setup_module()
+
 import pytest
 from assemblyline_v4_service.common.ontology_helper import *
 from assemblyline_v4_service.common.result import ResultSection

test/test_common/test_request.py

@@ -1,7 +1,7 @@
 import os
 import tempfile
 from logging import Logger
-from test.test_common import TESSERACT_LIST
+from test.test_common import TESSERACT_LIST, setup_module
 
 import pytest
 from assemblyline_v4_service.common.request import ServiceRequest
@@ -10,6 +10,8 @@ from assemblyline_v4_service.common.task import MaxExtractedExceeded, Task
 
 from assemblyline.odm.messages.task import Task as ServiceTask
 
+# Ensure service manifest is instantiated before importing from OCR submodule
+setup_module()
 
 @pytest.fixture
 def service_request():
@@ -105,10 +107,7 @@ def test_add_extracted(service_request):
 
 @pytest.mark.skipif(len(TESSERACT_LIST) < 1, reason="Requires tesseract-ocr apt package")
 def test_add_image(service_request):
-    if os.getcwd().endswith("/test"):
-        image_path = os.path.join(os.getcwd(), "test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
-    else:
-        image_path = os.path.join(os.getcwd(), "test/test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
+    image_path = os.path.join(os.path.dirname(__file__), "b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
 
     # Basic
     assert service_request.add_image(image_path, "image_name", "description of image") == {
@@ -196,10 +195,7 @@ def test_add_image(service_request):
     service_request.task.supplementary.clear()
 
     # Classification, OCR heuristic, OCR_IO, image with password
-    if os.getcwd().endswith("/test"):
-        image_path = os.path.join(os.getcwd(), "test_common/4031ed8786439eee24b87f84901e38038a76b8c55e9d87dd5a7d88df2806c1cf")
-    else:
-        image_path = os.path.join(os.getcwd(), "test/test_common/4031ed8786439eee24b87f84901e38038a76b8c55e9d87dd5a7d88df2806c1cf")
+    image_path = os.path.join(os.path.dirname(__file__), "4031ed8786439eee24b87f84901e38038a76b8c55e9d87dd5a7d88df2806c1cf")
     _, path = tempfile.mkstemp()
     ocr_io = open(path, "w")
     data = service_request.add_image(image_path, "image_name", "description of image", "TLP:A", ocr_heuristic_id, ocr_io)

test/test_common/test_result.py

@@ -1,6 +1,10 @@
 import os
 import tempfile
-from test.test_common import TESSERACT_LIST
+from test.test_common import TESSERACT_LIST, setup_module
+
+# Ensure service manifest is instantiated before importing from OCR submodule
+setup_module()
+
 
 import pytest
 from assemblyline_v4_service.common.request import ServiceRequest
@@ -9,7 +13,6 @@ from assemblyline_v4_service.common.task import Task
 
 from assemblyline.odm.messages.task import Task as ServiceTask
 
-
 @pytest.fixture
 def heuristic():
     return Heuristic(1)
@@ -592,10 +595,7 @@ def test_imagesectionbody_init(service_request):
 @pytest.mark.skipif(len(TESSERACT_LIST) < 1, reason="Requires tesseract-ocr apt package")
 def test_imagesectionbody_add_image(service_request):
     isb = ImageSectionBody(service_request)
-    if os.getcwd().endswith("/test"):
-        image_path = os.path.join(os.getcwd(), "test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
-    else:
-        image_path = os.path.join(os.getcwd(), "test/test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
+    image_path = os.path.join(os.path.dirname(__file__), "b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
 
     # Basic
     assert isb.add_image(image_path, "image_name", "description of image") is None
@@ -1230,10 +1230,8 @@ def test_resultimagesection_init(service_request):
 def test_resultimagesection_add_image(service_request):
     ris = ResultImageSection(service_request, "title_text_as_str")
 
-    if os.getcwd().endswith("/test"):
-        image_path = os.path.join(os.getcwd(), "test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
-    else:
-        image_path = os.path.join(os.getcwd(), "test/test_common/b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
+    image_path = os.path.join(os.path.dirname(__file__),
+                              "b32969aa664e3905c20f865cdd7b921f922678f5c3850c78e4c803fbc1757a8e")
 
     # Basic
     assert ris.add_image(image_path, "image_name", "description of image") is None