Merge pull request #695 from onekey-sec/skip-extension

Add ability to skip files based on extension.

tests/test_cli.py

@@ -302,3 +302,34 @@ def test_keep_extracted_chunks(
         process_file_mock.call_args.args[0].keep_extracted_chunks
         == keep_extracted_chunks
     ), fail_message
+
+
+@pytest.mark.parametrize(
+    "skip_extension, extracted_files_count",
+    [
+        pytest.param([], 5, id="skip-extension-empty"),
+        pytest.param([""], 5, id="skip-zip-extension-empty-suffix"),
+        pytest.param([".zip"], 1, id="skip-extension-zip"),
+        pytest.param([".rlib"], 5, id="skip-extension-rlib"),
+    ],
+)
+def test_skip_extension(
+    skip_extension: List[str], extracted_files_count: int, tmp_path: Path
+):
+    runner = CliRunner()
+    in_path = (
+        Path(__file__).parent
+        / "integration"
+        / "archive"
+        / "zip"
+        / "regular"
+        / "__input__"
+        / "apple.zip"
+    )
+    args = []
+    for suffix in skip_extension:
+        args += ["--skip-extension", suffix]
+    params = [*args, "--extract-dir", str(tmp_path), str(in_path)]
+    result = runner.invoke(unblob.cli.cli, params)
+    assert extracted_files_count == len(list(tmp_path.rglob("*")))
+    assert result.exit_code == 0

unblob/cli.py

@@ -22,6 +22,7 @@
 from .processing import (
     DEFAULT_DEPTH,
     DEFAULT_PROCESS_NUM,
+    DEFAULT_SKIP_EXTENSION,
     DEFAULT_SKIP_MAGIC,
     ExtractionConfig,
     process_file,
@@ -166,6 +167,15 @@ def __init__(
     show_default=True,
     multiple=True,
 )
+@click.option(
+    "--skip-extension",
+    "skip_extension",
+    type=click.STRING,
+    default=DEFAULT_SKIP_EXTENSION,
+    help="Skip processing files with given extension",
+    show_default=True,
+    multiple=True,
+)
 @click.option(
     "-p",
     "--process-num",
@@ -229,6 +239,7 @@ def cli(
     depth: int,
     entropy_depth: int,
     skip_magic: Iterable[str],
+    skip_extension: Iterable[str],
     skip_extraction: bool,  # noqa: FBT001
     keep_extracted_chunks: bool,  # noqa: FBT001
     handlers: Handlers,
@@ -254,6 +265,7 @@ def cli(
         entropy_plot=bool(verbose >= 3),
         skip_extraction=skip_extraction,
         skip_magic=skip_magic,
+        skip_extension=skip_extension,
         process_num=process_num,
         handlers=handlers,
         dir_handlers=dir_handlers,

unblob/processing.py

@@ -77,6 +77,7 @@
     "Windows Embedded CE binary image",
     "Intel serial flash for PCH ROM",
 )
+DEFAULT_SKIP_EXTENSION = (".rlib",)
 
 
 @attr.define(kw_only=True)
@@ -87,6 +88,7 @@ class ExtractionConfig:
     entropy_plot: bool = False
     max_depth: int = DEFAULT_DEPTH
     skip_magic: Iterable[str] = DEFAULT_SKIP_MAGIC
+    skip_extension: Iterable[str] = DEFAULT_SKIP_EXTENSION
     skip_extraction: bool = False
     process_num: int = DEFAULT_PROCESS_NUM
     keep_extracted_chunks: bool = False
@@ -292,9 +294,14 @@ def _process_task(self, result: TaskResult, task: Task):
         should_skip_file = any(
             magic.startswith(pattern) for pattern in self._config.skip_magic
         )
+        should_skip_file |= task.path.suffix in self._config.skip_extension
 
         if should_skip_file:
-            log.debug("Ignoring file based on magic", magic=magic)
+            log.debug(
+                "Ignoring file based on magic or extension.",
+                magic=magic,
+                extension=task.path.suffix,
+            )
             return
 
         _FileTask(self._config, task, stat_report.size, result).process()