refactor: shuffle code to their proper place

Author: e3krisztian

tests/handlers/filesystem/test_romfs.py

@@ -1,9 +1,7 @@
-from pathlib import Path
-
 import pytest
 
 from unblob.file_utils import File
-from unblob.handlers.filesystem.romfs import get_string, is_safe_path, valid_checksum
+from unblob.handlers.filesystem.romfs import get_string, valid_checksum
 
 
 @pytest.mark.parametrize(
@@ -44,23 +42,3 @@ def test_get_string(content, expected):
 )
 def test_valid_checksum(content, valid):
     assert valid_checksum(content) == valid
-
-
-@pytest.mark.parametrize(
-    "basedir, path, expected",
-    [
-        ("/lib/out", "/lib/out/file", True),
-        ("/lib/out", "file", True),
-        ("/lib/out", "dir/file", True),
-        ("/lib/out", "some/dir/file", True),
-        ("/lib/out", "some/dir/../file", True),
-        ("/lib/out", "some/dir/../../file", True),
-        ("/lib/out", "some/dir/../../../file", False),
-        ("/lib/out", "some/dir/../../../", False),
-        ("/lib/out", "some/dir/../../..", False),
-        ("/lib/out", "../file", False),
-        ("/lib/out", "/lib/out/../file", False),
-    ],
-)
-def test_is_safe_path(basedir, path, expected):
-    assert is_safe_path(Path(basedir), Path(path)) is expected

tests/test_file_utils.py

@@ -1,4 +1,5 @@
 import io
+from pathlib import Path
 from typing import List
 
 import pytest
@@ -14,13 +15,34 @@
     convert_int64,
     decode_multibyte_integer,
     get_endian,
+    is_safe_path,
     iterate_file,
     iterate_patterns,
     round_down,
     round_up,
 )
 
 
+@pytest.mark.parametrize(
+    "basedir, path, expected",
+    [
+        ("/lib/out", "/lib/out/file", True),
+        ("/lib/out", "file", True),
+        ("/lib/out", "dir/file", True),
+        ("/lib/out", "some/dir/file", True),
+        ("/lib/out", "some/dir/../file", True),
+        ("/lib/out", "some/dir/../../file", True),
+        ("/lib/out", "some/dir/../../../file", False),
+        ("/lib/out", "some/dir/../../../", False),
+        ("/lib/out", "some/dir/../../..", False),
+        ("/lib/out", "../file", False),
+        ("/lib/out", "/lib/out/../file", False),
+    ],
+)
+def test_is_safe_path(basedir, path, expected):
+    assert is_safe_path(Path(basedir), Path(path)) is expected
+
+
 @pytest.mark.parametrize(
     "size, alignment, result",
     [

unblob/extractor.py

@@ -5,7 +5,7 @@
 
 from structlog import get_logger
 
-from .file_utils import iterate_file
+from .file_utils import carve, is_safe_path
 from .models import Chunk, File, TaskResult, UnknownChunk, ValidChunk
 from .report import MaliciousSymlinkRemoved
 
@@ -14,12 +14,8 @@
 
 def carve_chunk_to_file(carve_path: Path, file: File, chunk: Chunk):
     """Extract valid chunk to a file, which we then pass to another tool to extract it."""
-    carve_path.parent.mkdir(parents=True, exist_ok=True)
     logger.debug("Carving chunk", path=carve_path)
-
-    with carve_path.open("xb") as f:
-        for data in iterate_file(file, chunk.start_offset, chunk.size):
-            f.write(data)
+    carve(carve_path, file, chunk.start_offset, chunk.size)
 
 
 def fix_permission(path: Path):
@@ -32,14 +28,6 @@ def fix_permission(path: Path):
         path.chmod(0o775)
 
 
-def is_safe_path(basedir: Path, path: Path) -> bool:
-    try:
-        basedir.joinpath(path).resolve().relative_to(basedir.resolve())
-    except ValueError:
-        return False
-    return True
-
-
 def is_recursive_link(path: Path) -> bool:
     try:
         path.resolve()

unblob/file_utils.py

@@ -15,6 +15,14 @@
 DEFAULT_BUFSIZE = shutil.COPY_BUFSIZE  # type: ignore
 
 
+def is_safe_path(basedir: Path, path: Path) -> bool:
+    try:
+        basedir.joinpath(path).resolve().relative_to(basedir.resolve())
+    except ValueError:
+        return False
+    return True
+
+
 class SeekError(ValueError):
     """Specific ValueError for File.seek."""
 
@@ -252,6 +260,15 @@ def iterate_file(
         yield data
 
 
+def carve(carve_path: Path, file: File, start_offset: int, size: int):
+    """Extract part of a file."""
+    carve_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with carve_path.open("xb") as f:
+        for data in iterate_file(file, start_offset, size):
+            f.write(data)
+
+
 def stream_scan(scanner, file: File):
     """Scan the whole file by increment of DEFAULT_BUFSIZE using Hyperscan's streaming mode."""
     scanner.scan(file, DEFAULT_BUFSIZE)