From f653520a73f018fcf253b0429e5d3fd7f312acb0 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 27 Jun 2024 14:44:00 +0000 Subject: [PATCH 1/7] Move from ID to UUID, make ID optional --- core/database_arango.py | 4 +- core/schemas/dfiq.py | 50 +++++--- core/web/apiv2/dfiq.py | 13 +- tests/apiv2/dfiq.py | 112 +++++++++++++----- tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml | 13 ++ tests/dfiq_test_data/F1005.yaml | 4 +- tests/dfiq_test_data/Q1020.10.yaml | 5 +- .../Q1020.10_no_indicators.yaml | 7 +- tests/dfiq_test_data/Q1020.10_no_parent.yaml | 61 ++++++++++ tests/dfiq_test_data/Q1020.yaml | 4 +- tests/dfiq_test_data/Q1020_no_parents.yaml | 11 ++ tests/dfiq_test_data/Q1020_uuid_parent.yaml | 12 ++ .../Q1020_uuid_scenario_parent.yaml | 12 ++ tests/dfiq_test_data/S1003.yaml | 4 +- tests/schemas/dfiq.py | 25 +++- 15 files changed, 272 insertions(+), 65 deletions(-) create mode 100644 tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml create mode 100644 tests/dfiq_test_data/Q1020.10_no_parent.yaml create mode 100644 tests/dfiq_test_data/Q1020_no_parents.yaml create mode 100644 tests/dfiq_test_data/Q1020_uuid_parent.yaml create mode 100644 tests/dfiq_test_data/Q1020_uuid_scenario_parent.yaml diff --git a/core/database_arango.py b/core/database_arango.py index e00765606..9f44d62b9 100644 --- a/core/database_arango.py +++ b/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: diff --git a/core/schemas/dfiq.py b/core/schemas/dfiq.py index 00078628d..532aa60b7 100644 --- a/core/schemas/dfiq.py +++ b/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,7 +245,7 @@ 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+" ) @@ -245,13 +253,14 @@ def from_yaml(cls: Type["DFIQFacet"], yaml_string: str) -> "DFIQFacet": 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,7 +276,7 @@ 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+" ) @@ -275,13 +284,14 @@ def from_yaml(cls: Type["DFIQQuestion"], yaml_string: str) -> "DFIQQuestion": 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), ) @@ -336,13 +346,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 +366,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), ) diff --git a/core/web/apiv2/dfiq.py b/core/web/apiv2/dfiq.py index 444ec47f5..d69fddbaf 100644 --- a/core/web/apiv2/dfiq.py +++ b/core/web/apiv2/dfiq.py @@ -77,13 +77,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: @@ -109,7 +116,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" ) diff --git a/tests/apiv2/dfiq.py b/tests/apiv2/dfiq.py index 57078bb45..b0c34c5c6 100644 --- a/tests/apiv2/dfiq.py +++ b/tests/apiv2/dfiq.py @@ -40,7 +40,7 @@ def test_new_dfiq_scenario(self) -> None: self.assertIsNotNone(data["created"]) self.assertEqual(data["name"], "scenario1") self.assertEqual(data["dfiq_id"], "S1003") - self.assertEqual(data["dfiq_version"], "1.0.0") + self.assertEqual(data["dfiq_version"], "1.1.0") self.assertEqual(data["description"], "Long description 1\n") self.assertEqual(data["type"], dfiq.DFIQType.scenario) self.assertEqual(data["dfiq_tags"], ["Tag1", "Tag2", "Tag3"]) @@ -49,7 +49,8 @@ def test_new_dfiq_facet(self) -> None: scenario = dfiq.DFIQScenario( name="mock_scenario", dfiq_id="S1003", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() @@ -70,7 +71,7 @@ def test_new_dfiq_facet(self) -> None: self.assertIsNotNone(data["created"]) self.assertEqual(data["name"], "facet1") self.assertEqual(data["dfiq_id"], "F1005") - self.assertEqual(data["dfiq_version"], "1.0.0") + self.assertEqual(data["dfiq_version"], "1.1.0") self.assertEqual(data["description"], "Long description of facet1\n") self.assertEqual(data["type"], dfiq.DFIQType.facet) self.assertEqual(data["dfiq_tags"], ["Web Browser"]) @@ -87,7 +88,8 @@ def test_new_dfiq_question(self) -> None: facet = dfiq.DFIQFacet( name="mock_facet", dfiq_id="F1005", - dfiq_version="1.0.0", + uuid="fake_facet_uuid", + dfiq_version="1.1.0", description="desc", parent_ids=["S1003"], dfiq_yaml="mock", @@ -109,7 +111,7 @@ def test_new_dfiq_question(self) -> None: self.assertIsNotNone(data["created"]) self.assertEqual(data["name"], "What is a question?") self.assertEqual(data["dfiq_id"], "Q1020") - self.assertEqual(data["dfiq_version"], "1.0.0") + self.assertEqual(data["dfiq_version"], "1.1.0") self.assertEqual(data["description"], None) self.assertEqual(data["type"], dfiq.DFIQType.question) self.assertEqual(data["dfiq_tags"], ["Web Browser"]) @@ -126,7 +128,8 @@ def test_new_dfiq_approach(self) -> None: question = dfiq.DFIQQuestion( name="mock_question", dfiq_id="Q1020", - dfiq_version="1.0.0", + uuid="fake_question_uuid", + dfiq_version="1.1.0", description="desc", parent_ids=["F1005"], dfiq_yaml="mock", @@ -148,7 +151,7 @@ def test_new_dfiq_approach(self) -> None: self.assertIsNotNone(data["created"]) self.assertEqual(data["name"], "Approach1") self.assertEqual(data["dfiq_id"], "Q1020.10") - self.assertEqual(data["dfiq_version"], "1.0.0") + self.assertEqual(data["dfiq_version"], "1.1.0") self.assertEqual(data["description"]["summary"], "Description for approach") self.assertEqual(data["type"], dfiq.DFIQType.approach) self.assertEqual(data["dfiq_tags"], ["Lots", "Of", "Tags"]) @@ -164,7 +167,8 @@ def test_dfiq_patch_updates_parents(self) -> None: scenario1 = dfiq.DFIQScenario( name="mock_scenario", dfiq_id="S1003", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid1", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() @@ -172,7 +176,8 @@ def test_dfiq_patch_updates_parents(self) -> None: scenario2 = dfiq.DFIQScenario( name="mock_scenario2", dfiq_id="S1222", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid2", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() @@ -180,7 +185,8 @@ def test_dfiq_patch_updates_parents(self) -> None: facet = dfiq.DFIQFacet( name="mock_facet", dfiq_id="F1005", - dfiq_version="1.0.0", + uuid="fake_facet_uuid", + dfiq_version="1.1.0", description="desc", parent_ids=["S1003"], dfiq_yaml="mock", @@ -215,7 +221,8 @@ def test_dfiq_patch_approach_updates_parents(self) -> None: dfiq.DFIQScenario( name="mock_scenario", dfiq_id="S1003", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() @@ -223,25 +230,28 @@ def test_dfiq_patch_approach_updates_parents(self) -> None: dfiq.DFIQFacet( name="mock_facet", dfiq_id="F1005", - dfiq_version="1.0.0", + uuid="fake_facet_uuid", + dfiq_version="1.1.0", description="desc", - parent_ids=["S1003"], + parent_ids=["fake_scenario_uuid"], dfiq_yaml="mock", ).save() question1 = dfiq.DFIQQuestion( name="mock_question", dfiq_id="Q1020", - dfiq_version="1.0.0", + uuid="bd46ce6e-c933-46e5-960c-36945aaef401", + dfiq_version="1.1.0", description="desc", - parent_ids=["F1005"], + parent_ids=["fake_facet_uuid"], dfiq_yaml="mock", ).save() question2 = dfiq.DFIQQuestion( name="mock_question2", + uuid="fake_question_uuid_2", dfiq_id="Q1022", - dfiq_version="1.0.0", + dfiq_version="1.1.0", description="desc", parent_ids=["F1005"], dfiq_yaml="mock", @@ -259,7 +269,7 @@ def test_dfiq_patch_approach_updates_parents(self) -> None: self.assertEqual(edges[0][0].description, "Uses DFIQ approach") self.assertEqual(total, 1) - approach.dfiq_id = "Q1022.10" + approach.parent_id = "fake_question_uuid_2" response = client.patch( f"/api/v2/dfiq/{approach.id}", json={ @@ -270,12 +280,14 @@ def test_dfiq_patch_approach_updates_parents(self) -> None: ) data = response.json() self.assertEqual(response.status_code, 200, data) - self.assertEqual(data["dfiq_id"], "Q1022.10") + self.assertEqual(data["dfiq_id"], "Q1020.10") self.assertEqual(data["id"], approach.id) vertices, edges, total = approach.neighbors() self.assertEqual(len(vertices), 1) self.assertEqual(vertices[f"dfiq/{question2.id}"].dfiq_id, "Q1022") + self.assertEqual(vertices[f"dfiq/{question2.id}"].uuid, "fake_question_uuid_2") + self.assertEqual(edges[0][0].type, "approach") self.assertEqual(edges[0][0].description, "Uses DFIQ approach") self.assertEqual(total, 1) @@ -283,16 +295,16 @@ def test_dfiq_patch_approach_updates_parents(self) -> None: def test_dfiq_patch_approach_updates_indicators(self) -> None: dfiq.DFIQScenario( name="mock_scenario", - dfiq_id="S1003", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() dfiq.DFIQFacet( name="mock_facet", - dfiq_id="F1005", - dfiq_version="1.0.0", + uuid="fake_facet_uuid", + dfiq_version="1.1.0", description="desc", parent_ids=["S1003"], dfiq_yaml="mock", @@ -300,8 +312,8 @@ def test_dfiq_patch_approach_updates_indicators(self) -> None: dfiq.DFIQQuestion( name="mock_question", - dfiq_id="Q1020", - dfiq_version="1.0.0", + uuid="bd46ce6e-c933-46e5-960c-36945aaef401", + dfiq_version="1.1.0", description="desc", parent_ids=["F1005"], dfiq_yaml="mock", @@ -350,16 +362,16 @@ def test_dfiq_patch_approach_updates_indicators(self) -> None: def test_dfiq_post_approach(self): dfiq.DFIQScenario( name="mock_scenario", - dfiq_id="S1003", - dfiq_version="1.0.0", + uuid="fake_scenario_uuid", + dfiq_version="1.1.0", description="desc", dfiq_yaml="mock", ).save() dfiq.DFIQFacet( name="mock_facet", - dfiq_id="F1005", - dfiq_version="1.0.0", + uuid="fake_facet_uuid", + dfiq_version="1.1.0", description="desc", parent_ids=["S1003"], dfiq_yaml="mock", @@ -367,8 +379,8 @@ def test_dfiq_post_approach(self): dfiq.DFIQQuestion( name="mock_question", - dfiq_id="Q1020", - dfiq_version="1.0.0", + uuid="bd46ce6e-c933-46e5-960c-36945aaef401", + dfiq_version="1.1.0", description="desc", parent_ids=["F1005"], dfiq_yaml="mock", @@ -393,6 +405,7 @@ def test_dfiq_post_approach(self): self.assertEqual(len(vertices), 1) approach.delete() + # Repeat the action, updating indicators response = client.post( "/api/v2/dfiq/from_yaml", json={ @@ -436,7 +449,12 @@ def test_wrong_parent_approach(self) -> None: ) data = response.json() self.assertEqual(response.status_code, 400, data) - self.assertEqual(data, {"detail": "Missing parent(s) ['Q1020'] for Q1020.10"}) + self.assertEqual( + data, + { + "detail": "Missing parent(s) ['bd46ce6e-c933-46e5-960c-36945aaef401'] for Q1020.10" + }, + ) def test_valid_dfiq_yaml(self) -> None: with open("tests/dfiq_test_data/S1003.yaml", "r") as f: @@ -499,6 +517,38 @@ def test_valid_dfiq_yaml(self) -> None: self.assertEqual(response.status_code, 200, data) self.assertEqual(data["valid"], True) + def test_standalone_question_creation(self): + with open("tests/dfiq_test_data/Q1020_no_parents.yaml", "r") as f: + yaml_string = f.read() + + response = client.post( + "/api/v2/dfiq/from_yaml", + json={ + "dfiq_yaml": yaml_string, + "dfiq_type": dfiq.DFIQType.question, + }, + ) + data = response.json() + self.assertEqual(response.status_code, 200, data) + self.assertIsNotNone(data["id"]) + self.assertEquals(data["parent_ids"], []) + + def test_standalone_approach_creation(self): + with open("tests/dfiq_test_data/Q1020.10_no_parent.yaml", "r") as f: + yaml_string = f.read() + + response = client.post( + "/api/v2/dfiq/from_yaml", + json={ + "dfiq_yaml": yaml_string, + "dfiq_type": dfiq.DFIQType.approach, + }, + ) + data = response.json() + self.assertEqual(response.status_code, 200, data) + self.assertIsNotNone(data["id"]) + self.assertIsNone(data["parent_id"]) + def test_upload_dfiq_archive(self): zip_archive = open("tests/dfiq_test_data/dfiq_test_data.zip", "rb") response = client.post( diff --git a/tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml b/tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml new file mode 100644 index 000000000..9b3cb8ef4 --- /dev/null +++ b/tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml @@ -0,0 +1,13 @@ +--- +display_name: scenario1 +type: scenario +description: > + Long description 1 +id: +uuid: 2ee16263-56f8-49a5-9b33-d1a2dd8b829c +internal: false +dfiq_version: 1.1.0 +tags: + - Tag1 + - Tag2 + - Tag3 diff --git a/tests/dfiq_test_data/F1005.yaml b/tests/dfiq_test_data/F1005.yaml index 4ea3a0f71..ee5de55f6 100644 --- a/tests/dfiq_test_data/F1005.yaml +++ b/tests/dfiq_test_data/F1005.yaml @@ -4,7 +4,9 @@ type: facet description: > Long description of facet1 id: F1005 -dfiq_version: 1.0.0 +uuid: b2bab31f-1670-4297-8cb1-685747a13468 +internal: false +dfiq_version: 1.1.0 tags: - Web Browser parent_ids: diff --git a/tests/dfiq_test_data/Q1020.10.yaml b/tests/dfiq_test_data/Q1020.10.yaml index 530022870..1553108a7 100644 --- a/tests/dfiq_test_data/Q1020.10.yaml +++ b/tests/dfiq_test_data/Q1020.10.yaml @@ -2,7 +2,10 @@ display_name: Approach1 type: approach id: Q1020.10 -dfiq_version: 1.0.0 +uuid: 292500f7-9d54-40ca-8254-34821e9b5c4e +parent_id: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 tags: - Lots - Of diff --git a/tests/dfiq_test_data/Q1020.10_no_indicators.yaml b/tests/dfiq_test_data/Q1020.10_no_indicators.yaml index 293f05efe..edd37f8be 100644 --- a/tests/dfiq_test_data/Q1020.10_no_indicators.yaml +++ b/tests/dfiq_test_data/Q1020.10_no_indicators.yaml @@ -2,13 +2,16 @@ display_name: Approach1 type: approach id: Q1020.10 -dfiq_version: 1.0.0 +uuid: fcbdb313-424a-436e-a877-130aeba3f134 +parent_id: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 tags: - Lots - Of - Tags description: - summary: Description for approach + summary: Descripion for approach details: > Details for approach references: diff --git a/tests/dfiq_test_data/Q1020.10_no_parent.yaml b/tests/dfiq_test_data/Q1020.10_no_parent.yaml new file mode 100644 index 000000000..66f3dbcf0 --- /dev/null +++ b/tests/dfiq_test_data/Q1020.10_no_parent.yaml @@ -0,0 +1,61 @@ +--- +display_name: Approach1 +type: approach +id: Q1020.10 +uuid: 292500f7-9d54-40ca-8254-34821e9b5c4e +parent_id: +internal: false +dfiq_version: 1.1.0 +tags: + - Lots + - Of + - Tags +description: + summary: Description for approach + details: > + Details for approach + references: + - "ref1" + - "ref2" + references_internal: null +view: + data: + - type: artifact + value: RandomArtifact + - type: description + value: Random description + notes: + covered: + - Covered1 + - Covered2 + - Covered3 + not_covered: + - Not covered1 + - Not covered2 + processors: + - name: processor1 + options: + - type: parsers + value: parser1option + analysis: + - name: OpenSearch + steps: + - description: random parser description + type: opensearch-query + value: data_type:("fs:stat") + - name: Python Notebook + steps: + - description: random step description + type: pandas + value: query('data_type in ("fs:stat")') + - name: processor2 + options: + - type: format + value: jsonl + analysis: + - name: analysis1 + steps: + - description: &filter-desc-processor2 > + something else + type: opensearch-query + value: data_type:"chrome:history:page_visited") diff --git a/tests/dfiq_test_data/Q1020.yaml b/tests/dfiq_test_data/Q1020.yaml index b92ded8b9..5878cc4f4 100644 --- a/tests/dfiq_test_data/Q1020.yaml +++ b/tests/dfiq_test_data/Q1020.yaml @@ -3,7 +3,9 @@ display_name: What is a question? type: question description: id: Q1020 -dfiq_version: 1.0.0 +uuid: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 tags: - Web Browser parent_ids: diff --git a/tests/dfiq_test_data/Q1020_no_parents.yaml b/tests/dfiq_test_data/Q1020_no_parents.yaml new file mode 100644 index 000000000..df8c2d5ca --- /dev/null +++ b/tests/dfiq_test_data/Q1020_no_parents.yaml @@ -0,0 +1,11 @@ +--- +display_name: What is a question? +type: question +description: +id: Q1020 +uuid: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 +tags: + - Web Browser +parent_ids: [] diff --git a/tests/dfiq_test_data/Q1020_uuid_parent.yaml b/tests/dfiq_test_data/Q1020_uuid_parent.yaml new file mode 100644 index 000000000..23a576f74 --- /dev/null +++ b/tests/dfiq_test_data/Q1020_uuid_parent.yaml @@ -0,0 +1,12 @@ +--- +display_name: What is a question? +type: question +description: +id: Q1020 +uuid: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 +tags: + - Web Browser +parent_ids: + - b2bab31f-1670-4297-8cb1-685747a13468 diff --git a/tests/dfiq_test_data/Q1020_uuid_scenario_parent.yaml b/tests/dfiq_test_data/Q1020_uuid_scenario_parent.yaml new file mode 100644 index 000000000..baf0371bf --- /dev/null +++ b/tests/dfiq_test_data/Q1020_uuid_scenario_parent.yaml @@ -0,0 +1,12 @@ +--- +display_name: What is a question? +type: question +description: +id: Q1020 +uuid: bd46ce6e-c933-46e5-960c-36945aaef401 +internal: false +dfiq_version: 1.1.0 +tags: + - Web Browser +parent_ids: + - S1003 diff --git a/tests/dfiq_test_data/S1003.yaml b/tests/dfiq_test_data/S1003.yaml index 6cac4c2b3..155690ecd 100644 --- a/tests/dfiq_test_data/S1003.yaml +++ b/tests/dfiq_test_data/S1003.yaml @@ -4,7 +4,9 @@ type: scenario description: > Long description 1 id: S1003 -dfiq_version: 1.0.0 +uuid: 2ee16263-56f8-49a5-9b33-d1a2dd8b829c +internal: false +dfiq_version: 1.1.0 tags: - Tag1 - Tag2 diff --git a/tests/schemas/dfiq.py b/tests/schemas/dfiq.py index dce5a0fba..736c0a341 100644 --- a/tests/schemas/dfiq.py +++ b/tests/schemas/dfiq.py @@ -28,7 +28,22 @@ def test_dfiq_scenario(self) -> None: self.assertIsNotNone(result.id) self.assertIsNotNone(result.created) self.assertEqual(result.name, "scenario1") - self.assertEqual(result.dfiq_version, "1.0.0") + self.assertEqual(result.dfiq_version, "1.1.0") + self.assertEqual(str(result.uuid), "2ee16263-56f8-49a5-9b33-d1a2dd8b829c") + self.assertEqual(result.description, "Long description 1\n") + self.assertEqual(result.type, DFIQType.scenario) + self.assertEqual(result.dfiq_tags, ["Tag1", "Tag2", "Tag3"]) + + def test_dfiq_scenario_no_id(self) -> None: + with open("tests/dfiq_test_data/DFIQ_Scenario_no_id.yaml", "r") as f: + yaml_string = f.read() + + result = DFIQScenario.from_yaml(yaml_string).save() + self.assertIsNotNone(result.id) + self.assertIsNotNone(result.created) + self.assertEqual(result.name, "scenario1") + self.assertEqual(result.dfiq_version, "1.1.0") + self.assertEqual(str(result.uuid), "2ee16263-56f8-49a5-9b33-d1a2dd8b829c") self.assertEqual(result.description, "Long description 1\n") self.assertEqual(result.type, DFIQType.scenario) self.assertEqual(result.dfiq_tags, ["Tag1", "Tag2", "Tag3"]) @@ -44,7 +59,8 @@ def test_dfiq_facet(self) -> None: self.assertEqual(result.name, "facet1") self.assertEqual(result.description, "Long description of facet1\n") self.assertEqual(result.dfiq_id, "F1005") - self.assertEqual(result.dfiq_version, "1.0.0") + self.assertEqual(result.dfiq_version, "1.1.0") + self.assertEqual(str(result.uuid), "b2bab31f-1670-4297-8cb1-685747a13468") self.assertEqual(result.dfiq_tags, ["Web Browser"]) self.assertEqual(result.parent_ids, ["S1003"]) self.assertEqual(result.type, DFIQType.facet) @@ -59,8 +75,9 @@ def test_dfiq_question(self) -> None: self.assertIsNotNone(result.created) self.assertEqual(result.name, "What is a question?") self.assertEqual(result.description, None) + self.assertEqual(str(result.uuid), "bd46ce6e-c933-46e5-960c-36945aaef401") self.assertEqual(result.dfiq_id, "Q1020") - self.assertEqual(result.dfiq_version, "1.0.0") + self.assertEqual(result.dfiq_version, "1.1.0") self.assertEqual(result.dfiq_tags, ["Web Browser"]) self.assertEqual(result.parent_ids, ["F1005"]) self.assertEqual(result.type, DFIQType.question) @@ -72,6 +89,8 @@ def test_dfiq_approach(self) -> None: result = DFIQApproach.from_yaml(yaml_string).save() self.assertIsNotNone(result.id) + self.assertEquals(result.uuid, "292500f7-9d54-40ca-8254-34821e9b5c4e") + self.assertEquals(result.parent_id, "bd46ce6e-c933-46e5-960c-36945aaef401") self.assertIsNotNone(result.created) self.assertEqual(result.name, "Approach1") self.assertEqual(result.description.summary, "Description for approach") From c69b6b545176087ad3b7189c0bdc65743ddf3a5a Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 27 Jun 2024 15:10:52 +0000 Subject: [PATCH 2/7] Remove approach summaries --- core/schemas/dfiq.py | 3 +-- tests/dfiq_test_data/Q1020.10.yaml | 1 - .../Q1020.10_no_indicators.yaml | 1 - tests/dfiq_test_data/Q1020.10_no_parent.yaml | 1 - tests/dfiq_test_data/dfiq_test_data.zip | Bin 1697 -> 1586 bytes tests/schemas/dfiq.py | 5 ++--- 6 files changed, 3 insertions(+), 8 deletions(-) diff --git a/core/schemas/dfiq.py b/core/schemas/dfiq.py index 532aa60b7..a4ebcb9c1 100644 --- a/core/schemas/dfiq.py +++ b/core/schemas/dfiq.py @@ -323,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 diff --git a/tests/dfiq_test_data/Q1020.10.yaml b/tests/dfiq_test_data/Q1020.10.yaml index 1553108a7..c3612db80 100644 --- a/tests/dfiq_test_data/Q1020.10.yaml +++ b/tests/dfiq_test_data/Q1020.10.yaml @@ -11,7 +11,6 @@ tags: - Of - Tags description: - summary: Description for approach details: > Details for approach references: diff --git a/tests/dfiq_test_data/Q1020.10_no_indicators.yaml b/tests/dfiq_test_data/Q1020.10_no_indicators.yaml index edd37f8be..3d0b53f70 100644 --- a/tests/dfiq_test_data/Q1020.10_no_indicators.yaml +++ b/tests/dfiq_test_data/Q1020.10_no_indicators.yaml @@ -11,7 +11,6 @@ tags: - Of - Tags description: - summary: Descripion for approach details: > Details for approach references: diff --git a/tests/dfiq_test_data/Q1020.10_no_parent.yaml b/tests/dfiq_test_data/Q1020.10_no_parent.yaml index 66f3dbcf0..8a401fa38 100644 --- a/tests/dfiq_test_data/Q1020.10_no_parent.yaml +++ b/tests/dfiq_test_data/Q1020.10_no_parent.yaml @@ -11,7 +11,6 @@ tags: - Of - Tags description: - summary: Description for approach details: > Details for approach references: diff --git a/tests/dfiq_test_data/dfiq_test_data.zip b/tests/dfiq_test_data/dfiq_test_data.zip index e70e59183c6724de8c6acee48d77aaad5b892843..bbe302982b842de565ded14504c3cedbf9fa27b1 100644 GIT binary patch literal 1586 zcmWIWW@Zs#U|`^2SW$S|a&WfsPlq!yRNrzDmn>bn^l7?|o+ zCg$dZhHx@4FRrgmV``{PE3M#WU}OPmVE_}peH-}>8SuE2*SFRygj#J#Nk~XYO7HP` z6f#YL(~s-b{l!6A8}Ht|JolUgb70io%a4q0ZzL`{P}@>D?T2yQ)eF;i1$W$O@_HMz z>9DdwoY%Y$`bQ^yTUh$cCaSt&#*>%vYu-jnvs*{JFfnD%_;AYHCbQB!Y{il}yQlA1 zC}uI)>t}!0u2p}4D<{Q5FTp)dJ2oj)`s}@ z&o&U)`&rvfX4h_ZC69c&kf_rB(CzW%PYwy}VV(VTM#!Jvd~Cwg6kON7@;1ufH}8Ir z=Oh&^(b8U-BHd__4>wA$omg~E$-E@z^v%8y%?H`j=+#OR_u;u$wg~@7>wEDMGPW=R&#s@1s1aA5U2v%YRg~R5nUinsHbD ztQTuOe4i2RrS9Atss1`dmHWA4!{(eSxBZ^6v%VF+<>p)+{qAYbpIbW)>@v_%pO$pa z>%n7|l;{ULSnpLgcs~s{wUZX-O0Rwp(OsJ2dE$oazRXlHn-r<@vz)%3o?GeK5uSK( zy0w4bnKe%*==^qgn`EZpJp5pc~g8z(m%$h#!>yi1#{`?ZZmi)l?oa}M0Tj{U#3+uNF zuq{nI`tx04{j%fH#!tfgZouwJqx=E}6yT~QN_|HQs}!uff#($fu3ZvJOPO){2E!a3Ez1icu0Mt~$4ct*gH zUQBxo`I-%QT+8Pl66fHO-5m0Y*UHdM&v5Nhvz8kHg}=?5a^AnTt$i+|a`DqL+q7Q= z2Ohe%E%Mm0W{q0$zNWWN^Uv{imaYksVxJ)L;ok4=_R`#%u#LVRlXVIouUdO_-{-mP zi?SlSs5wTdUdG2Sy8PfPz?0ReC)mn24-txf~nuP$B^r= zfdI?z>Z7#|ZsO%ypPEWM_4&NT*gKcJNNKK@m&qwGwdUS^OoC(4*R!+F|LNo_?V34T zq-o#PZs!gw)0MGaWfkm4+)r}q>k4@k{7}hukC>gQ}hwK5~j7)OOxQZtUpl=u$fcctXNh64br?6s$6jo>j6|(V|Z;NfKmYg8&H!KvJITf7#=`QWW*Vd8qdhaU&k~aIo5F*&&mc0SQa3x2Zr-paPbZR DL54ne literal 1697 zcmWIWW@h1H00EN~p%GvPl;CHOVMs~KEQ~KnEiQ>qNi0d!4-MgDVBVZ^Emdl}Pg-dO zHv=OJNHLfQ02(9$t>7C{aWq0d-F~p<#g=6LFNv9bZ7&KWKfIh(plt4oRBLh7{1A^XJ8{*wR z+dyF7XKgo`UAx(pJo4>AqDuQix39OJbx2?j>+G)hrpzn)3Y zR`FaERD3;QzI2Y~!Q|_f-L7+8_q{)I>||fb$_Lrg-d&EJ=9aoN{z?GH}*xns<2qbFDvr1pGsU#Qn-3vd$G=mxYEHJ#W$X3YIyDT8 zem7)T|F$jL^7ip9bM-aL?ws2BLHc$2wN_Efl_9}lO!DpuyIHw#P1@i^&U=hL5LE#+)iJNbHIOue9qpuC=woY1?+U82=r zJ{}f*sK0mp(@S0RzFfW)|Nl@G$DWfrCr3QVU6$DKdv~ir2DA3La{k-ackka?xo4x; z->+WrH)cJ1rj_phEqvCV)g{XBjU*2|@hPqP`}*)U)%fgUhO8O?PA_`As4OV->a4v} z-v_m`pYorzb$L+Pe|FRa67`VXya<@kyRavANCF|0*iAdF`3@QIIF!#nB<}H4VPlG8 zQbPJG3CW}bLQ67p9QN9VbiA!~;t05MYwq-zpNgEvrL9*fKK>(8Jogn}P3Ha2}}(18TxVw!w@UhYgTqN}Tbi fNe0>Y-wLv$BC=ixmj307GCqxLyGO0^U?m diff --git a/tests/schemas/dfiq.py b/tests/schemas/dfiq.py index 736c0a341..646d44fa7 100644 --- a/tests/schemas/dfiq.py +++ b/tests/schemas/dfiq.py @@ -89,11 +89,10 @@ def test_dfiq_approach(self) -> None: result = DFIQApproach.from_yaml(yaml_string).save() self.assertIsNotNone(result.id) - self.assertEquals(result.uuid, "292500f7-9d54-40ca-8254-34821e9b5c4e") - self.assertEquals(result.parent_id, "bd46ce6e-c933-46e5-960c-36945aaef401") + self.assertEqual(result.uuid, "292500f7-9d54-40ca-8254-34821e9b5c4e") + self.assertEqual(result.parent_id, "bd46ce6e-c933-46e5-960c-36945aaef401") self.assertIsNotNone(result.created) self.assertEqual(result.name, "Approach1") - self.assertEqual(result.description.summary, "Description for approach") self.assertEqual(result.description.details, "Details for approach\n") self.assertEqual(result.description.references, ["ref1", "ref2"]) From 4a40f1eb4466afd6ea7eab3937819853b4cc90d2 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 27 Jun 2024 15:11:08 +0000 Subject: [PATCH 3/7] Change deprecated test function --- tests/apiv2/dfiq.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/apiv2/dfiq.py b/tests/apiv2/dfiq.py index b0c34c5c6..da2f388c3 100644 --- a/tests/apiv2/dfiq.py +++ b/tests/apiv2/dfiq.py @@ -531,7 +531,7 @@ def test_standalone_question_creation(self): data = response.json() self.assertEqual(response.status_code, 200, data) self.assertIsNotNone(data["id"]) - self.assertEquals(data["parent_ids"], []) + self.assertEqual(data["parent_ids"], []) def test_standalone_approach_creation(self): with open("tests/dfiq_test_data/Q1020.10_no_parent.yaml", "r") as f: From 73decd3bd3c247f29dd1dff79b79601c534f23df Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 16 Aug 2024 09:05:23 +0000 Subject: [PATCH 4/7] New Config endpoint for DFIQ to compute data dynamically --- core/web/apiv2/dfiq.py | 26 ++++++++++++++++++++++++++ tests/apiv2/dfiq.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/core/web/apiv2/dfiq.py b/core/web/apiv2/dfiq.py index 742515d53..c3d484903 100644 --- a/core/web/apiv2/dfiq.py +++ b/core/web/apiv2/dfiq.py @@ -56,10 +56,36 @@ 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.""" diff --git a/tests/apiv2/dfiq.py b/tests/apiv2/dfiq.py index 65dc3da1f..ebea8a800 100644 --- a/tests/apiv2/dfiq.py +++ b/tests/apiv2/dfiq.py @@ -25,6 +25,38 @@ def setUp(self) -> None: ).json() client.headers = {"Authorization": "Bearer " + token_data["access_token"]} + def test_config(self) -> None: + dfiq.DFIQQuestion( + name="mock_question", + dfiq_id="Q1020", + uuid="bd46ce6e-c933-46e5-960c-36945aaef401", + dfiq_version="1.1.0", + description="desc", + parent_ids=["F1005"], + dfiq_yaml="mock", + ).save() + + with open("tests/dfiq_test_data/Q1020.10.yaml", "r") as f: + yaml_string = f.read() + + response = client.post( + "/api/v2/dfiq/from_yaml", + json={ + "dfiq_yaml": yaml_string, + "dfiq_type": dfiq.DFIQType.approach, + }, + ) + self.assertEqual(response.status_code, 200, response.json()) + + response = client.get("/api/v2/dfiq/config") + data = response.json() + + self.assertEqual(response.status_code, 200, data) + self.assertEqual(data["approach_data_sources"], ["artifact", "description"]) + self.assertEqual( + data["approach_analysis_step_types"], ["opensearch-query", "pandas"] + ) + def test_new_dfiq_scenario(self) -> None: with open("tests/dfiq_test_data/S1003.yaml", "r") as f: yaml_string = f.read() From de35d3c1f91d89dab05dd0adfd4154bef816b070 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 16 Aug 2024 09:05:33 +0000 Subject: [PATCH 5/7] update feed to tomchop's 1.1 for while changes are merged --- plugins/feeds/public/dfiq.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/feeds/public/dfiq.py b/plugins/feeds/public/dfiq.py index 61c5c7a00..95caddae1 100644 --- a/plugins/feeds/public/dfiq.py +++ b/plugins/feeds/public/dfiq.py @@ -19,7 +19,7 @@ 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") From 37a160f662e7c74f54a89020f23474bf74715c36 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 16 Aug 2024 11:34:19 +0000 Subject: [PATCH 6/7] Fix some leaky handles --- core/web/apiv2/dfiq.py | 34 +++++++++++++++---------------- plugins/feeds/public/artifacts.py | 10 ++++----- plugins/feeds/public/attack.py | 10 ++++----- plugins/feeds/public/dfiq.py | 6 +++--- tests/apiv2/dfiq.py | 10 ++++----- 5 files changed, 35 insertions(+), 35 deletions(-) diff --git a/core/web/apiv2/dfiq.py b/core/web/apiv2/dfiq.py index c3d484903..4d15bef32 100644 --- a/core/web/apiv2/dfiq.py +++ b/core/web/apiv2/dfiq.py @@ -89,10 +89,10 @@ async def config() -> DFIQConfigResponse: @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} @@ -148,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.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.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}/{obj.dfiq_id}.yaml", + f"{subdir}/{obj.type}/{obj.dfiq_id}.yaml", + ) return FileResponse(archive.name, media_type="application/zip", filename="dfiq.zip") diff --git a/plugins/feeds/public/artifacts.py b/plugins/feeds/public/artifacts.py index c730e864c..7dffde670 100644 --- a/plugins/feeds/public/artifacts.py +++ b/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 = {} diff --git a/plugins/feeds/public/attack.py b/plugins/feeds/public/attack.py index 56e502fdd..08898bdb3 100644 --- a/plugins/feeds/public/attack.py +++ b/plugins/feeds/public/attack.py @@ -193,11 +193,11 @@ def run(self): logging.info("No response: skipping MitreAttack update") return - tempdir = tempfile.TemporaryDirectory() - ZipFile(BytesIO(response.content)).extractall(path=tempdir.name) - enterprise_attack = os.path.join( - tempdir.name, f"cti-ATT-CK-{_VERSION}", "enterprise-attack" - ) + with tempfile.TemporaryDirectory() as tempdir: + ZipFile(BytesIO(response.content)).extractall(path=tempdir) + enterprise_attack = os.path.join( + tempdir, f"cti-ATT-CK-{_VERSION}", "enterprise-attack" + ) object_cache = {} diff --git a/plugins/feeds/public/dfiq.py b/plugins/feeds/public/dfiq.py index 95caddae1..1b596379b 100644 --- a/plugins/feeds/public/dfiq.py +++ b/plugins/feeds/public/dfiq.py @@ -25,9 +25,9 @@ def run(self): 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: diff --git a/tests/apiv2/dfiq.py b/tests/apiv2/dfiq.py index ebea8a800..550db7802 100644 --- a/tests/apiv2/dfiq.py +++ b/tests/apiv2/dfiq.py @@ -584,11 +584,11 @@ def test_standalone_approach_creation(self): self.assertIsNone(data["parent_id"]) def test_upload_dfiq_archive(self): - zip_archive = open("tests/dfiq_test_data/dfiq_test_data.zip", "rb") - response = client.post( - "/api/v2/dfiq/from_archive", - files={"archive": ("test_archive.zip", zip_archive, "application/zip")}, - ) + with open("tests/dfiq_test_data/dfiq_test_data.zip", "rb") as zip_archive: + response = client.post( + "/api/v2/dfiq/from_archive", + files={"archive": ("test_archive.zip", zip_archive, "application/zip")}, + ) data = response.json() self.assertEqual(response.status_code, 200, data) self.assertEqual(data, {"total_added": 4}) From 471b3c7d1bfe8b4d5c06f6b4f12b0051687761ad Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 16 Aug 2024 12:36:44 +0000 Subject: [PATCH 7/7] fix attack feed --- plugins/feeds/public/attack.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/plugins/feeds/public/attack.py b/plugins/feeds/public/attack.py index 08898bdb3..dfd1ac72a 100644 --- a/plugins/feeds/public/attack.py +++ b/plugins/feeds/public/attack.py @@ -193,11 +193,11 @@ def run(self): logging.info("No response: skipping MitreAttack update") return - with tempfile.TemporaryDirectory() as tempdir: - ZipFile(BytesIO(response.content)).extractall(path=tempdir) - enterprise_attack = os.path.join( - tempdir, f"cti-ATT-CK-{_VERSION}", "enterprise-attack" - ) + tempdir = tempfile.TemporaryDirectory() + ZipFile(BytesIO(response.content)).extractall(path=tempdir.name) + enterprise_attack = os.path.join( + tempdir.name, f"cti-ATT-CK-{_VERSION}", "enterprise-attack" + ) object_cache = {} @@ -255,6 +255,7 @@ def run(self): ) rel_count += 1 logging.info("Processed %s relationships", rel_count) + tempdir.cleanup() taskmanager.TaskManager.register_task(MitreAttack)