From bec9eb948b31931dd7abd40a5c6415866a95b8aa Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Mon, 23 Dec 2024 10:18:23 +0100 Subject: [PATCH] Fix leaky handles (#1210) --- .github/workflows/feedtest.yml | 9 ++- core/web/apiv2/dfiq.py | 104 +++++++++++++------------- plugins/feeds/public/artifacts.py | 10 +-- plugins/feeds/public/attack.py | 1 + plugins/feeds/public/dfiq.py | 14 ++-- plugins/feeds/public/signaturebase.py | 8 +- tests/apiv2/dfiq.py | 10 +-- 7 files changed, 82 insertions(+), 74 deletions(-) diff --git a/.github/workflows/feedtest.yml b/.github/workflows/feedtest.yml index b8c1cbbfb..42b9fdf50 100644 --- a/.github/workflows/feedtest.yml +++ b/.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' diff --git a/core/web/apiv2/dfiq.py b/core/web/apiv2/dfiq.py index c9278da76..228d92269 100644 --- a/core/web/apiv2/dfiq.py +++ b/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") 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..dfd1ac72a 100644 --- a/plugins/feeds/public/attack.py +++ b/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) diff --git a/plugins/feeds/public/dfiq.py b/plugins/feeds/public/dfiq.py index 81022a6d4..ad3ff4c06 100644 --- a/plugins/feeds/public/dfiq.py +++ b/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: diff --git a/plugins/feeds/public/signaturebase.py b/plugins/feeds/public/signaturebase.py index cbc6c9ef9..106b7ae77 100644 --- a/plugins/feeds/public/signaturebase.py +++ b/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) diff --git a/tests/apiv2/dfiq.py b/tests/apiv2/dfiq.py index cadfa4189..6a10d7c7a 100644 --- a/tests/apiv2/dfiq.py +++ b/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})