Merge pull request #627 from onekey-sec/626-handle-dirhandler-exception

fix(processing): handle when DirectoryHandler.calculate_multifile throws exception

Author: qkaiser

tests/test_processing.py

@@ -507,6 +507,16 @@ def calculate_multifile(self, file: Path) -> Optional[MultiFile]:
         )
 
 
+class ExceptionDirHandler(SplitDirHandler):
+    NAME = "exception-handler"
+    PATTERN = Glob("*.part0")
+    EXTRACTOR = None
+
+    def calculate_multifile(self, file: Path) -> Optional[MultiFile]:
+        del file
+        raise ValueError("Something bad happened")
+
+
 class DummyTestHandler(Handler):
     NAME = "dummy"
     PATTERNS = [Regex("AA")]
@@ -775,3 +785,27 @@ def test_multi_file_extraction_failed(
         for report in multi_file_reports
         for extraction_report in report.extraction_reports
     )
+
+
+def test_multi_file_calculate_exception(
+    multi_volume_zip: Path,
+    multi_file_extraction_config: ExtractionConfig,
+    extraction_root: Path,
+):
+    multi_file_extraction_config.dir_handlers = (ExceptionDirHandler,)
+    multi_file_extraction_config.handlers = (handlers.archive.zip.ZIPHandler,)
+
+    process_result = process_file(multi_file_extraction_config, multi_volume_zip)
+
+    task_result_by_path = {r.task.path: r for r in process_result.results}
+
+    directory = extraction_root / "input_extract"
+
+    multi_file_reports = task_result_by_path[directory].filter_reports(MultiFileReport)
+    assert not multi_file_reports
+    assert (directory / "test.part0") in task_result_by_path
+    assert (directory / "test.part1") in task_result_by_path
+    assert (directory / "test.part2") in task_result_by_path
+    assert (directory / "test2.part0") in task_result_by_path
+    assert (directory / "test2.part1") in task_result_by_path
+    assert (directory / "other") in task_result_by_path

unblob/processing.py

@@ -13,12 +13,13 @@
 from unblob.handlers import BUILTIN_DIR_HANDLERS, BUILTIN_HANDLERS, Handlers
 
 from .extractor import carve_unknown_chunk, carve_valid_chunk, fix_extracted_directory
-from .file_utils import iterate_file
+from .file_utils import InvalidInputFormat, iterate_file
 from .finder import search_chunks
 from .iter_utils import pairwise
 from .logging import noformat
 from .models import (
     Chunk,
+    DirectoryHandler,
     DirectoryHandlers,
     ExtractError,
     File,
@@ -31,6 +32,7 @@
 )
 from .pool import make_pool
 from .report import (
+    CalculateMultiFileExceptionReport,
     EntropyReport,
     ExtractDirectoryExistsReport,
     FileMagicReport,
@@ -324,7 +326,7 @@ def _process_directory(self) -> Tuple[Set[Path], Set[Path]]:
             dir_handler = dir_handler_class()
 
             for path in dir_handler.PATTERN.get_files(self.dir_task.path):
-                multi_file = dir_handler.calculate_multifile(path)
+                multi_file = self._calculate_multifile(dir_handler, path, self.result)
 
                 if multi_file is None:
                     continue
@@ -349,6 +351,32 @@ def _process_directory(self) -> Tuple[Set[Path], Set[Path]]:
                 processed_paths.update(multi_file.paths)
         return processed_paths, extract_dirs
 
+    @staticmethod
+    def _calculate_multifile(
+        dir_handler: DirectoryHandler, path: Path, task_result: TaskResult
+    ) -> Optional[MultiFile]:
+        try:
+            return dir_handler.calculate_multifile(path)
+        except InvalidInputFormat as exc:
+            logger.debug(
+                "Invalid MultiFile format",
+                exc_info=exc,
+                handler=dir_handler.NAME,
+                path=path,
+                _verbosity=2,
+            )
+        except Exception as exc:
+            error_report = CalculateMultiFileExceptionReport(
+                handler=dir_handler.NAME,
+                exception=exc,
+                path=path,
+            )
+            task_result.add_report(error_report)
+            logger.warning(
+                "Unhandled Exception during multi file calculation",
+                **error_report.asdict(),
+            )
+
     def _check_conflicting_files(
         self, multi_file: MultiFile, processed_paths: Set[Path]
     ):

unblob/report.py

@@ -78,6 +78,15 @@ class CalculateChunkExceptionReport(UnknownError):
     handler: str
 
 
+@attr.define(kw_only=True, frozen=True)
+class CalculateMultiFileExceptionReport(UnknownError):
+    """Describes an exception raised during calculate_chunk execution."""
+
+    path: Path
+    # Stored in `str` rather than `Handler`, because the pickle picks ups structs from `C_DEFINITIONS`
+    handler: str
+
+
 @attr.define(kw_only=True, frozen=True)
 class ExtractCommandFailedReport(ErrorReport):
     """Describes an error when failed to run the extraction command."""