From 586ae7746e25066961559e0da6a6453c24f68177 Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Mon, 26 Aug 2024 17:33:26 +0000 Subject: [PATCH 1/2] Added support for classification_aliases and metadata_suggestions --- .../v4_client/module/system.py | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/assemblyline_client/v4_client/module/system.py b/assemblyline_client/v4_client/module/system.py index 5bc17aa..374124b 100644 --- a/assemblyline_client/v4_client/module/system.py +++ b/assemblyline_client/v4_client/module/system.py @@ -1,6 +1,6 @@ from json import dumps -from assemblyline_client.v4_client.common.utils import api_path +from assemblyline_client.v4_client.common.utils import api_path, get_function_kwargs class System(object): @@ -49,3 +49,58 @@ def set_tag_safelist(self, yaml_file): yaml_file: New tag_safelist.yml file """ return self._connection.put(api_path('system', 'tag_safelist'), data=dumps(yaml_file)) + + def get_classification_aliases(self): + """\ +Get the current display aliases for the classification engine +""" + return self._connection.get(api_path('system', 'classification_aliases')) + + def set_classification_aliases(self, aliases): + """\ +Save display aliases for the classification engine + +Required: +aliases: A dictionary of alias name and short_names for a given classifications strings + +Example data block: +{ + "ORG_000000": {"name": "Communication Security Establishment", + "short_name": "CSE"}, + "ORG_000001": {"name": "Canadian Center for Cyber Security", + "short_name": "CCCS"}, + ... +} +""" + return self._connection.put(api_path('system', 'classification_aliases'), data=dumps(aliases)) + + def get_metadata_suggestions(self, key=None): + """\ +Get the current metadata suggestions + +Optional: +key: Get only the suggestions values for the given key +""" + kw = get_function_kwargs('key') + return self._connection.get(api_path('system', 'metadata_suggestions', **kw)) + + def set_metadata_suggestions(self, suggestions, key=None): + """\ +Set the metadata suggestions + +Required: +aliases: A dictionary with a list of suggestions for each key or + A list of suggestion if used on a specific key + +Example data block: +{ + "key_1": ["a", "b", "c"], + "key_2": ["d", "e", "f"], +} + +or + +["a", "b", "c"] # If use with a key +""" + kw = get_function_kwargs('key') + return self._connection.put(api_path('system', 'metadata_suggestions', **kw), data=dumps(suggestions)) From 025f88ea9b3eaabfd25959151df20589579158e6 Mon Sep 17 00:00:00 2001 From: Steve Garon Date: Tue, 27 Aug 2024 12:35:04 +0000 Subject: [PATCH 2/2] Fix key selection --- assemblyline_client/v4_client/module/system.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/assemblyline_client/v4_client/module/system.py b/assemblyline_client/v4_client/module/system.py index 374124b..94d8e59 100644 --- a/assemblyline_client/v4_client/module/system.py +++ b/assemblyline_client/v4_client/module/system.py @@ -1,6 +1,6 @@ from json import dumps -from assemblyline_client.v4_client.common.utils import api_path, get_function_kwargs +from assemblyline_client.v4_client.common.utils import api_path class System(object): @@ -81,7 +81,9 @@ def get_metadata_suggestions(self, key=None): Optional: key: Get only the suggestions values for the given key """ - kw = get_function_kwargs('key') + kw = {} + if key: + kw['key'] = key return self._connection.get(api_path('system', 'metadata_suggestions', **kw)) def set_metadata_suggestions(self, suggestions, key=None): @@ -102,5 +104,7 @@ def set_metadata_suggestions(self, suggestions, key=None): ["a", "b", "c"] # If use with a key """ - kw = get_function_kwargs('key') + kw = {} + if key: + kw['key'] = key return self._connection.put(api_path('system', 'metadata_suggestions', **kw), data=dumps(suggestions))