[client] Fix some bugs and add the tags management

Author: Samuel Hassine

examples/add_tag_to_malwre.py examples/add_tag_to_malware.pyexamples/add_tag_to_malwre.py renamed to examples/add_tag_to_malware.py

@@ -17,8 +17,18 @@
 )
 
 # Create the tag (if not exists)
+tag = opencti_api_client.tag.create(
+    tag_type='Malware-Type',
+    value='Ranswomware',
+    color='#ffa500',
+)
 
 # Add the tag
+opencti_api_client.stix_entity.add_tag(
+    id=malware['id'],
+    tag_id=tag['id']
+)
 
 # Print
+malware = opencti_api_client.malware.read(id=malware['id'])
 print(malware)

examples/create_incident_with_ttps_and_indicators.py

@@ -6,7 +6,7 @@
 from pycti import OpenCTIApiClient
 
 # Variables
-api_url = 'https://demo.opencti.io'
+api_url = 'http://localhost:4000'
 api_token = 'fa63eb1f-bf14-4777-9190-43b4571cbc8b'
 
 # OpenCTI initialization
@@ -166,4 +166,4 @@
 for object_ref in object_refs:
     opencti_api_client.report.add_stix_entity(id=report['id'], report=report, entity_id=object_ref)
 for observable_ref in observable_refs:
-    opencti_api_client.report.add_stix_observable(id=report['id'], report=report, entity_id=observable_ref)
+    opencti_api_client.report.add_stix_observable(id=report['id'], report=report, stix_observable_id=observable_ref)

pycti/__init__.py

@@ -6,6 +6,7 @@
 from pycti.connector.opencti_connector import OpenCTIConnector
 from pycti.connector.opencti_connector_helper import OpenCTIConnectorHelper, get_config_variable
 
+from pycti.entities.opencti_tag import Tag
 from pycti.entities.opencti_marking_definition import MarkingDefinition
 from pycti.entities.opencti_external_reference import ExternalReference
 from pycti.entities.opencti_kill_chain_phase import KillChainPhase

pycti/api/opencti_api_client.py

@@ -16,6 +16,7 @@
 from pycti.utils.constants import ObservableTypes
 from pycti.utils.opencti_stix2 import OpenCTIStix2
 
+from pycti.entities.opencti_tag import Tag
 from pycti.entities.opencti_marking_definition import MarkingDefinition
 from pycti.entities.opencti_external_reference import ExternalReference
 from pycti.entities.opencti_kill_chain_phase import KillChainPhase
@@ -78,6 +79,7 @@ def __init__(self, url, token, log_level='info', ssl_verify=False):
         self.stix2 = OpenCTIStix2(self)
 
         # Define the entities
+        self.tag = Tag(self)
         self.marking_definition = MarkingDefinition(self)
         self.external_reference = ExternalReference(self)
         self.kill_chain_phase = KillChainPhase(self)

pycti/entities/opencti_tag.py

@@ -0,0 +1,151 @@
+# coding: utf-8
+
+import json
+
+
+class Tag:
+    def __init__(self, opencti):
+        self.opencti = opencti
+        self.properties = """
+            id
+            tag_type
+            value
+            color
+            created_at
+            updated_at
+        """
+
+    """
+        List Tag objects
+
+        :param filters: the filters to apply
+        :param first: return the first n rows from the after ID (or the beginning if not set)
+        :param after: ID of the first row for pagination
+        :return List of Tag objects
+    """
+
+    def list(self, **kwargs):
+        filters = kwargs.get('filters', None)
+        first = kwargs.get('first', 500)
+        after = kwargs.get('after', None)
+        order_by = kwargs.get('orderBy', None)
+        order_mode = kwargs.get('orderMode', None)
+        self.opencti.log('info', 'Listing Tags with filters ' + json.dumps(filters) + '.')
+        query = """
+            query Tags($filters: [TagsFiltering], $first: Int, $after: ID, $orderBy: TagsOrdering, $orderMode: OrderingMode) {
+                tags(filters: $filters, first: $first, after: $after, orderBy: $orderBy, orderMode: $orderMode) {
+                    edges {
+                        node {
+                            """ + self.properties + """
+                        }
+                    }
+                    pageInfo {
+                        startCursor
+                        endCursor
+                        hasNextPage
+                        hasPreviousPage
+                        globalCount
+                    }                    
+                }
+            }
+        """
+        result = self.opencti.query(query, {
+            'filters': filters,
+            'first': first,
+            'after': after,
+            'orderBy': order_by,
+            'orderMode': order_mode
+        })
+        return self.opencti.process_multiple(result['data']['tags'])
+
+    """
+        Read a Tag object
+
+        :param id: the id of the Tag
+        :param filters: the filters to apply if no id provided
+        :return Marking-Definition object
+    """
+
+    def read(self, **kwargs):
+        id = kwargs.get('id', None)
+        filters = kwargs.get('filters', None)
+        if id is not None:
+            self.opencti.log('info', 'Reading Tag {' + id + '}.')
+            query = """
+                query Tag($id: String!) {
+                    tag(id: $id) {
+                        """ + self.properties + """
+                    }
+                }
+            """
+            result = self.opencti.query(query, {'id': id})
+            return self.opencti.process_multiple_fields(result['data']['tag'])
+        elif filters is not None:
+            result = self.list(filters=filters)
+            if len(result) > 0:
+                return result[0]
+            else:
+                return None
+        else:
+            self.opencti.log('error', 'Missing parameters: id or filters')
+            return None
+
+    """
+        Create a Tag object
+
+        :param tag_type: the tag type
+        :param value: the value
+        :param color: the color
+        :return Tag object
+    """
+
+    def create_raw(self, **kwargs):
+        tag_type = kwargs.get('tag_type', None)
+        value = kwargs.get('value', None)
+        color = kwargs.get('color', None)
+
+        if tag_type is not None and value is not None and color is not None:
+            query = """
+                mutation TagAdd($input: TagAddInput) {
+                    tagAdd(input: $input) {
+                        """ + self.properties + """
+                    }
+                }
+            """
+            result = self.opencti.query(query, {
+                'input': {
+                    'tag_type': tag_type,
+                    'value': value,
+                    'color': color
+                }
+            })
+            return self.opencti.process_multiple_fields(result['data']['tagAdd'])
+        else:
+            self.opencti.log('error', '[opencti_tag] Missing parameters: tag_type and value and color')
+
+    """
+        Create a Tag object only if it not exists, update it on request
+
+        :param tag_type: the tag type
+        :param value: the value
+        :param color: the color
+        :return Tag object
+    """
+
+    def create(self, **kwargs):
+        tag_type = kwargs.get('tag_type', None)
+        value = kwargs.get('value', None)
+        color = kwargs.get('color', None)
+
+        object_result = self.read(filters=[
+            {'key': 'tag_type', 'values': [tag_type]},
+            {'key': 'value', 'values': [value]}
+        ])
+        if object_result is not None:
+            return object_result
+        else:
+            return self.create_raw(
+                tag_type=tag_type,
+                value=value,
+                color=color,
+            )

pycti/utils/opencti_stix2.py

@@ -254,7 +254,7 @@ def extract_embedded_relationships(self, stix_object, types=None):
                         published=published,
                         report_class='Threat Report',
                         object_status=2,
-                        createdByRef=author,
+                        createdByRef=author['id'] if author is not None else None,
                         update=True
                     )
                     # Add marking
@@ -1423,8 +1423,10 @@ def import_bundle(self, stix_bundle, update=False, types=None) -> List:
         for item in stix_bundle['objects']:
             if item['type'] == 'relationship':
                 # Import only relationships between entities
-                if 'relationship' not in item[CustomProperties.SOURCE_REF] and \
-                        'relationship' not in item[CustomProperties.TARGET_REF]:
+                if (CustomProperties.SOURCE_REF not in item or 'relationship' not in item[
+                    CustomProperties.SOURCE_REF]) and (
+                        CustomProperties.TARGET_REF not in item or 'relationship' not in item[
+                    CustomProperties.TARGET_REF]):
                     self.import_relationship(item, update, types)
                     imported_elements.append({'id': item['id'], 'type': item['type']})
         end_time = time.time()
@@ -1434,8 +1436,8 @@ def import_bundle(self, stix_bundle, update=False, types=None) -> List:
         start_time = time.time()
         for item in stix_bundle['objects']:
             if item['type'] == 'relationship':
-                if 'relationship' in item[CustomProperties.SOURCE_REF] or \
-                        'relationship' in item[CustomProperties.TARGET_REF]:
+                if (CustomProperties.SOURCE_REF in item and 'relationship' in item[CustomProperties.SOURCE_REF]) or (
+                        CustomProperties.TARGET_REF in item and 'relationship' in item[CustomProperties.TARGET_REF]):
                     self.import_relationship(item, update, types)
                     imported_elements.append({'id': item['id'], 'type': item['type']})
         end_time = time.time()