assemblyline-v4-service 4.5.1.dev71__py3-none-any.whl → 4.5.1.dev75__py3-none-any.whl

assemblyline_v4_service/VERSION

@@ -1 +1 @@
-4.5.1.dev71
+4.5.1.dev75

assemblyline_v4_service/common/ocr.py

@@ -5,7 +5,7 @@ from typing import Dict, List, TextIO
 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,45 @@ 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}
 
+# Pre-compute the OCR_INDICATOR_* constants on module load so we only load the manifest once rather than per OCR request
+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()))
+# 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 +224,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.1.dev71.dist-info → assemblyline_v4_service-4.5.1.dev75.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-v4-service
-Version: 4.5.1.dev71
+Version: 4.5.1.dev75
 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.dev71.dist-info → assemblyline_v4_service-4.5.1.dev75.dist-info}/RECORD

@@ -1,4 +1,4 @@
-assemblyline_v4_service/VERSION,sha256=sgnbYqhBwfHaBVEZ9RudMJZ-VblAxWq7JSyPx5vGabE,12
+assemblyline_v4_service/VERSION,sha256=JZHoM4NWVTRtlPRRggVpGFHenhb-c_AXFCT7sYtWPQQ,12
 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
@@ -8,7 +8,7 @@ assemblyline_v4_service/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5
 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/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=dzx5n9fzOHCVX3NnGKONpGAF3Cq16hw_y7M7g_nsJ7A,8588
 assemblyline_v4_service/common/ontology_helper.py,sha256=QpwerYoS5hXjWzpx3Pmwv6j2330PQVYqxYGamjcpW3I,7890
 assemblyline_v4_service/common/request.py,sha256=XXBafAQCV43_OBLXOSHxYoDHmqwERBkNul8fb_X6Ves,11774
 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=wxSwnOj-_YOv2SuZjOJsw09q8A7p8GJmJuK4vozqCNg,11749
-test/test_common/test_result.py,sha256=Wm0Cs5kZRzlZr0jL-l8OTsYAvkoN2eaB3NkeXzvyssI,42208
+test/test_common/test_ocr.py,sha256=nel1GCkieDRW2F_6kYCbkIwB9Kwj_d2rJBgb8VZWXS8,1685
+test/test_common/test_ontology_helper.py,sha256=KhHEBg_ecJyQbDw79NMT4FzUyA4C1Aak3HEQCwBfM2s,7914
+test/test_common/test_request.py,sha256=PPhHfrwpwMdNZ33P1Z_0h1Zaz9ao9VFiDr_MJrBS3Lg,11492
+test/test_common/test_result.py,sha256=b96bCfyW0ukdTcCsl01jS_l5YhfzXFVYs_VPOwz7IEU,41982
 test/test_common/test_task.py,sha256=jnfF68EgJIu30Pz_4jiJHkncfI-3XpGaut5r79KIXOA,18718
 test/test_common/test_utils.py,sha256=TbnBxqpS_ZC5ptXR9XJX3xtbItD0mTbtiBxxdyP8J5k,5904
-assemblyline_v4_service-4.5.1.dev71.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
-assemblyline_v4_service-4.5.1.dev71.dist-info/METADATA,sha256=TCHvOZ28H5YkRA2nGTlA3jv9y8r8Ty2w6HPQ6o5vGSc,9498
-assemblyline_v4_service-4.5.1.dev71.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-assemblyline_v4_service-4.5.1.dev71.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
-assemblyline_v4_service-4.5.1.dev71.dist-info/RECORD,,
+assemblyline_v4_service-4.5.1.dev75.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
+assemblyline_v4_service-4.5.1.dev75.dist-info/METADATA,sha256=G0uwI6VSb0U0defzTlpY1uKuw0qji1s4OGUTxvPuZig,9498
+assemblyline_v4_service-4.5.1.dev75.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+assemblyline_v4_service-4.5.1.dev75.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
+assemblyline_v4_service-4.5.1.dev75.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,13 @@
 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
 
 @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")
+    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 +27,6 @@ def test_ocr_detections():
 
 
 def test_detections():
-    setup_module()
     # 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