From af2d02051d76784bb5b23ecf7f6b9435ae08b112 Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Mon, 30 Aug 2021 19:05:36 +0000 Subject: [PATCH 1/4] Zeroize sections if all heuristics are safe --- assemblyline_service_server/api/v1/task.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/assemblyline_service_server/api/v1/task.py b/assemblyline_service_server/api/v1/task.py index b8fced4..bd8d304 100644 --- 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]]: From 18a6f6953b0bc6b30f18fb803825913b6fc399af Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Tue, 31 Aug 2021 16:00:26 +0000 Subject: [PATCH 2/4] Add API to pull signature's safelist with test --- .../api/v1/safelist.py | 29 +++++++++++++++++++ test/test_srv_safelist.py | 15 ++++++++++ 2 files changed, 44 insertions(+) diff --git a/assemblyline_service_server/api/v1/safelist.py b/assemblyline_service_server/api/v1/safelist.py index b8d8fed..3261c5e 100644 --- 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/test/test_srv_safelist.py b/test/test_srv_safelist.py index 954f11c..e56bbac 100644 --- a/test/test_srv_safelist.py +++ b/test/test_srv_safelist.py @@ -93,3 +93,18 @@ 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.search = { + "offset": 0, + "rows": 1, + "total": 1, + "items": [{"signature": {"name": "test"}}] + } + + resp = client.get('/api/v1/safelist/', headers=headers) + assert resp.status_code == 200 + assert isinstance(resp.json['api_response'], list) + assert resp.json['api_response'] == ['test'] From 4cd8146e10fdd27cb0349773292c0bdcd7d2fa02 Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Tue, 31 Aug 2021 16:33:11 +0000 Subject: [PATCH 3/4] Fix typo --- test/test_srv_safelist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_srv_safelist.py b/test/test_srv_safelist.py index e56bbac..a12e540 100644 --- a/test/test_srv_safelist.py +++ b/test/test_srv_safelist.py @@ -104,7 +104,7 @@ def test_get_signature_safelist(client, storage): "items": [{"signature": {"name": "test"}}] } - resp = client.get('/api/v1/safelist/', headers=headers) + resp = client.get('/api/v1/safelist/signature/', headers=headers) assert resp.status_code == 200 assert isinstance(resp.json['api_response'], list) assert resp.json['api_response'] == ['test'] From 28b56993fd300042a620c8d9b290338b695797ed Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Tue, 31 Aug 2021 17:04:22 +0000 Subject: [PATCH 4/4] Fix tests --- test/test_srv_safelist.py | 25 +++++-------------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/test/test_srv_safelist.py b/test/test_srv_safelist.py index a12e540..1a05582 100644 --- 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) @@ -97,14 +87,9 @@ def test_get_full_safelist_specific(client, storage): # noinspection PyUnusedLocal def test_get_signature_safelist(client, storage): - storage.safelist.search = { - "offset": 0, - "rows": 1, - "total": 1, - "items": [{"signature": {"name": "test"}}] - } - - resp = client.get('/api/v1/safelist/signature/', headers=headers) + 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']