[client/connector] Introduce methods to store key/value in the Redis through the API (#32)

Author: Samuel Hassine

pycti/api/opencti_api_client.py

@@ -133,6 +133,37 @@ def health_check(self):
             return False
         return False
 
+    def get_key_value(self, key):
+        query = """
+            query KeyValue($key: String!) {
+                keyValue(key: $key) {
+                    key
+                    value
+                }
+            }
+        """
+        result = self.query(query, {'key': key})
+        return result['data']['keyValue']
+
+    def set_key_value(self, key, value, timeout=None):
+        logging.info('Setting KeyValue ' + key + ': ' + value + '...')
+        query = """
+            mutation KeyValueAdd($input: KeyValueAddInput) {
+                keyValueAdd(input: $input) {
+                   key
+                   value
+                }
+            }
+        """
+        result = self.query(query, {
+            'input': {
+                'key': key,
+                'value': value,
+                'timeout': timeout
+            }
+        })
+        return result['data']['keyValueAdd']
+
     def parse_stix(self, data):
         if 'createdByRef' in data and data['createdByRef'] is not None and 'node' in data['createdByRef']:
             data['createdByRef'] = data['createdByRef']['node']

pycti/connector/opencti_connector_helper.py

@@ -149,6 +149,12 @@ def log_info(self, msg):
     def date_now(self):
         return datetime.datetime.utcnow().replace(microsecond=0, tzinfo=datetime.timezone.utc).isoformat()
 
+    def set_key_value(self, key, value, timeout=None):
+        return self.api.set_key_value(key, value, timeout)
+
+    def get_key_value(self, key):
+        return self.api.get_key_value(key)['value']
+
     # Push Stix2 helper
     def send_stix2_bundle(self, bundle, entities_types=None):
         if entities_types is None: