assemblyline-core 4.5.1.dev171__tar.gz → 4.5.1.dev173__tar.gz

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-core
-Version: 4.5.1.dev171
+Version: 4.5.1.dev173
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

assemblyline-core-4.5.1.dev173/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.1.dev173

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core/replay/client.py

@@ -7,8 +7,10 @@ from assemblyline.common.bundling import create_bundle, import_bundle
 from assemblyline.odm import Model
 from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline.remote.datatypes.hash import Hash
+from assemblyline_core.replay.replay import INPUT_TYPES
 from assemblyline_core.badlist_client import BadlistClient
 from assemblyline_core.safelist_client import SafelistClient
+from assemblyline_core.signature_client import SignatureClient
 
 EMPTY_WAIT_TIME = int(os.environ.get('EMPTY_WAIT_TIME', '30'))
 REPLAY_REQUESTED = 'requested'
@@ -18,7 +20,7 @@ REPLAY_DONE = 'done'
 
 class ClientBase(object):
     def __init__(self, log, lookback_time='*',
-                 alert_fqs=None, badlist_fqs=None, safelist_fqs=None, submission_fqs=None, workflow_fqs=None):
+                 alert_fqs=None, badlist_fqs=None, safelist_fqs=None, signature_fqs=None, submission_fqs=None, workflow_fqs=None):
         # Set logger
         self.log = log
 
@@ -31,6 +33,7 @@ class ClientBase(object):
         self.alert_fqs = alert_fqs or []
         self.badlist_fqs = badlist_fqs or []
         self.safelist_fqs = safelist_fqs or []
+        self.signature_fqs = signature_fqs or []
         self.submission_fqs = submission_fqs or []
         self.workflow_fqs = workflow_fqs or []
 
@@ -202,15 +205,18 @@ class ClientBase(object):
             if once:
                 break
 
-    def setup_workflow_input_queue(self, once=False):
-        self._setup_checkpoint_based_input_queue("workflow", "workflow_id", "last_edit", once)
-
     def setup_badlist_input_queue(self, once=False):
         self._setup_checkpoint_based_input_queue("badlist", "id", "updated", once)
 
     def setup_safelist_input_queue(self, once=False):
         self._setup_checkpoint_based_input_queue("safelist", "id", "updated", once)
 
+    def setup_signature_input_queue(self, once=False):
+        self._setup_checkpoint_based_input_queue("signature", "id", "last_modified", once)
+
+    def setup_workflow_input_queue(self, once=False):
+        self._setup_checkpoint_based_input_queue("workflow", "workflow_id", "last_edit", once)
+
     def _query(self, collection, query, filter_queries=[], rows=None, track_total_hits=False):
         raise NotImplementedError()
 
@@ -232,6 +238,9 @@ class ClientBase(object):
     def get_next_safelist(self):
         return self.get_next_message("safelist")
 
+    def get_next_signature(self):
+        return self.get_next_message("signature")
+
     def get_next_submission(self):
         return self.get_next_message("submission")
 
@@ -321,6 +330,9 @@ class APIClient(ClientBase):
                     elif collection == "safelist":
                         data['enabled'] = obj["enabled"]
                         self.al_client.safelist.add_update(data)
+                    elif collection == "signature":
+                        data['status'] = obj["status"]
+                        self.al_client.signature.add_update(data)
                 except ClientError as e:
                     if e.status_code == 404:
                         # The document doesn't exist in the system, therefore create it
@@ -330,6 +342,8 @@ class APIClient(ClientBase):
                             self.al_client.badlist.add_update(data)
                         elif collection == "safelist":
                             self.al_client.safelist.add_update(data)
+                        elif collection == "signature":
+                            self.al_client.signature.add_update(data)
                         return
                     raise
 
@@ -363,7 +377,7 @@ class DirectClient(ClientBase):
         self.datastore = forge.get_datastore(config=config)
         self.queues = {
             queue_type: NamedQueue(f"replay_{queue_type}", host=redis)
-            for queue_type in ['alert', 'file', 'submission', 'safelist', 'badlist', 'workflow']
+            for queue_type in INPUT_TYPES + ['file']
         }
         self.checkpoint_hash = Hash('replay_checkpoints', redis_persist)
 
@@ -431,6 +445,11 @@ class DirectClient(ClientBase):
                         # Preserve the system's enabled state of the item
                         data['enabled'] = obj["enabled"]
                     es_collection.save(id, SafelistClient._merge_hashes(data, obj))
+                elif collection == "signature":
+                    if obj:
+                        # Preserve the system's status state of the item
+                        data['status'] = obj["status"]
+                    es_collection.save(id, data)
             es_collection.commit()
 
     def set_single_object_complete(self, collection, id):

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core/replay/creator/run.py

@@ -1,8 +1,7 @@
 import os
 
 from assemblyline_core.replay.client import APIClient, DirectClient
-from assemblyline_core.replay.replay import ReplayBase
-
+from assemblyline_core.replay.replay import ReplayBase, INPUT_TYPES
 
 class ReplayCreator(ReplayBase):
     def __init__(self):
@@ -16,12 +15,9 @@ class ReplayCreator(ReplayBase):
         os.makedirs(self.replay_config.creator.working_directory, exist_ok=True)
 
         # Load client
-        client_config = dict(lookback_time=self.replay_config.creator.lookback_time,
-                             alert_fqs=self.replay_config.creator.alert_input.filter_queries,
-                             badlist_fqs=self.replay_config.creator.badlist_input.filter_queries,
-                             safelist_fqs=self.replay_config.creator.safelist_input.filter_queries,
-                             submission_fqs=self.replay_config.creator.submission_input.filter_queries,
-                             workflow_fqs=self.replay_config.creator.workflow_input.filter_queries)
+        client_config = {f'{input_type}_fqs': getattr(self.replay_config.creator, f'{input_type}_input').filter_queries
+                         for input_type in INPUT_TYPES}
+        client_config['lookback_time'] = self.replay_config.creator.lookback_time
 
         if self.replay_config.creator.client.type == 'direct':
             self.log.info("Using direct database access client")
@@ -36,20 +32,9 @@ class ReplayCreator(ReplayBase):
 
     def try_run(self):
         threads = {}
-        if self.replay_config.creator.alert_input.enabled:
-            threads['Load Alerts'] = self.client.setup_alert_input_queue
-
-        if self.replay_config.creator.badlist_input.enabled:
-            threads['Load Badlist Items'] = self.client.setup_badlist_input_queue
-
-        if self.replay_config.creator.safelist_input.enabled:
-            threads['Load Safelist Items'] = self.client.setup_safelist_input_queue
-
-        if self.replay_config.creator.submission_input.enabled:
-            threads['Load Submissions'] = self.client.setup_submission_input_queue
-
-        if self.replay_config.creator.workflow_input.enabled:
-            threads['Load Workflows'] = self.client.setup_workflow_input_queue
+        for input_type in INPUT_TYPES:
+            if getattr(self.replay_config.creator, f'{input_type}_input').enabled:
+                threads[f'Load {input_type.capitalize()}s'] = getattr(self.client, f'setup_{input_type}_input_queue')
 
         if threads:
             self.maintain_threads(threads)

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core/replay/creator/run_worker.py

@@ -4,7 +4,7 @@ import os
 from assemblyline.filestore import FileStore
 from assemblyline.common.isotime import now_as_iso
 from assemblyline_core.replay.client import APIClient, DirectClient
-from assemblyline_core.replay.replay import ReplayBase
+from assemblyline_core.replay.replay import ReplayBase, INPUT_TYPES
 
 REPLAY_BATCH_SIZE = int(os.environ.get("REPLAY_BATCH_SIZE", "1000"))
 
@@ -39,7 +39,7 @@ class ReplayCreatorWorker(ReplayBase):
             raise ValueError(f'Invalid client type ({self.replay_config.creator.client.type}). '
                              'Must be either \'api\' or \'direct\'.')
 
-    def process_alerts(self, once=False):
+    def process_alert(self, once=False):
         while self.running:
             # Process alerts found
             alert = self.client.get_next_alert()
@@ -67,7 +67,7 @@ class ReplayCreatorWorker(ReplayBase):
             if once:
                 break
 
-    def process_submissions(self, once=False):
+    def process_submission(self, once=False):
         while self.running:
             # Process submissions found
             submission = self.client.get_next_submission()
@@ -151,31 +151,19 @@ class ReplayCreatorWorker(ReplayBase):
     def process_safelist(self, once=False):
         self._process_json_exports("safelist", "id", "updated", once)
 
+    def process_signature(self, once=False):
+        self._process_json_exports("signature", "id", "last_modified", once)
+
     def process_workflow(self, once=False):
         self._process_json_exports("workflow", "id", "last_edit", once)
 
     def try_run(self):
         threads = {}
-        if self.replay_config.creator.alert_input.enabled:
-            for ii in range(self.replay_config.creator.alert_input.threads):
-                threads[f'Alert process thread #{ii}'] = self.process_alerts
-
-        if self.replay_config.creator.badlist_input.enabled:
-            for ii in range(self.replay_config.creator.badlist_input.threads):
-                threads[f'Badlist process thread #{ii}'] = self.process_badlist
-
-        if self.replay_config.creator.safelist_input.enabled:
-            for ii in range(self.replay_config.creator.safelist_input.threads):
-                threads[f'Safelist process thread #{ii}'] = self.process_safelist
-
-        if self.replay_config.creator.submission_input.enabled:
-            for ii in range(self.replay_config.creator.submission_input.threads):
-                threads[f'Submission process thread #{ii}'] = self.process_submissions
-
-        if self.replay_config.creator.workflow_input.enabled:
-            for ii in range(self.replay_config.creator.workflow_input.threads):
-                threads[f'Workflow process thread #{ii}'] = self.process_workflow
-
+        for input_type in INPUT_TYPES:
+            input_config = getattr(self.replay_config.creator, f"{input_type}_input")
+            if input_config.enabled:
+                for ii in range(input_config.threads):
+                    threads[f"{input_type.capitalize()} process thread #{ii}"] = getattr(self, f"process_{input_type}")
         if threads:
             self.maintain_threads(threads)
         else:

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core/replay/replay.py

@@ -11,6 +11,7 @@ from assemblyline.odm.models.replay import ReplayConfig
 from assemblyline_core.server_base import ServerBase
 
 CONFIG_PATH = os.environ.get('REPLAY_CONFIG_PATH', '/etc/assemblyline/replay.yml')
+INPUT_TYPES = ['alert', 'badlist', 'safelist', 'signature', 'submission', 'workflow']
 
 
 class ReplayBase(ServerBase):

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core/updater/helper.py

@@ -227,7 +227,7 @@ def _get_dockerhub_tags(image_name, update_channel, proxies=None):
         elif resp.status_code == 429:
             # Based on https://docs.docker.com/docker-hub/api/latest/#tag/rate-limiting
             # We've hit the rate limit so we have to wait and try again later
-            time.sleep(int(time.time()) - int(resp.headers['retry-after']))
+            time.sleep(int(resp.headers['retry-after']) - int(time.time()))
         else:
             break
 

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/assemblyline_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-core
-Version: 4.5.1.dev171
+Version: 4.5.1.dev173
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

{assemblyline-core-4.5.1.dev171 → assemblyline-core-4.5.1.dev173}/test/test_replay.py

@@ -2,11 +2,12 @@ import collections
 import json
 import os
 import random
+import time
 
 import pytest
 
 from assemblyline.common import forge
-from assemblyline.odm.random_data import create_alerts, wipe_alerts, wipe_submissions, create_submission, create_badlists, create_safelists, create_workflows, wipe_badlist, wipe_safelist, wipe_workflows
+from assemblyline.odm.random_data import create_alerts, wipe_alerts, wipe_submissions, create_submission, create_badlists, create_safelists, create_workflows, wipe_badlist, wipe_safelist, wipe_workflows, create_signatures, wipe_signatures
 from assemblyline_core.replay.creator.run import ReplayCreator
 from assemblyline_core.replay.creator.run_worker import ReplayCreatorWorker
 from assemblyline_core.replay.loader.run import ReplayLoader
@@ -15,6 +16,7 @@ from assemblyline_core.replay.loader.run_worker import ReplayLoaderWorker
 NUM_ALERTS = 1
 NUM_BADLIST_ITEMS = 1
 NUM_SAFELIST_ITEMS = 1
+NUM_SIGNATURES = 1
 NUM_SUBMISSIONS = 1
 NUM_WORKFLOWS = 1
 
@@ -48,14 +50,24 @@ def datastore(request, datastore_connection, fs):
     wipe_badlist(datastore_connection)
     wipe_safelist(datastore_connection)
     wipe_submissions(datastore_connection, fs)
+    wipe_signatures(datastore_connection)
     wipe_workflows(datastore_connection)
 
     for _ in range(NUM_SUBMISSIONS):
         all_submissions.append(create_submission(datastore_connection, fs))
     create_alerts(datastore_connection, alert_count=NUM_ALERTS,
                   submission_list=all_submissions)
-    create_safelists(datastore_connection, count=NUM_SAFELIST_ITEMS)
     create_badlists(datastore_connection, count=NUM_BADLIST_ITEMS)
+
+    # Generate all signatures from testing set, but only keep what's being asked to limit to
+    create_signatures(datastore_connection)
+    data_collections["signature"] = \
+        datastore_connection.signature.search("*", rows=NUM_SIGNATURES, fl="id,*")['items']
+    wipe_signatures(datastore_connection)
+    for sig in data_collections["signature"]:
+        datastore_connection.signature.save(sig.id, sig)
+
+    create_safelists(datastore_connection, count=NUM_SAFELIST_ITEMS)
     create_workflows(datastore_connection, count=NUM_WORKFLOWS)
     for alert in datastore_connection.alert.stream_search("id:*", fl="*"):
         all_alerts.append(alert)
@@ -70,6 +82,7 @@ def datastore(request, datastore_connection, fs):
         wipe_alerts(datastore_connection)
         wipe_badlist(datastore_connection)
         wipe_safelist(datastore_connection)
+        wipe_signatures(datastore_connection)
         wipe_submissions(datastore_connection, fs)
         wipe_workflows(datastore_connection)
 
@@ -141,7 +154,7 @@ def test_replay_single_alert(config, datastore, creator, creator_worker, loader,
         'alert_id'] == alert.alert_id
 
     # Test replay creator worker
-    creator_worker.process_alerts(once=True)
+    creator_worker.process_alert(once=True)
     datastore.alert.commit()
     assert creator_worker.client.queues['alert'].length() == 0
     assert datastore.alert.get(alert.alert_id, as_obj=False)['metadata']['replay'] == 'done'
@@ -190,7 +203,7 @@ def test_replay_single_submission(config, datastore, creator, creator_worker, lo
     assert creator.client.queues['submission'].peek_next()['sid'] == sub['sid']
 
     # Test replay creator worker
-    creator_worker.process_submissions(once=True)
+    creator_worker.process_submission(once=True)
     datastore.submission.commit()
     assert creator_worker.client.queues['submission'].length() == 0
     assert datastore.submission.get(sub['sid'], as_obj=False)['metadata']['replay'] == 'done'
@@ -222,7 +235,7 @@ def test_replay_single_submission(config, datastore, creator, creator_worker, lo
     assert 'replay' not in loaded_submission['metadata']
 
 
-@pytest.mark.parametrize("collection", ["badlist", "safelist", "workflow"])
+@pytest.mark.parametrize("collection", ["badlist", "safelist", "signature", "workflow"])
 def test_replay_single_data_collection(datastore, creator, creator_worker, loader, loader_worker, collection):
     output_dir = creator.replay_config.creator.output_filestore.replace('file://', '')
     input_dir = loader.replay_config.loader.input_directory

assemblyline-core-4.5.1.dev171/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.1.dev171