From 3903342fbc7a2acdc9b7c575f921a6333c04cfc2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A1szl=C3=B3=20Vask=C3=B3?=
 <1771332+vlaci@users.noreply.github.com>
Date: Thu, 22 Feb 2024 18:00:48 +0100
Subject: [PATCH] wip: move sandbox config to cli

we should(?) have tests where sandboxing is enabled.
---
 tests/test_cli.py    |  1 +
 unblob/cli.py        | 23 +++++++++++++++++++++
 unblob/processing.py | 48 ++++++++++++++++++++------------------------
 3 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index 3a00145889..1874037a43 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -270,6 +270,7 @@ def test_archive_success(
         handlers=BUILTIN_HANDLERS,
         verbose=expected_verbosity,
         progress_reporter=expected_progress_reporter,
+        sandbox_access_restrictions=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)
diff --git a/unblob/cli.py b/unblob/cli.py
index fd65675052..c9dcd2449f 100755
--- a/unblob/cli.py
+++ b/unblob/cli.py
@@ -11,6 +11,7 @@
 from rich.style import Style
 from rich.table import Table
 from structlog import get_logger
+from unblob_native.sandbox import AccessFS
 
 from unblob.models import DirectoryHandlers, Handlers, ProcessResult
 from unblob.plugins import UnblobPluginManager
@@ -277,6 +278,27 @@ def cli(
     extra_magics_to_skip = () if clear_skip_magics else DEFAULT_SKIP_MAGIC
     skip_magic = tuple(sorted(set(skip_magic).union(extra_magics_to_skip)))
 
+    sandbox_access_restrictions = [
+        # Python, shared libraries and so on
+        AccessFS.read("/"),
+        # Multiprocessing
+        AccessFS.read_write("/dev/shm"),  # noqa: S108
+        # Extracted contents
+        AccessFS.read_write(extract_root.as_posix()),
+        AccessFS.make_dir(extract_root.parent.as_posix()),
+    ]
+
+    if report_file:
+        sandbox_access_restrictions += [
+            AccessFS.read_write(report_file),
+            AccessFS.make_reg(report_file.parent),
+        ]
+    if log_path:
+        sandbox_access_restrictions += [
+            AccessFS.read_write(log_path),
+            AccessFS.make_reg(log_path.parent),
+        ]
+
     config = ExtractionConfig(
         extract_root=extract_root,
         force_extract=force,
@@ -294,6 +316,7 @@ def cli(
         progress_reporter=NullProgressReporter
         if verbose
         else RichConsoleProgressReporter,
+        sandbox_access_restrictions=sandbox_access_restrictions,
     )
 
     logger.info("Start processing file", file=file)
diff --git a/unblob/processing.py b/unblob/processing.py
index 44c5b3eb52..1392d62700 100644
--- a/unblob/processing.py
+++ b/unblob/processing.py
@@ -1,6 +1,6 @@
+import functools
 import multiprocessing
 import shutil
-import sys
 from operator import attrgetter
 from pathlib import Path
 from typing import Iterable, List, Optional, Sequence, Set, Tuple, Type, Union
@@ -104,6 +104,7 @@ class ExtractionConfig:
     dir_handlers: DirectoryHandlers = BUILTIN_DIR_HANDLERS
     verbose: int = 1
     progress_reporter: Type[ProgressReporter] = NullProgressReporter
+    sandbox_access_restrictions: List[AccessFS] = []
 
     def get_extract_dir_for(self, path: Path) -> Path:
         """Return extraction dir under root with the name of path."""
@@ -117,28 +118,28 @@ def get_extract_dir_for(self, path: Path) -> Path:
         return extract_dir.expanduser().resolve()
 
 
-def sandbox(extract_dir: Path, report_file: Optional[Path]):
-    restrictions = [
-        AccessFS.read("/"),
-        AccessFS.read_write("/dev/shm"),  # noqa: S108
-        AccessFS.read_write(extract_dir.as_posix()),
-        AccessFS.make_dir(extract_dir.parent.as_posix()),
-    ]
+def call_once(fn):
+    fn.__called_once__ = False
 
-    if report_file:
-        restrictions += [
-            AccessFS.read_write(report_file),
-            AccessFS.make_reg(report_file.parent),
-        ]
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        if fn.__called_once__:
+            return
+        fn.__called_once__ = True
+        fn(*args, **kwargs)
 
-    if "pytest" in sys.modules:
-        restrictions += [
-            AccessFS.read_write("/tmp"),  # noqa: S108
-            AccessFS.read_write("/build"),
-            AccessFS.read_write(Path(__file__).parent.parent.resolve().as_posix()),
-        ]
+    return wrapper
 
-    restrict_access(*restrictions)
+
+@call_once
+def try_enter_sandbox(config: ExtractionConfig):
+    if not config.sandbox_access_restrictions:
+        return
+    try:
+        restrict_access(*config.sandbox_access_restrictions)
+    except SandboxError:
+        logger.warning("Sandboxing FS access is unavailable on this system, skipping.")
+        restrict_access(*config.sandbox_access_restrictions)
 
 
 @terminate_gracefully
@@ -165,12 +166,7 @@ def process_file(
         )
         return ProcessResult()
 
-    try:
-        if not hasattr(process_file, "_sandboxed"):
-            sandbox(extract_dir, report_file)
-            process_file._sandboxed = True  # noqa: SLF001
-    except SandboxError:
-        logger.warning("Sandboxing FS access is unavailable on this system, skipping.")
+    try_enter_sandbox(config)
 
     process_result = _process_task(config, task)