Use tmpdir with enough free space.

7 files changed, 93 insertions, 8 deletions

diff --git a/src/sec_certs/dataset/dataset.py b/src/sec_certs/dataset/dataset.py
index be070580..446cc5ff 100644
--- a/src/sec_certs/dataset/dataset.py
+++ b/src/sec_certs/dataset/dataset.py
@@ -182,10 +182,12 @@ class Dataset(Generic[CertSubType], ComplexSerializableType, ABC):
             if not path.is_dir():
                 raise ValueError("Path needs to be a directory.")
         if artifacts:
-            with tempfile.TemporaryDirectory() as tmp_dir:
+            fsize = helpers.query_file_size(str(archive_url))
+            base_tmpdir = tempfile.gettempdir() if fsize is None else helpers.tempdir_for(fsize)
+            with tempfile.TemporaryDirectory(dir=base_tmpdir) as tmp_dir:
                 dset_path = Path(tmp_dir) / "dataset.tar.gz"
                 res = helpers.download_file(
-                    archive_url,
+                    str(archive_url),
                     dset_path,
                     show_progress_bar=True,
                     progress_bar_desc=progress_bar_desc,
@@ -201,7 +203,7 @@ class Dataset(Generic[CertSubType], ComplexSerializableType, ABC):
             with tempfile.TemporaryDirectory() as tmp_dir:
                 dset_path = Path(tmp_dir) / "dataset.json"
                 helpers.download_file(
-                    snapshot_url,
+                    str(snapshot_url),
                     dset_path,
                     show_progress_bar=True,
                     progress_bar_desc=progress_bar_desc,
diff --git a/src/sec_certs/heuristics/__init__.py b/src/sec_certs/heuristics/__init__.py
new file mode 100644
index 00000000..ce714058
--- /dev/null
+++ b/src/sec_certs/heuristics/__init__.py
@@ -0,0 +1 @@
+"""This package provides heuristics for extracting information from certificates."""
diff --git a/src/sec_certs/utils/__init__.py b/src/sec_certs/utils/__init__.py
index 45b35aea..02ccc56d 100644
--- a/src/sec_certs/utils/__init__.py
+++ b/src/sec_certs/utils/__init__.py
@@ -1 +1 @@
-"""This package provides utilities used throught the framework."""
+"""This package provides utilities used throughout the framework."""
diff --git a/src/sec_certs/utils/helpers.py b/src/sec_certs/utils/helpers.py
index d34e9f0e..63de6fd9 100644
--- a/src/sec_certs/utils/helpers.py
+++ b/src/sec_certs/utils/helpers.py
@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import hashlib
 import logging
+import os
 import re
+import shutil
 import time
 from collections.abc import Collection
 from contextlib import nullcontext
@@ -28,7 +30,72 @@ _PROXIES = {
 }
 
 
-def download_file(
+def tempdirs() -> list[str]:
+    """Get a list of potential temporary directory bases, as tempfile.gettempdir() does."""
+    dirlist = []
+
+    # First, try the environment.
+    for envname in "TMPDIR", "TEMP", "TMP":
+        dirname = os.getenv(envname)
+        if dirname:
+            dirlist.append(dirname)
+
+    # Failing that, try OS-specific locations.
+    if os.name == "nt":
+        dirlist.extend(
+            [
+                str(Path(r"~\AppData\Local\Temp").expanduser()),
+                os.path.expandvars(r"%SYSTEMROOT%\Temp"),
+                r"c:\temp",
+                r"c:\tmp",
+                r"\temp",
+                r"\tmp",
+            ]
+        )
+    else:
+        dirlist.extend(["/tmp", "/var/tmp", "/usr/tmp"])
+
+    # As a last resort, the current directory.
+    try:
+        dirlist.append(str(Path.cwd()))
+    except (AttributeError, OSError):
+        dirlist.append(os.curdir)
+
+    return dirlist
+
+
+def tempdir_for(size: int) -> str:
+    """
+    Find a temporary directory base that fits the given size.
+
+    :param size: the minimum size required
+    :returns: the name of the directory
+    :raises OSError: if no suitable temporary directory is found
+    """
+    for dirname in tempdirs():
+        try:
+            usage = shutil.disk_usage(dirname)
+        except OSError:
+            continue
+        if usage.free > size:
+            return dirname
+    raise OSError("No suitable temporary directory found with enough space")
+
+
+def query_file_size(url: str) -> int | None:
+    """Use a HEAD request to query the file size of a remote file."""
+    try:
+        r = requests.head(url, timeout=constants.REQUEST_TIMEOUT)
+        if r.status_code == requests.codes.ok:
+            return int(r.headers.get("content-length", 0))
+    except requests.exceptions.Timeout:
+        return None
+    except Exception as e:
+        logger.error(f"Failed to query file size from {url}; {e}")
+    return None
+
+
+def download_file(  # noqa: C901
     url: str,
     output: Path,
     delay: float = 0,
@@ -36,13 +103,17 @@ def download_file(
     progress_bar_desc: str | None = None,
     proxy: bool = False,
 ) -> str | int:
+    """Download a file from a URL to a local path."""
     try:
-        time.sleep(delay)
+        proxied = False
         if proxy:
             for upstream in _PROXIES:
                 if upstream in url:
+                    proxied = True
                     url = url.replace(upstream, _PROXIES[upstream])
                     break
+        if not proxied:
+            time.sleep(delay)
         # See https://github.com/psf/requests/issues/3953 for header justification
         r = requests.get(
             url,
diff --git a/tests/cc/test_cc_dataset.py b/tests/cc/test_cc_dataset.py
index 4c988b5b..cf2468a9 100644
--- a/tests/cc/test_cc_dataset.py
+++ b/tests/cc/test_cc_dataset.py
@@ -78,6 +78,12 @@ def test_download_and_convert_pdfs(toy_dataset: CCDataset, data_dir: Path):
         )
 
 
+@pytest.mark.slow
+def test_from_web():
+    dset = CCDataset.from_web()
+    assert len(dset) > 6000
+
+
 def test_dataset_to_json(toy_dataset: CCDataset, data_dir: Path, tmp_path: Path):
     toy_dataset.to_json(tmp_path / "dset.json")
 
diff --git a/tests/cc/test_cc_maintenance_updates.py b/tests/cc/test_cc_maintenance_updates.py
index 05148cd7..176a9244 100644
--- a/tests/cc/test_cc_maintenance_updates.py
+++ b/tests/cc/test_cc_maintenance_updates.py
@@ -77,9 +77,8 @@ def test_to_pandas(mu_dset: CCDatasetMaintenanceUpdates):
     assert set(df.columns) == set(CCMaintenanceUpdate.pandas_columns) - {"dgst"}
 
 
-@pytest.mark.skip(reason="Will work only with fresh snapshot on sec-certs.org")
+@pytest.mark.slow
 def test_from_web():
     dset = CCDatasetMaintenanceUpdates.from_web()
     assert dset is not None
     assert len(dset) >= 492  # Contents as of November 2022, maintenances should not disappear
-    assert "cert_8f08cacb49a742fb_update_559ed93dd80320b5" in dset  # random cert verified to be present
diff --git a/tests/fips/test_fips_dataset.py b/tests/fips/test_fips_dataset.py
index 81ad019d..c9785c8a 100644
--- a/tests/fips/test_fips_dataset.py
+++ b/tests/fips/test_fips_dataset.py
@@ -35,6 +35,12 @@ def test_dataset_to_json(toy_dataset: FIPSDataset, data_dir: Path, tmp_path: Pat
     assert data == template_data
 
 
+@pytest.mark.slow
+def test_from_web():
+    dset = FIPSDataset.from_web()
+    assert len(dset) > 4000
+
+
 def test_dataset_from_json(toy_dataset: FIPSDataset, data_dir: Path, tmp_path: Path):
     assert toy_dataset == FIPSDataset.from_json(data_dir / "toy_dataset.json")