Leakyhandles

core/database_arango.py

@@ -115,9 +115,7 @@ def connect(
         self.db.collection("indicators").add_persistent_index(
             fields=["name", "type"], unique=True
         )
-        self.db.collection("dfiq").add_persistent_index(
-            fields=["dfiq_id", "type"], unique=True
-        )
+        self.db.collection("dfiq").add_persistent_index(fields=["uuid"], unique=True)
 
     def clear(self, truncate=True):
         if not self.db:

core/schemas/dfiq.py

@@ -102,7 +102,8 @@ class DFIQBase(YetiModel, database_arango.ArangoYetiConnector):
     _root_type: Literal["dfiq"] = "dfiq"
 
     name: str = Field(min_length=1)
-    dfiq_id: str = Field(min_length=1)
+    uuid: str  # = Field(default_factory=lambda: str(uuid.uuid4()))
+    dfiq_id: str | None = None
     dfiq_version: str = Field(min_length=1)
     dfiq_tags: list[str] | None = None
     contributors: list[str] | None = None
@@ -142,7 +143,7 @@ def parse_yaml(cls, yaml_string: str) -> dict[str, Any]:
         if "id" not in yaml_data:
             raise ValueError(f"Invalid DIFQ YAML (missing 'id' attribute): {yaml_data}")
 
-        if not re.match("^\d+\.\d+\.\d+$", str(yaml_data.get("dfiq_version", ""))):
+        if not re.match(r"^\d+\.\d+\.\d+$", str(yaml_data.get("dfiq_version", ""))):
             raise ValueError(f"Invalid DFIQ version: {yaml_data['dfiq_version']}")
 
         return yaml_data
@@ -156,27 +157,31 @@ def to_yaml(self) -> str:
         dump = self.model_dump(
             exclude={"created", "modified", "id", "root_type", "dfiq_yaml"}
         )
-        dump.pop("internal")
         dump["type"] = dump["type"].removeprefix("DFIQType.")
         dump["display_name"] = dump.pop("name")
         dump["tags"] = dump.pop("dfiq_tags")
         dump["id"] = dump.pop("dfiq_id")
+        dump["uuid"] = dump.pop("uuid")
         if dump["contributors"] is None:
             dump.pop("contributors")
         return yaml.dump(dump)
 
     def update_parents(self) -> None:
         intended_parent_ids = None
-        if hasattr(self, "parent_ids"):
+        if getattr(self, "parent_ids", []):
             intended_parent_ids = self.parent_ids
-        elif self.type == DFIQType.approach:
-            intended_parent_ids = [self.dfiq_id.split(".")[0]]
+        elif self.type == DFIQType.approach and self.parent_id:
+            intended_parent_ids = [self.parent_id]
         else:
             return
 
-        intended_parents = [
-            DFIQBase.find(dfiq_id=parent_id) for parent_id in intended_parent_ids
-        ]
+        intended_parents = []
+        for parent_id in intended_parent_ids:
+            parent = DFIQBase.find(dfiq_id=parent_id)
+            if not parent:
+                parent = DFIQBase.find(uuid=parent_id)
+            intended_parents.append(parent)
+
         if not all(intended_parents):
             raise ValueError(
                 f"Missing parent(s) {intended_parent_ids} for {self.dfiq_id}"
@@ -190,7 +195,9 @@ def update_parents(self) -> None:
                     continue
                 if rel.target != self.extended_id:
                     continue
-                if vertices[rel.source].dfiq_id not in intended_parent_ids:
+                if (
+                    vertices[rel.source].dfiq_id and vertices[rel.source].uuid
+                ) not in intended_parent_ids:
                     rel.delete()
 
         for parent in intended_parents:
@@ -209,19 +216,20 @@ def from_yaml(cls: Type["DFIQScenario"], yaml_string: str) -> "DFIQScenario":
         if yaml_data["type"] != "scenario":
             raise ValueError(f"Invalid type for DFIQ scenario: {yaml_data['type']}")
         # use re.match to check that DFIQ Ids for scenarios start with S[0-1]\d+
-        if not re.match(r"^S[0-1]\d+$", yaml_data["id"] or ""):
+        if yaml_data.get("id") and not re.match(r"^S[0-1]\d+$", yaml_data["id"] or ""):
             raise ValueError(
                 f"Invalid DFIQ ID for scenario: {yaml_data['id']}. Must be in the format S[0-1]\d+"
             )
         return cls(
             name=yaml_data["display_name"],
             description=yaml_data["description"],
+            uuid=yaml_data["uuid"],
             dfiq_id=yaml_data["id"],
             dfiq_version=yaml_data["dfiq_version"],
             dfiq_tags=yaml_data.get("tags"),
             contributors=yaml_data.get("contributors"),
             dfiq_yaml=yaml_string,
-            internal=yaml_data["id"][1] == "0",
+            internal=yaml_data.get("internal", True),
         )
 
 
@@ -237,21 +245,22 @@ def from_yaml(cls: Type["DFIQFacet"], yaml_string: str) -> "DFIQFacet":
         yaml_data = cls.parse_yaml(yaml_string)
         if yaml_data["type"] != "facet":
             raise ValueError(f"Invalid type for DFIQ facet: {yaml_data['type']}")
-        if not re.match(r"^F[0-1]\d+$", yaml_data["id"] or ""):
+        if yaml_data.get("id") and not re.match(r"^F[0-1]\d+$", yaml_data["id"] or ""):
             raise ValueError(
                 f"Invalid DFIQ ID for facet: {yaml_data['id']}. Must be in the format F[0-1]\d+"
             )
 
         return cls(
             name=yaml_data["display_name"],
             description=yaml_data.get("description"),
+            uuid=yaml_data["uuid"],
             dfiq_id=yaml_data["id"],
             dfiq_version=yaml_data["dfiq_version"],
             dfiq_tags=yaml_data.get("tags"),
             contributors=yaml_data.get("contributors"),
             parent_ids=yaml_data["parent_ids"],
             dfiq_yaml=yaml_string,
-            internal=yaml_data["id"][1] == "0",
+            internal=yaml_data.get("internal", True),
         )
 
 
@@ -267,21 +276,22 @@ def from_yaml(cls: Type["DFIQQuestion"], yaml_string: str) -> "DFIQQuestion":
         yaml_data = cls.parse_yaml(yaml_string)
         if yaml_data["type"] != "question":
             raise ValueError(f"Invalid type for DFIQ question: {yaml_data['type']}")
-        if not re.match(r"^Q[0-1]\d+$", yaml_data["id"] or ""):
+        if yaml_data.get("id") and not re.match(r"^Q[0-1]\d+$", yaml_data["id"] or ""):
             raise ValueError(
                 f"Invalid DFIQ ID for question: {yaml_data['id']}. Must be in the format Q[0-1]\d+"
             )
 
         return cls(
             name=yaml_data["display_name"],
             description=yaml_data.get("description"),
+            uuid=yaml_data["uuid"],
             dfiq_id=yaml_data["id"],
             dfiq_version=yaml_data["dfiq_version"],
             dfiq_tags=yaml_data.get("tags"),
             contributors=yaml_data.get("contributors"),
             parent_ids=yaml_data["parent_ids"],
             dfiq_yaml=yaml_string,
-            internal=yaml_data["id"][1] == "0",
+            internal=yaml_data.get("internal", True),
         )
 
 
@@ -313,8 +323,7 @@ class DFIQProcessors(BaseModel):
 
 
 class DFIQApproachDescription(BaseModel):
-    summary: str = Field(min_length=1)
-    details: str = Field(min_length=1)
+    details: str = ""
     references: list[str] = []
     references_internal: list[str] | None = None
 
@@ -336,13 +345,14 @@ class DFIQApproach(DFIQBase):
     description: DFIQApproachDescription
     view: DFIQApproachView
     type: Literal[DFIQType.approach] = DFIQType.approach
+    parent_id: str | None = None
 
     @classmethod
     def from_yaml(cls: Type["DFIQApproach"], yaml_string: str) -> "DFIQApproach":
         yaml_data = cls.parse_yaml(yaml_string)
         if yaml_data["type"] != "approach":
             raise ValueError(f"Invalid type for DFIQ approach: {yaml_data['type']}")
-        if not re.match(r"^Q[0-1]\d+\.\d+$", yaml_data["id"]):
+        if yaml_data.get("id") and not re.match(r"^Q[0-1]\d+\.\d+$", yaml_data["id"]):
             raise ValueError(
                 f"Invalid DFIQ ID for approach: {yaml_data['id']}. Must be in the format Q[0-1]\d+.\d+"
             )
@@ -355,17 +365,18 @@ def from_yaml(cls: Type["DFIQApproach"], yaml_string: str) -> "DFIQApproach":
                 f"Invalid DFIQ view for approach (has to be an object): {yaml_data['view']}"
             )
 
-        internal = bool(re.match(r"^Q[0-1]\d+\.0\d+$", yaml_data["id"]))
         return cls(
             name=yaml_data["display_name"],
             description=DFIQApproachDescription(**yaml_data["description"]),
             view=DFIQApproachView(**yaml_data["view"]),
+            uuid=yaml_data["uuid"],
             dfiq_id=yaml_data["id"],
             dfiq_version=yaml_data["dfiq_version"],
             dfiq_tags=yaml_data.get("tags"),
+            parent_id=yaml_data.get("parent_id"),
             contributors=yaml_data.get("contributors"),
             dfiq_yaml=yaml_string,
-            internal=internal,
+            internal=yaml_data.get("internal", True),
         )
 
 

core/web/apiv2/dfiq.py

@@ -56,17 +56,43 @@ class DFIQSearchResponse(BaseModel):
     total: int
 
 
+class DFIQConfigResponse(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    approach_data_sources: list[str]
+    approach_analysis_step_types: list[str]
+
+
 # API endpoints
 router = APIRouter()
 
 
+@router.get("/config")
+async def config() -> DFIQConfigResponse:
+    all_approaches = dfiq.DFIQApproach.list()
+
+    data_sources = set()
+    analysis_step_types = set()
+
+    for approach in all_approaches:
+        data_sources.update({data.type for data in approach.view.data})
+        for processor in approach.view.processors:
+            for analysis in processor.analysis:
+                analysis_step_types.update({step.type for step in analysis.steps})
+
+    return DFIQConfigResponse(
+        approach_data_sources=sorted(list(data_sources)),
+        approach_analysis_step_types=sorted(list(analysis_step_types)),
+    )
+
+
 @router.post("/from_archive")
 async def from_archive(archive: UploadFile) -> dict[str, int]:
     """Uncompresses a ZIP archive and processes the DFIQ content inside it."""
-    tempdir = tempfile.TemporaryDirectory()
-    contents = await archive.read()
-    ZipFile(BytesIO(contents)).extractall(path=tempdir.name)
-    total_added = dfiq.read_from_data_directory(tempdir.name)
+    with tempfile.TemporaryDirectory() as tempdir:
+        contents = await archive.read()
+        ZipFile(BytesIO(contents)).extractall(path=tempdir)
+        total_added = dfiq.read_from_data_directory(tempdir)
     return {"total_added": total_added}
 
 
@@ -78,13 +104,20 @@ async def new_from_yaml(request: NewDFIQRequest) -> dfiq.DFIQTypes:
     except ValueError as error:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(error))
 
-    # Ensure there is not an object with the same ID:
-    if dfiq.DFIQBase.find(dfiq_id=new.dfiq_id):
+    # Ensure there is not an object with the same ID or UUID
+
+    if new.dfiq_id and dfiq.DFIQBase.find(dfiq_id=new.dfiq_id):
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=f"DFIQ with id {new.dfiq_id} already exists",
         )
 
+    if dfiq.DFIQBase.find(uuid=new.uuid):
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"DFIQ with uuid {new.uuid} already exists",
+        )
+
     new = new.save()
 
     try:
@@ -115,19 +148,19 @@ async def to_archive(request: DFIQSearchRequest) -> FileResponse:
         aliases=request.filter_aliases,
     )
 
-    tempdir = tempfile.TemporaryDirectory()
-    for obj in dfiq_objects:
-        with open(f"{tempdir.name}/{obj.dfiq_id}.yaml", "w") as f:
-            f.write(obj.to_yaml())
+    with tempfile.TemporaryDirectory() as tempdir:
+        for obj in dfiq_objects:
+            with open(f"{tempdir}/{obj.dfiq_id}.yaml", "w") as f:
+                f.write(obj.to_yaml())
 
-    with tempfile.NamedTemporaryFile(delete=False) as archive:
-        with ZipFile(archive, "w") as zipf:
-            for obj in dfiq_objects:
-                subdir = "internal" if obj.internal else "public"
-                zipf.write(
-                    f"{tempdir.name}/{obj.dfiq_id}.yaml",
-                    f"{subdir}/{obj.type}/{obj.dfiq_id}.yaml",
-                )
+        with tempfile.NamedTemporaryFile(delete=False) as archive:
+            with ZipFile(archive, "w") as zipf:
+                for obj in dfiq_objects:
+                    subdir = "internal" if obj.internal else "public"
+                    zipf.write(
+                        f"{tempdir}/{obj.dfiq_id}.yaml",
+                        f"{subdir}/{obj.type}/{obj.dfiq_id}.yaml",
+                    )
 
     return FileResponse(archive.name, media_type="application/zip", filename="dfiq.zip")
 
@@ -144,7 +177,7 @@ async def validate_dfiq_yaml(request: DFIQValidateRequest) -> DFIQValidateRespon
     except KeyError as error:
         return DFIQValidateResponse(valid=False, error=f"Invalid DFIQ type: {error}")
 
-    if request.check_id and dfiq.DFIQBase.find(dfiq_id=obj.dfiq_id):
+    if request.check_id and obj.dfiq_id and dfiq.DFIQBase.find(dfiq_id=obj.dfiq_id):
         return DFIQValidateResponse(
             valid=False, error=f"DFIQ with id {obj.dfiq_id} already exists"
         )

plugins/feeds/public/artifacts.py

@@ -30,11 +30,11 @@ def run(self):
             logging.info("No response: skipping ForensicArtifact update")
             return
 
-        tempdir = tempfile.TemporaryDirectory()
-        ZipFile(BytesIO(response.content)).extractall(path=tempdir.name)
-        artifacts_datadir = os.path.join(
-            tempdir.name, "artifacts-main", "artifacts", "data"
-        )
+        with tempfile.TemporaryDirectory() as tempdir:
+            ZipFile(BytesIO(response.content)).extractall(path=tempdir)
+            artifacts_datadir = os.path.join(
+                tempdir, "artifacts-main", "artifacts", "data"
+            )
 
         data_files_glob = glob.glob(os.path.join(artifacts_datadir, "*.yaml"))
         artifacts_dict = {}

plugins/feeds/public/attack.py

@@ -255,6 +255,7 @@ def run(self):
                         )
                         rel_count += 1
         logging.info("Processed %s relationships", rel_count)
+        tempdir.cleanup()
 
 
 taskmanager.TaskManager.register_task(MitreAttack)

plugins/feeds/public/dfiq.py

@@ -19,15 +19,15 @@ class DFIQFeed(task.FeedTask):
 
     def run(self):
         response = self._make_request(
-            "https://github.com/google/dfiq/archive/refs/heads/main.zip"
+            "https://github.com/tomchop/dfiq/archive/refs/heads/dfiq1.1.zip"
         )
         if not response:
             logging.info("No response: skipping DFIQ update")
             return
 
-        tempdir = tempfile.TemporaryDirectory()
-        ZipFile(BytesIO(response.content)).extractall(path=tempdir.name)
-        dfiq.read_from_data_directory(tempdir.name)
+        with tempfile.TemporaryDirectory() as tempdir:
+            ZipFile(BytesIO(response.content)).extractall(path=tempdir)
+            dfiq.read_from_data_directory(tempdir)
 
         extra_dirs = yeti_config.get("dfiq", "extra_dirs")
         if not extra_dirs: