Merge pull request #53 from CybercentreCanada/feature/safelist_signat…

…ures Feature/safelist signatures
CybercentreCanada · Sep 2, 2021 · 2be1ab9 · 2be1ab9
2 parents 9faca52 + 28b5699
commit 2be1ab9
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 14 deletions.
diff --git a/assemblyline_service_server/api/v1/safelist.py b/assemblyline_service_server/api/v1/safelist.py
@@ -97,3 +97,32 @@ def get_safelist_for_tags(**_):
             output['match'][sl['tag']['type']].append(sl['tag']['value'])
 
     return make_api_response(output)
+
+
+@safelist_api.route("/signatures/", methods=["GET"])
+@api_login()
+def get_safelist_for_signatures(**_):
+    """
+    Get the safelist for all heuristic's signatures
+
+    Variables:
+    None
+
+    Arguments:
+    None
+
+    Data Block:
+    None
+
+    API call example:
+    GET /api/v1/safelist/signatures/
+
+    Result example:
+    ["McAfee.Eicar", "Avira.Eicar", ...]
+    """
+    output = [
+        item['signature']['name']
+        for item in STORAGE.safelist.stream_search(
+            "type:signature AND enabled:true", fl="signature.name", as_obj=False)]
+
+    return make_api_response(output)
diff --git a/assemblyline_service_server/api/v1/task.py b/assemblyline_service_server/api/v1/task.py
@@ -7,7 +7,7 @@
 from assemblyline.common.constants import SERVICE_STATE_HASH, ServiceStatus
 from assemblyline.common.dict_utils import flatten, unflatten
 from assemblyline.common.forge import CachedObject
-from assemblyline.common.heuristics import service_heuristic_to_result_heuristic, InvalidHeuristicException
+from assemblyline.common.heuristics import HeuristicHandler, InvalidHeuristicException
 from assemblyline.common.isotime import now_as_iso
 from assemblyline.odm import construct_safe
 from assemblyline.odm.messages.service_heartbeat import Metrics
@@ -26,6 +26,7 @@
 status_table = ExpiringHash(SERVICE_STATE_HASH, ttl=60*30)
 dispatch_client = DispatchClient(STORAGE)
 heuristics = cast(Dict[str, Heuristic], CachedObject(get_heuristics, refresh=300))
+heuristic_hander = HeuristicHandler(STORAGE)
 tag_safelister = CachedObject(forge.get_tag_safelister,
                               kwargs=dict(log=LOGGER, config=config, datastore=STORAGE),
                               refresh=300)
@@ -214,12 +215,14 @@ def handle_task_result(exec_time: int, task: ServiceTask, result: Dict[str, Any]
     # Add scores to the heuristics, if any section set a heuristic
     total_score = 0
     for section in result['result']['sections']:
+        zeroize_on_sig_safe = section.pop('zeroize_on_sig_safe', True)
         section['tags'] = flatten(section['tags'])
         if section.get('heuristic'):
             heur_id = f"{client_info['service_name'].upper()}.{str(section['heuristic']['heur_id'])}"
             section['heuristic']['heur_id'] = heur_id
             try:
-                section['heuristic'], new_tags = service_heuristic_to_result_heuristic(section['heuristic'], heuristics)
+                section['heuristic'], new_tags = heuristic_hander.service_heuristic_to_result_heuristic(
+                    section['heuristic'], heuristics, zeroize_on_sig_safe)
                 for tag in new_tags:
                     section['tags'].setdefault(tag[0], [])
                     if tag[1] not in section['tags'][tag[0]]:

diff --git a/test/test_srv_safelist.py b/test/test_srv_safelist.py
@@ -56,12 +56,7 @@ def test_safelist_missing(client, storage):
 
 # noinspection PyUnusedLocal
 def test_get_full_safelist(client, storage):
-    storage.safelist.search = {
-        "offset": 0,
-        "rows": 0,
-        "total": 0,
-        "items": []
-    }
+    storage.safelist.stream_search.return_value = []
 
     resp = client.get('/api/v1/safelist/', headers=headers)
     assert resp.status_code == 200
@@ -73,12 +68,7 @@ def test_get_full_safelist(client, storage):
 
 # noinspection PyUnusedLocal
 def test_get_full_safelist_specific(client, storage):
-    storage.safelist.search = {
-        "offset": 0,
-        "rows": 0,
-        "total": 0,
-        "items": []
-    }
+    storage.safelist.stream_search.return_value = []
 
     tag_type = "network.dynamic.domain"
     resp = client.get(f'/api/v1/safelist/?tags={tag_type}', headers=headers)
@@ -93,3 +83,13 @@ def test_get_full_safelist_specific(client, storage):
 
     for k in resp.json['api_response']['regex']:
         assert k == tag_type
+
+
+# noinspection PyUnusedLocal
+def test_get_signature_safelist(client, storage):
+    storage.safelist.stream_search.return_value = [{"signature": {"name": "test"}}]
+
+    resp = client.get('/api/v1/safelist/signatures/', headers=headers)
+    assert resp.status_code == 200
+    assert isinstance(resp.json['api_response'], list)
+    assert resp.json['api_response'] == ['test']