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

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

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

assemblyline-core-4.5.1.dev177/assemblyline_core/VERSION

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

{assemblyline-core-4.5.1.dev173 → assemblyline-core-4.5.1.dev177}/assemblyline_core/signature_client.py

@@ -17,12 +17,13 @@ CLASSIFICATION = forge.get_classification()
 class SignatureClient:
     """A helper class to simplify signature management for privileged services and service-server."""
 
-    def __init__(self, datastore: AssemblylineDatastore = None, config=None):
+    def __init__(self, datastore: AssemblylineDatastore = None, config=None, classification_replace_map={}):
         self.log = logging.getLogger('assemblyline.signature_client')
         self.config = config or forge.CachedObject(forge.get_config)
         self.datastore = datastore or forge.get_datastore(self.config)
         self.service_list = forge.CachedObject(self.datastore.list_all_services, kwargs=dict(as_obj=False, full=True))
         self.delimiters = forge.CachedObject(self._get_signature_delimiters)
+        self.classification_replace_map = classification_replace_map
 
     def _get_signature_delimiters(self):
         signature_delimiters = {}
@@ -39,6 +40,22 @@ class SignatureClient:
             delimiter = SIGNATURE_DELIMITERS.get(delimiter_type, '\n\n')
         return {'type': delimiter_type, 'delimiter': delimiter}
 
+    def _update_classification(self, signature):
+        classification = signature['classification']
+        # Update classification of signatures based on rewrite definition
+        for term, replacement in self.classification_replace_map.items():
+            if replacement.startswith('_'):
+                # Replace with known field in Signature model
+                # Otherwise replace with literal
+                if signature.get(replacement[1:]):
+                    replacement = signature[replacement[1:]]
+
+            classification = classification.replace(term, replacement)
+
+        # Save the (possibly) updated classfication
+        signature['classification'] = classification
+
+
     def add_update(self, data, dedup_name=True):
         if data.get('type', None) is None or data['name'] is None or data['data'] is None:
             raise ValueError("Signature id, name, type and data are mandatory fields.")
@@ -79,6 +96,8 @@ class SignatureClient:
             # Preserve signature stats
             data['stats'] = old['stats']
 
+        self._update_classification(data)
+
         # Save the signature
         success = self.datastore.signature.save(key, data)
         return success, key, op
@@ -126,6 +145,8 @@ class SignatureClient:
                 # Preserve signature stats
                 rule['stats'] = old_data[key]['stats']
 
+            self._update_classification(rule)
+
             plan.add_upsert_operation(key, rule)
 
         if not plan.empty:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: assemblyline-core
-Version: 4.5.1.dev173
+Version: 4.5.1.dev177
 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-4.5.1.dev177}/test/test_signature_client.py

@@ -103,3 +103,16 @@ def test_download_signatures(client):
 def test_update_available(client):
     assert client.update_available()
     assert not client.update_available(since='2030-01-01T00:00:00.000000Z')
+
+def test_update_classification(client):
+    sig = client.datastore.signature.search("*", rows=1, as_obj=False)['items'][0]
+
+    # Update classification with literal string
+    client.classification_replace_map = {"TLP:C": "TLP:A//TEST"}
+    client._update_classification(sig)
+    assert sig['classification'] == "TLP:A//TEST"
+
+    # Update classification with value from another field within the signature
+    client.classification_replace_map = {"TEST": "_source"}
+    client._update_classification(sig)
+    assert sig['classification'] == f"TLP:A//{sig['source']}"

assemblyline-core-4.5.1.dev173/assemblyline_core/VERSION

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