feat(processing): add unblob dedicated temporary directory handling

Author: qkaiser

python/unblob/processing.py

@@ -1,6 +1,10 @@
+import atexit
 import multiprocessing
+import os
 import shutil
+import tempfile
 from collections.abc import Iterable, Sequence
+from contextlib import contextmanager
 from operator import attrgetter
 from pathlib import Path
 from typing import Optional, Union
@@ -100,6 +104,21 @@ class ExtractionConfig:
     dir_handlers: DirectoryHandlers = BUILTIN_DIR_HANDLERS
     verbose: int = 1
     progress_reporter: type[ProgressReporter] = NullProgressReporter
+    tmp_dir: Path = attrs.field(
+        factory=lambda: Path(tempfile.mkdtemp(prefix="unblob-tmp-"))
+    )
+
+    def __attrs_post_init__(self):
+        atexit.register(self._cleanup_tmp_dir)
+
+    def _cleanup_tmp_dir(self):
+        if isinstance(self.tmp_dir, Path) and self.tmp_dir.exists():
+            try:
+                shutil.rmtree(self.tmp_dir)
+            except Exception as e:
+                logger.warning(
+                    "Failed to clean up tmp_dir", tmp_dir=self.tmp_dir, exc_info=e
+                )
 
     def _get_output_path(self, path: Path) -> Path:
         """Return path under extract root."""
@@ -227,6 +246,24 @@ def write_json_report(report_file: Path, process_result: ProcessResult):
         logger.info("JSON report written", path=report_file)
 
 
+@contextmanager
+def tmp_dir(path):
+    """Context manager that temporarily sets all temp directory env vars."""
+    tmp_vars = ("TMP", "TMPDIR", "TEMP", "TEMPDIR")
+    saved = {}
+    try:
+        for var in tmp_vars:
+            saved[var] = os.environ.get(var)
+            os.environ[var] = str(path)
+        yield
+    finally:
+        for var, original in saved.items():
+            if original is None:
+                os.environ.pop(var, None)
+            else:
+                os.environ[var] = original
+
+
 class Processor:
     def __init__(self, config: ExtractionConfig):
         self._config = config
@@ -244,7 +281,8 @@ def __init__(self, config: ExtractionConfig):
     def process_task(self, task: Task) -> TaskResult:
         result = TaskResult(task=task)
         try:
-            self._process_task(result, task)
+            with tmp_dir(self._config.tmp_dir.as_posix()):
+                self._process_task(result, task)
         except Exception as exc:
             self._process_error(result, exc)
         return result

python/unblob/sandbox.py

@@ -55,6 +55,10 @@ def __init__(
             AccessFS.remove_file(config.extract_root),
             AccessFS.make_dir(config.extract_root.parent),
             AccessFS.read_write(log_path),
+            # Allow access to the managed temp directory for handlers
+            AccessFS.read_write(config.tmp_dir),
+            AccessFS.remove_dir(config.tmp_dir),
+            AccessFS.remove_file(config.tmp_dir),
             *extra_passthrough,
         ]
 

tests/test_cli.py

@@ -274,6 +274,7 @@ def test_archive_success(
         verbose=expected_verbosity,
         progress_reporter=expected_progress_reporter,
     )
+    config.tmp_dir = mock.ANY
     process_file_mock.assert_called_once_with(config, in_path, None)
     logger_config_mock.assert_called_once_with(expected_verbosity, tmp_path, log_path)
 

tests/test_sandbox.py

@@ -34,6 +34,8 @@ def test_necessary_resources_can_be_created_in_sandbox(
 ):
     directory_in_extract_root = extraction_config.extract_root / "path" / "to" / "dir"
     file_in_extract_root = directory_in_extract_root / "file"
+    file_in_tmp_dir = extraction_config.tmp_dir / "tmp_file"
+    directory_in_tmp_dir = extraction_config.tmp_dir / "tmp_dir"
 
     sandbox.run(extraction_config.extract_root.mkdir, parents=True)
     sandbox.run(directory_in_extract_root.mkdir, parents=True)
@@ -45,6 +47,12 @@ def test_necessary_resources_can_be_created_in_sandbox(
     log_path.touch()
     sandbox.run(log_path.write_text, "log line")
 
+    sandbox.run(directory_in_tmp_dir.mkdir, parents=True)
+    sandbox.run(file_in_tmp_dir.touch)
+    sandbox.run(file_in_tmp_dir.write_text, "tmp file content")
+    sandbox.run(file_in_tmp_dir.unlink)
+    sandbox.run(directory_in_tmp_dir.rmdir)
+
 
 def test_access_outside_sandbox_is_not_possible(sandbox: Sandbox, tmp_path: Path):
     unrelated_dir = tmp_path / "unrelated" / "path"