Dont add ext def to report (#130)

* remove ext defs from report.object_refs #126

* make `incident_classification` a list #127

---------

Co-authored-by: Fadl <[email protected]>

Author: himynamesdave

txt2stix/ai_extractor/utils.py

@@ -33,7 +33,7 @@ class RelationshipList(BaseModel):
 class DescribesIncident(BaseModel):
     describes_incident: bool = Field(description="does the <document> include malware analysis, APT group reports, data breaches and vulnerabilities?")
     explanation: str = Field(description="Two or three sentence summary of the incidents it describes OR summary of what it describes instead of an incident")
-    incident_classification : str = Field(description="One of valid incident classifications that best describes this document/report")
+    incident_classification : list[str] = Field(description="All the valid incident classifications that describe this document/report")
 
 class AttackFlowItem(BaseModel):
     position : int = Field(description="order of object starting at 0")

txt2stix/stix.py

@@ -240,6 +240,7 @@ def __init__(
             confidence=confidence,
         )
         self.report.object_refs.clear()  # clear object refs
+        self.added_objects = set()
         self.set_defaults()
 
     def set_defaults(self):
@@ -264,18 +265,20 @@ def add_extension(self, object):
             logger.info(f'getting extension definition for "{_type}" from `{url}`')
             self.EXTENSION_MAPPING[_type] = self.load_object_from_json(url)
             extension = self.EXTENSION_MAPPING[_type]
-            self.add_ref(extension)
+            self.add_ref(extension, is_report_object=False)
 
     @staticmethod
     def load_object_from_json(url):
         resp = requests.get(url)
         return dict_to_stix2(resp.json())
 
-    def add_ref(self, sdo):
+    def add_ref(self, sdo, is_report_object=True):
         self.add_extension(sdo)
         sdo_id = sdo["id"]
-        if sdo_id not in self.report.object_refs:
-            self.report.object_refs.append(sdo_id)
+        if sdo_id not in self.added_objects:
+            self.added_objects.add(sdo_id)
+            if is_report_object:
+                self.report.object_refs.append(sdo_id)
             self.bundle.objects.append(sdo)
 
         sdo_value = ""
@@ -426,5 +429,6 @@ def flow_objects(self, objects):
         for obj in objects:
             if obj['id'] == self.report.id:
                 continue
-            self.add_ref(obj)
+            is_report_object = obj['type'] != "extension-definition"
+            self.add_ref(obj, is_report_object=is_report_object)
         self._flow_objects = objects