assemblyline-core 4.5.1.dev475__tar.gz → 4.5.1.dev480__tar.gz

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/PKG-INFO

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

assemblyline_core-4.5.1.dev480/assemblyline_core/VERSION

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

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/assemblyline_core/plumber/run_plumber.py

@@ -11,15 +11,18 @@ import warnings
 from typing import Optional
 
 from assemblyline.common.constants import service_queue_name
-from assemblyline.common.forge import get_service_queue
-from assemblyline.common.isotime import now_as_iso
+from assemblyline.common.forge import get_service_queue, get_config
+from assemblyline.common.isotime import DAY_IN_SECONDS, now_as_iso
+from assemblyline.odm.models.apikey import get_apikey_id
 from assemblyline.odm.models.error import Error
 from assemblyline.odm.models.service import Service
+from assemblyline.odm.models.user import load_roles, load_roles_form_acls
 from assemblyline.remote.datatypes import retry_call
 from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline_core.dispatching.client import DispatchClient
 from assemblyline_core.server_base import CoreBase, ServiceStage
 
+
 DAY = 60 * 60 * 24
 TASK_DELETE_CHUNK = 10000
 
@@ -53,6 +56,8 @@ class Plumber(CoreBase):
                                      name="redis_notification_queue_cleanup")
         nq_thread.start()
 
+        ua_thread = threading.Thread(target=self.user_apikey_cleanup, daemon=True, name="user_apikey_cleanup")
+        ua_thread.start()
         self.service_queue_plumbing()
 
     def service_queue_plumbing(self):
@@ -189,6 +194,63 @@ class Plumber(CoreBase):
         self.log.info(f"Done watching {service_name} service queue")
 
 
+
+    def user_apikey_cleanup(self):
+        query = "id:*"
+        offset = 0
+        rows = 100
+        total = 1
+        cur_total = 0
+
+        config = get_config()
+        apikey_max_dtl = config.auth.apikey_max_dtl
+
+        expiry_ts = now_as_iso(apikey_max_dtl * DAY_IN_SECONDS) if apikey_max_dtl is not None else None
+
+        while cur_total < total:
+            result = self.datastore.user.search(query, offset=offset, rows=rows)
+            total = result.get('total', 0)
+            cur_total = cur_total + (result.get("count", total))
+
+            # check for API keys in total
+            users = result.get('items', [])
+
+            for u in users:
+                uname = u['uname']
+                user = self.datastore.user.get(uname)
+                apikeys = user.apikeys
+
+                for key in apikeys:
+                    old_apikey = apikeys[key]
+                    key_id = get_apikey_id(key, uname)
+
+                    roles = None
+                    if old_apikey['acl'] == ["C"]:
+
+                        roles = [r for r in old_apikey['roles']
+                                    if r in load_roles(user['type'], user['roles'])]
+
+                    else:
+                        roles = [r for r in load_roles_form_acls(old_apikey['acl'], roles)
+                                if r in load_roles(user['type'], user['roles'])]
+                    new_apikey = {
+                        "password": old_apikey['password'],
+                        "acl": old_apikey['acl'],
+                        "uname": uname,
+                        "key_name": key,
+                        "roles": roles,
+                        "expiry_ts": expiry_ts
+                    }
+                    self.datastore.apikey.save(key_id, new_apikey)
+
+                user['apikeys'] = {}
+                self.datastore.user.save(uname, user)
+
+
+
+
+
+
 if __name__ == '__main__':
     with Plumber() as server:
         server.serve_forever()

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/assemblyline_core/workflow/run_workflow.py

@@ -59,7 +59,7 @@ class WorkflowManager(ServerBase):
 
         return ret_val
 
-    def try_run(self):
+    def try_run(self, run_once=False):
         self.datastore.alert.commit()
         while self.running:
             self.heartbeat()
@@ -220,6 +220,8 @@ class WorkflowManager(ServerBase):
             else:
                 self.log.info("Skipping all workflows since there where no new alerts in the specified time period.")
 
+            if run_once:
+                break
             time.sleep(30)
             self.start_ts = end_ts
 

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/assemblyline_core.egg-info/SOURCES.txt

@@ -84,4 +84,5 @@ test/test_simulation.py
 test/test_tasking_client.py
 test/test_vacuum.py
 test/test_worker_ingest.py
-test/test_worker_submit.py
+test/test_worker_submit.py
+test/test_workflow.py

{assemblyline_core-4.5.1.dev475 → assemblyline_core-4.5.1.dev480}/test/test_tasking_client.py

@@ -1,5 +1,3 @@
-
-import pytest
 from assemblyline_core.tasking_client import TaskingClient
 
 from assemblyline.odm.models.service import Service

assemblyline_core-4.5.1.dev480/test/test_workflow.py

@@ -0,0 +1,45 @@
+import pytest
+import random
+from assemblyline_core.workflow.run_workflow import WorkflowManager
+
+from assemblyline.common.isotime import now_as_iso
+from assemblyline.odm.models.workflow import Workflow
+from assemblyline.odm.random_data import create_alerts, wipe_alerts, wipe_workflows
+from assemblyline.odm.randomizer import random_minimal_obj
+
+
+@pytest.fixture(scope="module")
+def manager(datastore_connection):
+    try:
+        create_alerts(datastore_connection)
+        wipe_workflows(datastore_connection)
+        datastore_connection.alert.update_by_query("*", [(datastore_connection.alert.UPDATE_SET, 'reporting_ts', now_as_iso())])
+        datastore_connection.alert.commit()
+        yield WorkflowManager()
+    finally:
+        wipe_alerts(datastore_connection)
+
+def test_workflow(manager, datastore_connection):
+    # Create workflow that targets alerts based on YARA rule association
+    workflow = random_minimal_obj(Workflow)
+
+    yara_rule = random.choice(list(datastore_connection.alert.facet("al.yara").keys()))   
+    workflow.query = f'al.yara:"{yara_rule}"'
+    workflow.workflow_id = "AL_TEST"
+    workflow.labels = ["AL_TEST"]
+    workflow.priority = "LOW"
+    workflow.status = "MALICIOUS"    
+    datastore_connection.workflow.save(workflow.workflow_id, workflow)
+    datastore_connection.workflow.commit()
+
+    # Run Workflow manager to process new workflow against existing alerts
+    manager.running = True
+    manager.get_last_reporting_ts = lambda x: "now/d+1d"
+    manager.try_run(run_once=True)
+    datastore_connection.alert.commit()
+    
+    # Assert that custom labels were applied to alerts
+    assert datastore_connection.alert.search("label:AL_TEST", track_total_hits=True)['total']
+
+    # Assert that the change has been record in the alerts' event history
+    assert datastore_connection.alert.search(f"events.entity_id:{workflow.workflow_id}", track_total_hits=True)['total']

assemblyline_core-4.5.1.dev475/assemblyline_core/VERSION

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