feat: print output directory path at the end of processing

onekey-sec · Nov 27, 2024 · 579e6a7 · 579e6a7
1 parent 6db4508
commit 579e6a7
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 20 deletions.
diff --git a/unblob/cli.py b/unblob/cli.py
@@ -322,14 +322,25 @@ def cli(
 
     logger.info("Start processing file", file=file)
     process_results = process_file(config, file, report_file)
+
     if verbose == 0:
         if skip_extraction:
             print_scan_report(process_results)
         else:
             print_report(process_results)
+        print_output_dir(process_results)
     return process_results
 
 
+def print_output_dir(process_results: ProcessResult):
+    path = process_results.get_output_dir()
+
+    if path:
+        click.echo(f"Find output at {path}")
+    else:
+        click.echo("Too bad! Nothing was extracted.")
+
+
 cli.context_class = UnblobContext
 
 
@@ -453,35 +464,36 @@ def print_report(reports: ProcessResult):
             valid_size += size
         total_size += size
 
-    if total_size == 0:
-        return
+    console = Console()
 
-    summary = Panel(
-        f"""Extracted files: [#00FFC8]{total_files}[/#00FFC8]
+    if total_size:
+        summary = Panel(
+            f"""Extracted files: [#00FFC8]{total_files}[/#00FFC8]
 Extracted directories: [#00FFC8]{total_dirs}[/#00FFC8]
 Extracted links: [#00FFC8]{total_links}[/#00FFC8]
 Extraction directory size: [#00FFC8]{human_size(extracted_size)}[/#00FFC8]
 Chunks identification ratio: [#00FFC8]{(valid_size/total_size) * 100:0.2f}%[/#00FFC8]""",
-        subtitle="Summary",
-        title=f"unblob ({get_version()})",
-    )
+            subtitle="Summary",
+            title=f"unblob ({get_version()})",
+        )
 
-    console = Console()
-    console.print(summary)
+        console.print(summary)
 
-    chunks_table = Table(title="Chunks distribution")
-    chunks_table.add_column("Chunk type", justify="left", style="#00FFC8", no_wrap=True)
-    chunks_table.add_column("Size", justify="center", style="#00FFC8", no_wrap=True)
-    chunks_table.add_column("Ratio", justify="center", style="#00FFC8", no_wrap=True)
+        # fmt: off
+        chunks_table = Table(title="Chunks distribution")
+        chunks_table.add_column("Chunk type", justify="left", style="#00FFC8", no_wrap=True)
+        chunks_table.add_column("Size", justify="center", style="#00FFC8", no_wrap=True)
+        chunks_table.add_column("Ratio", justify="center", style="#00FFC8", no_wrap=True)
+        # fmt: on
 
-    for handler, size in sorted(
-        chunks_distribution.items(), key=lambda item: item[1], reverse=True
-    ):
-        chunks_table.add_row(
-            handler.upper(), human_size(size), f"{(size/total_size) * 100:0.2f}%"
-        )
+        for handler, size in sorted(
+            chunks_distribution.items(), key=lambda item: item[1], reverse=True
+        ):
+            chunks_table.add_row(
+                handler.upper(), human_size(size), f"{(size/total_size) * 100:0.2f}%"
+            )
 
-    console.print(chunks_table)
+        console.print(chunks_table)
 
     if len(reports.errors):
         errors_table = Table(title="Encountered errors")

diff --git a/unblob/models.py b/unblob/models.py
@@ -13,6 +13,7 @@
 from .identifiers import new_id
 from .parser import hexstring2regex
 from .report import (
+    CarveDirectoryReport,
     ChunkReport,
     ErrorReport,
     MultiFileReport,
@@ -238,6 +239,20 @@ def register(self, result: TaskResult):
     def to_json(self, indent="  "):
         return to_json(self.results, indent=indent)
 
+    def get_output_dir(self) -> Optional[Path]:
+        try:
+            top_result = self.results[0]
+            if carves := top_result.filter_reports(CarveDirectoryReport):
+                # we have a top level carve
+                return carves[0].carve_dir
+
+            # we either have an extraction,
+            # and the extract directory registered as subtask
+            return top_result.subtasks[0].path
+        except IndexError:
+            # or no extraction
+            return None
+
 
 class _JSONEncoder(json.JSONEncoder):
     def default(self, obj):