assemblyline-v4-service 4.5.1.dev453__py3-none-any.whl → 4.5.1.dev455__py3-none-any.whl

assemblyline_v4_service/VERSION

@@ -1 +1 @@
-4.5.1.dev453
+4.5.1.dev455

assemblyline_v4_service/dev/run_service_once.py

@@ -1,15 +1,15 @@
 import argparse
+import cProfile
+import importlib
 import json
 import logging
 import os
 import pprint
 import shutil
 import tempfile
-import yaml
-import cProfile
+from typing import Dict, Union
 
 from cart import get_metadata_only, unpack_stream
-from typing import Union, Dict
 
 from assemblyline.common import forge
 from assemblyline.common.heuristics import HeuristicHandler, InvalidHeuristicException
@@ -19,12 +19,14 @@ from assemblyline.common.uid import get_random_id
 from assemblyline.odm.messages.task import Task as ServiceTask
 from assemblyline.odm.models.result import Result
 from assemblyline.odm.models.service import Service
+from assemblyline_v4_service.common.base import ServiceBase
 from assemblyline_v4_service.common.helper import get_heuristics, get_service_manifest
+from assemblyline_v4_service.dev.updater import load_rules
 
 
 class RunService:
     def __init__(self):
-        self.service = None
+        self.service: ServiceBase = None
         self.service_class = None
         self.submission_params = None
         self.file_dir = None
@@ -41,6 +43,10 @@ class RunService:
 
         self.load_service_manifest()
 
+        if self.service.service_attributes.update_config:
+            # Download required signatures and process them for the service run
+            load_rules(self.service)
+
         if not os.path.isfile(FILE_PATH):
             LOG.info(f"File not found: {FILE_PATH}")
             return
@@ -179,6 +185,10 @@ class RunService:
         LOG.info(f"Cleaning up file used for temporary processing: {target_file}")
         os.unlink(target_file)
 
+        if self.service.rules_directory:
+            LOG.info("Cleaning up downloaded signatures..")
+            shutil.rmtree(self.service.rules_directory)
+
         LOG.info(f"Moving {result_json} to the working directory: {working_dir}/result.json")
         shutil.move(result_json, os.path.join(working_dir, 'result.json'))
 

assemblyline_v4_service/dev/updater.py

@@ -0,0 +1,150 @@
+import importlib
+import inspect
+import json
+import os
+import tempfile
+import threading
+
+from assemblyline.common.isotime import now_as_iso
+from assemblyline.odm.models.service import SIGNATURE_DELIMITERS
+from assemblyline_v4_service.common.base import ServiceBase
+from assemblyline_v4_service.updater.client import (
+    BadlistClient,
+    SafelistClient,
+    SignatureClient,
+    UpdaterClient,
+)
+from assemblyline_v4_service.updater.updater import (
+    SIGNATURES_META_FILENAME,
+    SOURCE_STATUS_KEY,
+    ServiceUpdater,
+    UniqueQueue,
+)
+
+
+class TestSignatureClient(SignatureClient):
+    def __init__(self, output_directory: str):
+        self.sync = False
+        self.output_directory = output_directory
+
+    def add_update_many(self, source, sig_type, data, dedup_name=True):
+        os.makedirs(os.path.join(self.output_directory, sig_type, source), exist_ok=True)
+        for d in data:
+            with open(os.path.join(self.output_directory, sig_type, source, d['name']), 'w') as f:
+                json.dump(d, f)
+
+        return {'success': len(data)}
+
+class TestBadlistClient(BadlistClient):
+    def __init__(self, output_directory: str):
+        self.sync = False
+        self.output_directory = output_directory
+
+    def add_update_many(self, list_of_badlist_objects):
+        return {'success': len(list_of_badlist_objects)}
+
+class TestSafelistClient(SafelistClient):
+    def __init__(self, output_directory: str):
+        self.sync = False
+        self.output_directory = output_directory
+
+    def add_update_many(self, list_of_safelist_objects):
+        return {'success': len(list_of_safelist_objects)}
+
+class TestUpdaterClient(UpdaterClient):
+    def __init__(self, output_directory: str):
+        self._sync = False
+        self._classification_override = False
+        self.signature = TestSignatureClient(output_directory)
+        self.badlist = TestBadlistClient(output_directory)
+        self.safelist = TestSafelistClient(output_directory)
+
+def load_rules(service: ServiceBase):
+    with tempfile.TemporaryDirectory() as latest_updates_dir:
+        updater_module = importlib.import_module(service.service_attributes.dependencies['updates'].container.command[-1])
+        # Find the UpdaterServer class
+        for v in updater_module.__dict__.values():
+            if inspect.isclass(v) and issubclass(v, ServiceUpdater) and v != ServiceUpdater:
+                updater_class = v
+                break
+
+
+        # Implement a class to be used with RunServiceOnce without a dependency on Assemblyline
+        class TestServiceUpdater(updater_class):
+            def __init__(self, *args, **kwargs):
+                self.update_data_hash = {}
+                self._current_source = ""
+                self.log = service.log
+                self._service = service.service_attributes
+                self.update_queue = UniqueQueue()
+                self.updater_type = self._service.name.lower()
+                self.delimiter = self._service.update_config.signature_delimiter
+                self.default_pattern = self._service.update_config.default_pattern
+                self.signatures_meta = {}
+                [self.update_queue.put(update.name) for update in self._service.update_config.sources]
+
+                self.latest_updates_dir = latest_updates_dir
+                self.client = TestUpdaterClient(latest_updates_dir)
+                self.source_update_flag = threading.Event()
+                self.local_update_flag = threading.Event()
+                self.local_update_start = threading.Event()
+
+            def set_source_update_time(self, update_time: float): ...
+
+            def set_source_extra(self, extra_data): ...
+
+            def set_active_config_hash(self, config_hash: int): ...
+
+            # Keep a record of the source status as a dictionary
+            def push_status(self, state: str, message: str):
+                # Push current state of updater with source
+                self.log.debug(f"Pushing state for {self._current_source}: [{state}] {message}")
+                self.update_data_hash[f'{self._current_source}.{SOURCE_STATUS_KEY}'] = \
+                    dict(state=state, message=message, ts=now_as_iso())
+
+            def do_source_update(self):
+                super().do_source_update(self._service)
+
+            def do_local_update(self):
+                if self._service.update_config.generates_signatures:
+                    signaure_data = []
+                    updatepath = os.path.join(self.latest_updates_dir, self.updater_type)
+                    for source in os.listdir(updatepath):
+                        sourcepath = os.path.join(updatepath, source)
+
+                        for file in os.listdir(sourcepath):
+                            # Save signatures to disk
+                            filepath = os.path.join(sourcepath, file)
+                            with open(filepath) as f:
+                                data = json.load(f)
+
+                                signaure_data.append(data.pop('data'))
+                                self.signatures_meta[data['signature_id']] = data
+
+                            if self.delimiter != "file":
+                                os.remove(filepath)
+
+                        if self.delimiter != "file":
+                            # Render the response when calling `client.signature.download`
+                            os.removedirs(sourcepath)
+                            with open(os.path.join(self.latest_updates_dir, source), 'w') as f:
+                                f.write(SIGNATURE_DELIMITERS[self.delimiter].join(signaure_data))
+
+                else:
+                    self.signatures_meta = {
+                        source.name: {'classification': source['default_classification'].value}
+                        for source in self._service.update_config.sources
+                    }
+
+
+
+        # Initialize updater, download signatures, and load them into the service
+        updater = TestServiceUpdater()
+        updater.do_source_update()
+        updater.do_local_update()
+        rules_directory = updater.prepare_output_directory()
+        service.signatures_meta = updater.signatures_meta
+        service.rules_directory = rules_directory
+        service.rules_list = [os.path.join(rules_directory, i) for i  in os.listdir(rules_directory)
+                            if i != SIGNATURES_META_FILENAME]
+        service._load_rules()

assemblyline_v4_service/updater/updater.py

@@ -1,37 +1,44 @@
 from __future__ import annotations
-import shutil
-from typing import Optional, Any, Tuple, List
-import typing
-import os
-import logging
-import time
+
+import hashlib
 import json
-import tempfile
+import logging
+import os
+import shutil
+import subprocess
 import tarfile
+import tempfile
 import threading
-import subprocess
-import hashlib
+import time
+import typing
 from io import BytesIO
 from queue import Queue
+from typing import Any, List, Optional, Tuple
 from zipfile import ZipFile
 
-from assemblyline.common import forge, log as al_log
+from assemblyline_core.server_base import ServiceStage, ThreadedCoreBase
+
+from assemblyline.common import forge
+from assemblyline.common import log as al_log
 from assemblyline.common.isotime import epoch_to_iso, now_as_iso
 from assemblyline.odm.messages.changes import Operation, ServiceChange, SignatureChange
-from assemblyline.remote.datatypes.events import EventSender, EventWatcher
-
-from assemblyline_core.server_base import ThreadedCoreBase, ServiceStage
 from assemblyline.odm.models.service import Service, UpdateSource
+from assemblyline.remote.datatypes.events import EventSender, EventWatcher
 from assemblyline.remote.datatypes.hash import Hash
-
 from assemblyline_v4_service.common.base import SIGNATURES_META_FILENAME
 from assemblyline_v4_service.updater.client import UpdaterClient
-from assemblyline_v4_service.updater.helper import url_download, git_clone_repo, SkipSource, filter_downloads
+from assemblyline_v4_service.updater.helper import (
+    SkipSource,
+    filter_downloads,
+    git_clone_repo,
+    url_download,
+)
 
 if typing.TYPE_CHECKING:
     import redis
-    from assemblyline.odm.models.config import Config
+
     from assemblyline.datastore.helper import AssemblylineDatastore
+    from assemblyline.odm.models.config import Config
     RedisType = redis.Redis[typing.Any]
 
 SERVICE_PULL_INTERVAL = 1200
@@ -65,10 +72,10 @@ class ServiceUpdater(ThreadedCoreBase):
                  shutdown_timeout: float = None, config: Config = None,
                  datastore: AssemblylineDatastore = None,
                  redis: RedisType = None, redis_persist: RedisType = None,
-                 default_pattern=".*", downloadable_signature_statuses=['DEPLOYED', 'NOISY']):
+                 downloadable_signature_statuses=['DEPLOYED', 'NOISY']):
 
         self.updater_type = os.environ['SERVICE_PATH'].split('.')[-1].lower()
-        self.default_pattern = default_pattern
+        self.default_pattern = None
 
         if not logger:
             al_log.init_logging(f'updater.{self.updater_type}', log_level=os.environ.get('LOG_LEVEL', "WARNING"))
@@ -247,6 +254,10 @@ class ServiceUpdater(ThreadedCoreBase):
         # Download the service object from datastore
         self._service = self.datastore.get_service_with_delta(SERVICE_NAME)
 
+        # Set default pattern if not already set
+        if not self.default_pattern:
+            self.default_pattern = self._service.update_config.default_pattern
+
         # Update signature client with any changes to classification rewrites
         self.client.signature.classification_replace_map = \
             self._service.config.get('updater', {}).get('classification_replace', {})

{assemblyline_v4_service-4.5.1.dev453.dist-info → assemblyline_v4_service-4.5.1.dev455.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-assemblyline_v4_service/VERSION,sha256=ebZs1rTqnub4bIpqlps2JL3Zu7V-GpViEdOtKpjFTdM,13
+assemblyline_v4_service/VERSION,sha256=wHs-DuDtnVx9REMerMGYAVk7KidZ-A0fOnJzsGVi8YI,13
 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
@@ -15,14 +15,15 @@ assemblyline_v4_service/common/result.py,sha256=9AqM6qCYiia_Bpyn_fBFhzNQMcqJbtFS
 assemblyline_v4_service/common/task.py,sha256=dJsvRpW0x88CCF_LW6w87jQ_UKTVaOs2Gb117IDNiU8,14233
 assemblyline_v4_service/common/utils.py,sha256=k2__d-V5LjB6o2IKbjVe7tJWKcKuUHto5TyT5oKhIa0,3890
 assemblyline_v4_service/dev/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-assemblyline_v4_service/dev/run_service_once.py,sha256=1sdxkRBI7x9Sht8EVCTMIW3MgSxIeV6Ufw77Wx9Vf20,10150
+assemblyline_v4_service/dev/run_service_once.py,sha256=wx_82hGwavxWhMsmNDfmlYs4KbcLnqeMg91ZGfXTZV0,10621
+assemblyline_v4_service/dev/updater.py,sha256=b-FK6XPRZbETbl-SIYEhnYGT-W7EcQhnxwD6x2NMC7g,6411
 assemblyline_v4_service/updater/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 assemblyline_v4_service/updater/__main__.py,sha256=9Os-u8Tf7MD73JSrUSPmOaErTgfvesNLiEeszU4ujXA,133
 assemblyline_v4_service/updater/app.py,sha256=Mtmx4bkXfP4nFqqa5q15jW8QIXr4JK84lCovxAVyvPs,3317
 assemblyline_v4_service/updater/client.py,sha256=tLY84gaGdFBVIDaMgRHIEa7x2S8jBl7lQLzp4seC6aI,11200
 assemblyline_v4_service/updater/gunicorn_config.py,sha256=p3j2KPBeD5jvMw9O5i7vAtlRgPSVVxIG9AO0DfN82J8,1247
 assemblyline_v4_service/updater/helper.py,sha256=Zy6OBmbTh0YurW0MnM0wM92vaKYMbo_MKnafe_5ONUI,10034
-assemblyline_v4_service/updater/updater.py,sha256=HqRVUs5nNIfKmHu7g1GUkgC7bvYWs2EC9ZFy1kj-Ltw,31609
+assemblyline_v4_service/updater/updater.py,sha256=kli-5v1uVmk2FARAI9DsZ9YM4EhgirkmWJaMJWdm9GI,31795
 test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 test/conftest.py,sha256=W3SieQpZsZpGEmtLqY4aIlxREDSsHceyCrFcFsWUM0U,1851
 test/test_healthz.py,sha256=DkeLUlrb7rGx3nZ04aADU9HXXu5mZTf_DBwT0xhzIv4,7
@@ -38,8 +39,8 @@ test/test_common/test_request.py,sha256=Ceyds8BNO1O0f1kH1VEb84faJcaupvSjVKIrGdHe
 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.dev453.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
-assemblyline_v4_service-4.5.1.dev453.dist-info/METADATA,sha256=zXKPiqpBNkTlyOEai0ypH-hT3VkhfZfex8xS-8aoFN0,5603
-assemblyline_v4_service-4.5.1.dev453.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-assemblyline_v4_service-4.5.1.dev453.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
-assemblyline_v4_service-4.5.1.dev453.dist-info/RECORD,,
+assemblyline_v4_service-4.5.1.dev455.dist-info/LICENCE.md,sha256=NSkYo9EH8h5oOkzg4VhjAHF4339MqPP2cQ8msTPgl-c,1396
+assemblyline_v4_service-4.5.1.dev455.dist-info/METADATA,sha256=p06bT3Ftf1HlTzKgyrzS-Bxv-5hdvlipGbHjJdYVE_w,5603
+assemblyline_v4_service-4.5.1.dev455.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+assemblyline_v4_service-4.5.1.dev455.dist-info/top_level.txt,sha256=LpTOEaVCatkrvbVq3EZseMSIa2PQZU-2rhuO_FTpZgY,29
+assemblyline_v4_service-4.5.1.dev455.dist-info/RECORD,,