From 586ae7746e25066961559e0da6a6453c24f68177 Mon Sep 17 00:00:00 2001
From: Steve Garon <steve.garon@cyber.gc.ca>
Date: Mon, 26 Aug 2024 17:33:26 +0000
Subject: [PATCH] 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))