Merge pull request #301 from IoT-Inspector/skip-processing

Skip processing

Dockerfile

@@ -19,7 +19,8 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     p7zip-full \
     unar \
     xz-utils \
-    zlib1g-dev
+    zlib1g-dev \
+    libmagic1
 RUN curl -L -o sasquatch_1.0_amd64.deb https://github.com/IoT-Inspector/sasquatch/releases/download/sasquatch-v1.0/sasquatch_1.0_amd64.deb \
     && dpkg -i sasquatch_1.0_amd64.deb \
     && rm -f sasquatch_1.0_amd64.deb

default.nix

@@ -13,6 +13,7 @@
 , sasquatch
 , simg2img
 , unar
+, file
 }:
 
 let
@@ -55,6 +56,12 @@ let
           self.python-lzo
         ];
       });
+
+      file-magic = super.file-magic.overridePythonAttrs (_: {
+        patchPhase = ''
+          substituteInPlace magic.py --replace "find_library('magic')" "'${file}/lib/libmagic.so'"
+        '';
+      });
     });
 
     python = python3;

pyproject.toml

@@ -26,6 +26,7 @@ jefferson = { git = "https://github.com/IoT-Inspector/jefferson.git", rev = "1a0
 yaffshiv = { git = "https://github.com/IoT-Inspector/yaffshiv.git", rev = "24e6e453a36a02144ae2d159eb3229f9c6312828" }
 plotext = "^4.1.5"
 pluggy = "^1.0.0"
+file-magic = "^0.4.0"
 
 
 [tool.poetry.dev-dependencies]

unblob/cli.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import sys
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 
 import click
 from structlog import get_logger
@@ -16,6 +16,7 @@
 from .processing import (
     DEFAULT_DEPTH,
     DEFAULT_PROCESS_NUM,
+    DEFAULT_SKIP_MAGIC,
     ExtractionConfig,
     process_file,
 )
@@ -116,6 +117,16 @@ def __init__(
     help="Load plugins from the provided path.",
     show_default=True,
 )
+@click.option(
+    "-S",
+    "--skip-magic",
+    "skip_magic",
+    type=click.STRING,
+    default=DEFAULT_SKIP_MAGIC,
+    help="Skip processing files with given magic prefix",
+    show_default=True,
+    multiple=True,
+)
 @click.option(
     "-p",
     "--process-num",
@@ -146,6 +157,7 @@ def cli(
     extract_root: Path,
     depth: int,
     entropy_depth: int,
+    skip_magic: Iterable[str],
     process_num: int,
     keep_extracted_chunks: bool,
     verbose: int,
@@ -164,6 +176,7 @@ def cli(
         max_depth=depth,
         entropy_depth=entropy_depth,
         entropy_plot=bool(verbose >= 3),
+        skip_magic=skip_magic,
         process_num=process_num,
         handlers=handlers,
         keep_extracted_chunks=keep_extracted_chunks,

unblob/processing.py

@@ -3,9 +3,10 @@
 import statistics
 from operator import attrgetter
 from pathlib import Path
-from typing import List
+from typing import Iterable, List
 
 import attr
+import magic
 import plotext as plt
 from structlog import get_logger
 
@@ -26,6 +27,7 @@
 
 DEFAULT_DEPTH = 10
 DEFAULT_PROCESS_NUM = multiprocessing.cpu_count()
+DEFAULT_SKIP_MAGIC = ("ELF", "JPEG", "GIF", "PNG")
 
 
 @attr.define(kw_only=True)
@@ -34,6 +36,7 @@ class ExtractionConfig:
     entropy_depth: int
     entropy_plot: bool = False
     max_depth: int = DEFAULT_DEPTH
+    skip_magic: Iterable[str] = DEFAULT_SKIP_MAGIC
     process_num: int = DEFAULT_PROCESS_NUM
     keep_extracted_chunks: bool = False
     extract_suffix: str = "_extract"
@@ -122,11 +125,20 @@ def _process_task(self, result: TaskResult, task: Task):
             log.debug("Ignoring empty file")
             return
 
-        filetask = _FileTask(self._config, task, size, result)
-        filetask.process()
-        # ensure that the root extraction directory is created even for empty extractions
-        if task.depth == 0:
-            filetask.extract_dir.mkdir(parents=True, exist_ok=True)
+        elif self._should_skip_magic(task):
+            log.debug("Ignoring file based on magic")
+            return
+
+        _FileTask(self._config, task, size, result).process()
+
+    def _should_skip_magic(self, task: Task):
+        detect = magic.detect_from_filename(task.path)
+        logger.debug(
+            "Detected file-magic", magic=detect.name, path=task.path, _verbosity=2
+        )
+        return any(
+            detect.name.startswith(pattern) for pattern in self._config.skip_magic
+        )
 
 
 class _FileTask:
@@ -168,6 +180,8 @@ def process(self):
                 # calculate entropy for whole files which produced no valid chunks
                 self._calculate_entropies([self.task.path])
 
+        self._ensure_root_extract_dir()
+
     def _process_chunks(
         self, file, outer_chunks: List[ValidChunk], unknown_chunks: List[UnknownChunk]
     ):
@@ -179,6 +193,11 @@ def _process_chunks(
         for chunk in outer_chunks:
             self._extract_chunk(file, chunk)
 
+    def _ensure_root_extract_dir(self):
+        # ensure that the root extraction directory is created even for empty extractions
+        if self.task.depth == 0:
+            self.extract_dir.mkdir(parents=True, exist_ok=True)
+
     def _calculate_entropies(self, paths: List[Path]):
         if self.task.depth < self.config.entropy_depth:
             for path in paths: