Fix leaky handles (#1210)

.github/workflows/feedtest.yml

@@ -6,8 +6,13 @@ on:
   pull_request:
 
 jobs:
+  debug:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Debug
+      run: echo "${{ toJson(github) }}"
   testfeeds:
-    if: ${{ contains(github.event.pull_request.labels.*.name, 'tasks:feed') || github.event_type == 'schedule' }}
+    if: ${{ github.event_type == 'schedule' || contains(github.event.pull_request.labels.*.name, 'tasks:feed') }}
     runs-on: ubuntu-latest
     env:
       YETI_REDIS_HOST: localhost
@@ -43,6 +48,8 @@ jobs:
       run: cp yeti.conf.sample yeti.conf
     - name: Start redis & arangodb conainers
       run: docker compose -f extras/docker/dev/docker-compose.yaml up -d redis arangodb
+    - name: Start event consumer
+      run: poetry run python -m core.events.consumers events & sleep 5
     - name: Run test feeds
       run: |
         poetry run python -m unittest discover -s tests/ -p 'feeds.py'

core/web/apiv2/dfiq.py

@@ -96,12 +96,12 @@ def config() -> DFIQConfigResponse:
 @router.post("/from_archive")
 def from_archive(httpreq: Request, archive: UploadFile) -> dict[str, int]:
     """Uncompresses a ZIP archive and processes the DFIQ content inside it."""
-    tempdir = tempfile.TemporaryDirectory()
-    contents = archive.file.read()
-    ZipFile(BytesIO(contents)).extractall(path=tempdir.name)
-    total_added = dfiq.read_from_data_directory(
-        f"{tempdir.name}/*/*.yaml", username=httpreq.state.username
-    )
+    with tempfile.TemporaryDirectory() as tempdir:
+        contents = archive.file.read()
+        ZipFile(BytesIO(contents)).extractall(path=tempdir)
+        total_added = dfiq.read_from_data_directory(
+            f"{tempdir}/*/*.yaml", username=httpreq.state.username
+        )
     return {"total_added": total_added}
 
 
@@ -179,54 +179,54 @@ def to_archive(request: DFIQSearchRequest) -> FileResponse:
         dfiq.DFIQType.question: "questions",
     }
 
-    tempdir = tempfile.TemporaryDirectory()
-    public_objs = []
-    internal_objs = []
-    for obj in dfiq_objects:
-        if obj.dfiq_tags and "internal" in obj.dfiq_tags:
-            internal_objs.append(obj)
-        else:
-            if obj.type == dfiq.DFIQType.question:
-                public_version = obj.model_copy()
-                internal_approaches = False
-                for approach in obj.approaches:
-                    if "internal" in approach.tags:
-                        internal_approaches = True
-                        break
-                if internal_approaches:
-                    public_version.approaches = [
-                        a for a in obj.approaches if "internal" not in a.tags
-                    ]
-                    public_objs.append(public_version)
-                    internal_objs.append(obj)
+    with tempfile.TemporaryDirectory() as tempdir:
+        public_objs = []
+        internal_objs = []
+        for obj in dfiq_objects:
+            if obj.dfiq_tags and "internal" in obj.dfiq_tags:
+                internal_objs.append(obj)
+            else:
+                if obj.type == dfiq.DFIQType.question:
+                    public_version = obj.model_copy()
+                    internal_approaches = False
+                    for approach in obj.approaches:
+                        if "internal" in approach.tags:
+                            internal_approaches = True
+                            break
+                    if internal_approaches:
+                        public_version.approaches = [
+                            a for a in obj.approaches if "internal" not in a.tags
+                        ]
+                        public_objs.append(public_version)
+                        internal_objs.append(obj)
+                    else:
+                        public_objs.append(obj)
                 else:
                     public_objs.append(obj)
-            else:
-                public_objs.append(obj)
-
-    for dir_name in ["public", "internal"]:
-        os.makedirs(f"{tempdir.name}/{dir_name}")
-
-    for obj in public_objs:
-        with open(f"{tempdir.name}/public/{obj.uuid}.yaml", "w") as f:
-            f.write(obj.to_yaml())
-
-    for obj in internal_objs:
-        with open(f"{tempdir.name}/internal/{obj.uuid}.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 public_objs:
-                zipf.write(
-                    f"{tempdir.name}/public/{obj.uuid}.yaml",
-                    f"public/{_TYPE_TO_DUMP_DIR[obj.type]}/{obj.uuid}.yaml",
-                )
-            for obj in internal_objs:
-                zipf.write(
-                    f"{tempdir.name}/internal/{obj.uuid}.yaml",
-                    f"internal/{_TYPE_TO_DUMP_DIR[obj.type]}/{obj.uuid}.yaml",
-                )
+
+        for dir_name in ["public", "internal"]:
+            os.makedirs(f"{tempdir}/{dir_name}")
+
+        for obj in public_objs:
+            with open(f"{tempdir}/public/{obj.uuid}.yaml", "w") as f:
+                f.write(obj.to_yaml())
+
+        for obj in internal_objs:
+            with open(f"{tempdir}/internal/{obj.uuid}.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 public_objs:
+                    zipf.write(
+                        f"{tempdir}/public/{obj.uuid}.yaml",
+                        f"public/{_TYPE_TO_DUMP_DIR[obj.type]}/{obj.uuid}.yaml",
+                    )
+                for obj in internal_objs:
+                    zipf.write(
+                        f"{tempdir}/internal/{obj.uuid}.yaml",
+                        f"internal/{_TYPE_TO_DUMP_DIR[obj.type]}/{obj.uuid}.yaml",
+                    )
 
     return FileResponse(archive.name, media_type="application/zip", filename="dfiq.zip")
 

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

@@ -28,13 +28,13 @@ 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(
-            os.path.join(tempdir.name, "*", "dfiq", "data", "*", "*.yaml"),
-            "DFIQFeed",
-            overwrite=True,
-        )
+        with tempfile.TemporaryDirectory() as tempdir:
+            ZipFile(BytesIO(response.content)).extractall(path=tempdir)
+            dfiq.read_from_data_directory(
+                os.path.join(tempdir, "*", "dfiq", "data", "*", "*.yaml"),
+                "DFIQFeed",
+                overwrite=True,
+            )
 
         extra_dirs = yeti_config.get("dfiq", "extra_dirs")
         if not extra_dirs:

plugins/feeds/public/signaturebase.py

@@ -37,9 +37,9 @@ def run(self):
             logging.info("No response: skipping Neo23x0 Signature base update")
             return
 
-        tempdir = tempfile.TemporaryDirectory()
-        ZipFile(BytesIO(response.content)).extractall(path=tempdir.name)
-        rules_path = os.path.join(tempdir.name, "signature-base-master", "yara")
+        with tempfile.TemporaryDirectory() as tempdir:
+            ZipFile(BytesIO(response.content)).extractall(path=tempdir)
+            rules_path = os.path.join(tempdir, "signature-base-master", "yara")
 
         for file in glob.glob(f"{rules_path}/*.yar"):
             with open(file, "r") as f:
@@ -58,7 +58,7 @@ def run(self):
                 location="filesystem",
             ).save()
 
-        yara_object.tag(["Neo23x0", "signature-base"])
+            yara_object.tag(["Neo23x0", "signature-base"])
 
 
 taskmanager.TaskManager.register_task(Neo23x0SignatureBase)

tests/apiv2/dfiq.py

@@ -478,11 +478,11 @@ def test_standalone_question_creation(self):
         self.assertEqual(data["parent_ids"], [])
 
     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": 3})