Merge pull request #646 from onekey-sec/645-fix-permissions

fix(extractor): improve post-extraction permission fixing.

tests/test_extractor.py

@@ -1,8 +1,11 @@
+import itertools
 from pathlib import Path, PosixPath
 
 import pytest
 
 from unblob.extractor import (
+    DIR_PERMISSION_MASK,
+    FILE_PERMISSION_MASK,
     carve_unknown_chunk,
     fix_extracted_directory,
     fix_permission,
@@ -31,15 +34,21 @@ def test_carve_unknown_chunk(tmp_path: Path):
 
 
 def test_fix_permission(tmpdir: Path):
-    tmpdir = PosixPath(tmpdir)
+    tmpdir = PosixPath(tmpdir / "dir")
     tmpfile = PosixPath(tmpdir / "file.txt")
-    tmpfile.touch()
-    tmpdir.chmod(0o777)
-    tmpfile.chmod(0o777)
-    fix_permission(tmpdir)
-    fix_permission(tmpfile)
-    assert (tmpdir.stat().st_mode & 0o777) == 0o775
-    assert (tmpfile.stat().st_mode & 0o777) == 0o644
+
+    for user, group, others in itertools.product(range(8), repeat=3):
+        permission = (user << 6) + (group << 3) + others
+        tmpdir.mkdir()
+        tmpfile.touch()
+        tmpfile.chmod(permission)
+        tmpdir.chmod(permission)
+        fix_permission(tmpdir)
+        fix_permission(tmpfile)
+        assert (tmpdir.stat().st_mode & 0o777) == permission | DIR_PERMISSION_MASK
+        assert (tmpfile.stat().st_mode & 0o777) == permission | FILE_PERMISSION_MASK
+        tmpfile.unlink()
+        tmpdir.rmdir()
 
 
 def test_fix_extracted_directory(tmpdir: Path, task_result: TaskResult):

unblob/extractor.py

@@ -11,6 +11,9 @@
 
 logger = get_logger()
 
+FILE_PERMISSION_MASK = 0o644
+DIR_PERMISSION_MASK = 0o775
+
 
 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."""
@@ -19,13 +22,20 @@ def carve_chunk_to_file(carve_path: Path, file: File, chunk: Chunk):
 
 
 def fix_permission(path: Path):
+    if not path.exists():
+        return
+
     if path.is_symlink():
         return
 
+    mode = path.stat().st_mode
+
     if path.is_file():
-        path.chmod(0o644)
+        mode |= FILE_PERMISSION_MASK
     elif path.is_dir():
-        path.chmod(0o775)
+        mode |= DIR_PERMISSION_MASK
+
+    path.chmod(mode)
 
 
 def is_recursive_link(path: Path) -> bool: