fix: cleanup PyPI package temp directory and prevent leaks

src/macaron/malware_analyzer/pypi_heuristics/metadata/type_stub_file.py

@@ -1,11 +1,12 @@
-# Copyright (c) 2024 - 2025, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 - 2026, Oracle and/or its affiliates. All rights reserved.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/.
 
 """This analyzer checks if a PyPI package has minimal .pyi stub content."""
 
 import logging
 import os
 
+from macaron.errors import SourceCodeError
 from macaron.json_tools import JsonType
 from macaron.malware_analyzer.pypi_heuristics.base_analyzer import BaseHeuristicAnalyzer
 from macaron.malware_analyzer.pypi_heuristics.heuristics import HeuristicResult, Heuristics
@@ -42,15 +43,15 @@ def analyze(self, pypi_package_json: PyPIPackageJsonAsset) -> tuple[HeuristicRes
         # TODO: .pyi stub files may be present in both source distributions (sdist) and wheels.
         # Currently, we only check the sdist, which can lead to false positives in this heuristic.
         # To improve accuracy, we should also check for stub files in the wheel distribution.
-        result = pypi_package_json.download_sourcecode()
-        if not result:
+        try:
+            with pypi_package_json.sourcecode():
+                file_count = sum(
+                    sum(1 for f in files if f.endswith(".pyi"))
+                    for _, _, files in os.walk(pypi_package_json.package_sourcecode_path)
+                )
+        except SourceCodeError:
             return HeuristicResult.SKIP, {"message": "No source code files have been downloaded.", "pyi_files": 0}
 
-        file_count = sum(
-            sum(1 for f in files if f.endswith(".pyi"))
-            for _, _, files in os.walk(pypi_package_json.package_sourcecode_path)
-        )
-
         if file_count >= self.FILES_THRESHOLD:
             return HeuristicResult.PASS, {"message": "Package has sufficient pyi files", "pyi_files": file_count}
 

src/macaron/slsa_analyzer/package_registry/pypi_registry.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import bisect
+import copy
 import hashlib
 import logging
 import os
@@ -208,6 +209,7 @@ def cleanup_sourcecode_directory(
             logger.debug(error_message)
         try:
             shutil.rmtree(directory, onerror=_handle_temp_dir_clean)
+            logger.debug("Successfully cleaned up temporary directory %s.", directory)
         except SourceCodeError as tempdir_exception:
             tempdir_exception_msg = (
                 f"Unable to cleanup temporary directory {directory} for source code: {tempdir_exception}"
@@ -266,27 +268,49 @@ def download_package_sourcecode(self, url: str) -> str:
         source_file = os.path.join(temp_dir, file_name)
         timeout = defaults.getint("downloads", "timeout", fallback=120)
         size_limit = defaults.getint("downloads", "max_download_size", fallback=10000000)
-        if not download_file_with_size_limit(url, {}, source_file, timeout, size_limit):
-            self.cleanup_sourcecode_directory(temp_dir, "Could not download the file.")
+
+        download_succeeded = False
+        try:
+            download_succeeded = download_file_with_size_limit(url, {}, source_file, timeout, size_limit)
+        except requests.exceptions.RequestException as error:
+            self.cleanup_sourcecode_directory(
+                temp_dir, f"Error downloading source code from file {file_name}: {error}", error
+            )
+        if not download_succeeded:
+            self.cleanup_sourcecode_directory(
+                temp_dir, f"Error downloading source code from file {file_name}."
+            )
 
         if not tarfile.is_tarfile(source_file):
             self.cleanup_sourcecode_directory(temp_dir, f"Unable to extract source code from file {file_name}")
 
         try:
             with tarfile.open(source_file, "r:gz") as sourcecode_tar:
-                sourcecode_tar.extractall(temp_dir, filter="data")
+                members = sourcecode_tar.getmembers()
+                if members and all(
+                    member.name == package_name or member.name.startswith(f"{package_name}/")
+                    for member in members
+                ):
+                    # Most sdists wrap their contents in a single package-version directory.
+                    # Strip that wrapper during extraction so the returned temp directory
+                    # contains the package files directly and cleanup removes the whole tree.
+                    members_to_extract = []
+                    for member in members:
+                        if member.name == package_name:
+                            continue
+                        stripped_member = copy.copy(member)
+                        stripped_member.name = member.name.removeprefix(f"{package_name}/")
+                        members_to_extract.append(stripped_member)
+                    members = members_to_extract
+
+                sourcecode_tar.extractall(temp_dir, members=members, filter="data")
         except tarfile.TarError as tar_error:
             self.cleanup_sourcecode_directory(
                 temp_dir, f"Error extracting source code tar file: {tar_error}", tar_error
             )
 
         os.remove(source_file)
 
-        extracted_dir = os.listdir(temp_dir)
-        if len(extracted_dir) == 1 and extracted_dir[0] == package_name:
-            # Structure used package name and version as top-level directory.
-            temp_dir = os.path.join(temp_dir, extracted_dir[0])
-
         logger.debug("Temporary download and unzip of %s stored in %s", file_name, temp_dir)
         return temp_dir
 
@@ -321,8 +345,17 @@ def download_package_wheel(self, url: str) -> str:
         timeout = defaults.getint("downloads", "timeout", fallback=120)
         size_limit = defaults.getint("downloads", "max_download_size", fallback=10000000)
 
-        if not download_file_with_size_limit(url, {}, wheel_file, timeout, size_limit):
-            self.cleanup_sourcecode_directory(temp_dir, "Could not download the file.")
+        download_succeeded = False
+        try:
+            download_succeeded = download_file_with_size_limit(url, {}, wheel_file, timeout, size_limit)
+        except requests.exceptions.RequestException as error:
+            self.cleanup_sourcecode_directory(
+                temp_dir, f"Error downloading wheel from file {file_name}: {error}", error
+            )
+        if not download_succeeded:
+            self.cleanup_sourcecode_directory(
+                temp_dir, f"Error downloading wheel from file {file_name}."
+            )
 
         # Wheel is a zip
         if not zipfile.is_zipfile(wheel_file):
@@ -943,10 +976,12 @@ def wheel(self, download_binaries: bool) -> Generator[None]:
             raise WheelTagError("Macaron does not currently support analysis of non-pure Python wheels.")
         if not self.download_wheel():
             raise SourceCodeError("Unable to download requested wheel.")
-        yield
-        if self.wheel_path:
-            # Name for cleanup_sourcecode_directory could be refactored here
-            PyPIRegistry.cleanup_sourcecode_directory(self.wheel_path)
+        try:
+            yield
+        finally:
+            if self.wheel_path:
+                # Name for cleanup_sourcecode_directory could be refactored here
+                PyPIRegistry.cleanup_sourcecode_directory(self.wheel_path)
 
     def download_wheel(self) -> bool:
         """Download and extract wheel metadata to a temporary directory.
@@ -996,9 +1031,11 @@ def sourcecode(self) -> Generator[None]:
         """Download and cleanup source code of the package with a context manager."""
         if not self.download_sourcecode():
             raise SourceCodeError("Unable to download package source code.")
-        yield
-        if self.package_sourcecode_path:
-            PyPIRegistry.cleanup_sourcecode_directory(self.package_sourcecode_path)
+        try:
+            yield
+        finally:
+            if self.package_sourcecode_path:
+                PyPIRegistry.cleanup_sourcecode_directory(self.package_sourcecode_path)
 
     def download_sourcecode(self) -> bool:
         """Get the source code of the package and store it in a temporary directory.

tests/malware_analyzer/pypi/test_type_stub_file.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024 - 2025, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2024 - 2026, Oracle and/or its affiliates. All rights reserved.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/.
 
 """Tests for the TypeStubFileAnalyzer heuristic."""
@@ -7,6 +7,7 @@
 
 import pytest
 
+from macaron.errors import SourceCodeError
 from macaron.malware_analyzer.pypi_heuristics.heuristics import HeuristicResult
 from macaron.malware_analyzer.pypi_heuristics.metadata.type_stub_file import TypeStubFileAnalyzer
 
@@ -26,7 +27,7 @@ def test_analyze_sufficient_files_pass(analyzer: TypeStubFileAnalyzer, pypi_pack
         result, _ = analyzer.analyze(pypi_package_json)
 
     assert result == HeuristicResult.PASS
-    pypi_package_json.download_sourcecode.assert_called_once()
+    pypi_package_json.sourcecode.assert_called_once()
 
 
 def test_analyze_exactly_threshold_files_pass(analyzer: TypeStubFileAnalyzer, pypi_package_json: MagicMock) -> None:
@@ -64,7 +65,7 @@ def test_analyze_no_files_fail(analyzer: TypeStubFileAnalyzer, pypi_package_json
 
 def test_analyze_download_failed_raises_error(analyzer: TypeStubFileAnalyzer, pypi_package_json: MagicMock) -> None:
     """Test the analyzer when source code download fails."""
-    pypi_package_json.download_sourcecode.return_value = False
+    pypi_package_json.sourcecode.side_effect = SourceCodeError("download failed")
     assert (
         HeuristicResult.SKIP,
         {"message": "No source code files have been downloaded.", "pyi_files": 0},

tests/slsa_analyzer/package_registry/test_pypi_registry.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2026 - 2026, Oracle and/or its affiliates. All rights reserved.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/.
+
+"""Tests for the PyPI package registry."""
+
+import os
+import shutil
+import tarfile
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import requests
+
+from macaron.errors import InvalidHTTPResponseError, SourceCodeError
+from macaron.slsa_analyzer.package_registry import pypi_registry
+from macaron.slsa_analyzer.package_registry.pypi_registry import PyPIInspectorAsset, PyPIPackageJsonAsset, PyPIRegistry
+
+
+def _raise_during_sourcecode_context(asset: PyPIPackageJsonAsset) -> None:
+    """Raise an error inside the sourcecode context manager."""
+    with asset.sourcecode():
+        raise SourceCodeError("analysis failed")
+
+
+def test_download_package_sourcecode_flattens_top_level_directory(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Downloaded sdists with a single top-level package directory should clean up from the temp root."""
+    package_name = "example-1.0.0"
+    pyproject_path = os.path.join(tmp_path, "pyproject.toml")
+    with open(pyproject_path, "w", encoding="utf-8") as pyproject_file:
+        pyproject_file.write("[build-system]\n")
+    archive_path = os.path.join(tmp_path, f"{package_name}.tar.gz")
+    with tarfile.open(archive_path, "w:gz") as archive:
+        archive.add(pyproject_path, arcname=os.path.join(package_name, "pyproject.toml"))
+
+    temp_root = os.path.join(tmp_path, "downloads")
+    os.makedirs(temp_root)
+
+    def mkdtemp(prefix: str) -> str:
+        return os.path.join(temp_root, f"{prefix}abcdef")
+
+    def download_file(_url: str, _headers: dict, dest: str, _timeout: int, _size_limit: int) -> bool:
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        shutil.copyfile(archive_path, dest)
+        return True
+
+    monkeypatch.setattr(pypi_registry.tempfile, "mkdtemp", mkdtemp)
+    monkeypatch.setattr(pypi_registry, "download_file_with_size_limit", download_file)
+
+    source_path = PyPIRegistry().download_package_sourcecode(f"https://example.test/{package_name}.tar.gz")
+
+    assert source_path == os.path.join(temp_root, f"{package_name}_abcdef")
+    assert os.path.exists(os.path.join(source_path, "pyproject.toml"))
+    assert not os.path.exists(os.path.join(source_path, package_name))
+
+    PyPIRegistry.cleanup_sourcecode_directory(source_path)
+    assert not os.path.exists(source_path)
+
+
+def test_sourcecode_context_cleans_up_when_analysis_raises(tmp_path: Path) -> None:
+    """The sourcecode context manager must remove downloads even if the caller fails."""
+    source_path = os.path.join(tmp_path, "example-1.0.0_abcdef")
+    os.makedirs(source_path)
+
+    registry = PyPIRegistry()
+    registry.download_package_sourcecode = MagicMock(return_value=str(source_path))  # type: ignore[method-assign]
+    asset = PyPIPackageJsonAsset("example", "1.0.0", False, registry, {}, PyPIInspectorAsset("", [], {}))
+    asset.get_sourcecode_url = MagicMock(  # type: ignore[method-assign]
+        return_value="https://example.test/example-1.0.0.tar.gz"
+    )
+
+    with pytest.raises(SourceCodeError, match="analysis failed"):
+        _raise_during_sourcecode_context(asset)
+
+    assert not os.path.exists(source_path)
+
+
+def test_download_package_sourcecode_cleans_up_when_download_raises(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The sdist temp directory should be removed if the download helper raises."""
+    package_name = "example-1.0.0"
+    source_path = os.path.join(tmp_path, f"{package_name}_abcdef")
+
+    def mkdtemp(prefix: str) -> str:
+        path = os.path.join(tmp_path, f"{prefix}abcdef")
+        os.makedirs(path)
+        return path
+
+    def download_file(_url: str, _headers: dict, _dest: str, _timeout: int, _size_limit: int) -> bool:
+        raise requests.exceptions.ConnectionError("download crashed")
+
+    monkeypatch.setattr(pypi_registry.tempfile, "mkdtemp", mkdtemp)
+    monkeypatch.setattr(pypi_registry, "download_file_with_size_limit", download_file)
+
+    with pytest.raises(InvalidHTTPResponseError, match="download crashed"):
+        PyPIRegistry().download_package_sourcecode(f"https://example.test/{package_name}.tar.gz")
+
+    assert not os.path.exists(source_path)
+
+
+def test_download_package_wheel_cleans_up_when_download_raises(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The wheel temp directory should be removed if the download helper raises."""
+    wheel_name = "example-1.0.0-py3-none-any"
+    wheel_path = os.path.join(tmp_path, f"{wheel_name}_abcdef")
+
+    def mkdtemp(prefix: str) -> str:
+        path = os.path.join(tmp_path, f"{prefix}abcdef")
+        os.makedirs(path)
+        return path
+
+    def download_file(_url: str, _headers: dict, _dest: str, _timeout: int, _size_limit: int) -> bool:
+        raise requests.exceptions.ConnectionError("download crashed")
+
+    monkeypatch.setattr(pypi_registry.tempfile, "mkdtemp", mkdtemp)
+    monkeypatch.setattr(pypi_registry, "download_file_with_size_limit", download_file)
+
+    with pytest.raises(InvalidHTTPResponseError, match="download crashed"):
+        PyPIRegistry().download_package_wheel(f"https://example.test/{wheel_name}.whl")
+
+    assert not os.path.exists(wheel_path)