Merge pull request #3370 from trailofbits/cli-refactor

CLI refactor and libmagic optimizations

Author: ESultanik

polyfile/__main__.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+import atexit
 from enum import Enum
 from pathlib import Path
 from pdb import Pdb
@@ -9,6 +10,7 @@
 from .magic import (
     AbsoluteOffset, FailedTest, InvalidOffsetError, MagicMatcher, MagicTest, Offset, TestResult, TEST_TYPES
 )
+from .profiling import Profiler, Unprofiled, unprofiled
 from .repl import ANSIColor, ANSIWriter, arg_completer, command, ExitREPL, log, REPL, SetCompleter
 from .wildcards import Wildcard
 
@@ -383,6 +385,20 @@ def value(self, new_value):
             self.debugger._instrument()
 
 
+class Profile(BooleanVariable):
+    def __init__(self, value: bool, debugger: "Debugger"):
+        self.debugger: Debugger = debugger
+        self._registered_callback = False
+        super().__init__(value)
+
+    @Variable.value.setter
+    def value(self, new_value):
+        Variable.value.fset(self, new_value)
+        if new_value and not self._registered_callback:
+            self._registered_callback = True
+            atexit.register(Debugger.profile_command, self.debugger, "")
+
+
 class Debugger(REPL):
     def __init__(self, break_on_parsing: bool = True):
         super().__init__(name="the debugger")
@@ -399,13 +415,17 @@ def __init__(self, break_on_parsing: bool = True):
         self.repl_test: Optional[MagicTest] = None
         self.instrumented_parsers: List[InstrumentedParser] = []
         self.break_on_submatching: BreakOnSubmatching = BreakOnSubmatching(break_on_parsing, self)
+        self.profile: Profile = Profile(False, self)
         self.variables_by_name: Dict[str, Variable] = {
-            "break_on_parsing": self.break_on_submatching
+            "break_on_parsing": self.break_on_submatching,
+            "profile": self.profile
         }
         self.variable_descriptions: Dict[str, str] = {
             "break_on_parsing": "Break when a PolyFile parser is about to be invoked and debug using PDB (default=True;"
-                                " disable from the command line with `--no-debug-python`)"
+                                " disable from the command line with `--no-debug-python`)",
+            "profile": "Profile the performance of each magic test that is run (default=False)"
         }
+        self.profile_results: Dict[Union[MagicTest, Type[Parser]], float] = {}
         self._pdb: Optional[Pdb] = None
 
     def save_context(self):
@@ -442,7 +462,7 @@ def _instrument(self):
             if "test" in test.__dict__:
                 # this class actually implements the test() function
                 self.instrumented_tests.append(InstrumentedTest(test, self))
-        if self.break_on_submatching:
+        if self.break_on_submatching.value:
             for parsers in PARSERS.values():
                 for parser in parsers:
                     self.instrumented_parsers.append(InstrumentedParser(parser, self))
@@ -506,6 +526,8 @@ def write_test(self, test: MagicTest, is_current_test: bool = False):
             indent = ""
             self.write(f"{'>' * test.level}{test.offset!s}\t")
             self.write(test.message, color=ANSIColor.BLUE, bold=True)
+        if self.profile.value and test in self.profile_results:
+            self.write(f"\t⏱  {int(self.profile_results[test] + 0.5)}ms")
         if test.mime is not None:
             self.write(f"\n  {indent}!:mime ", dim=True)
             self.write(test.mime, color=ANSIColor.BLUE)
@@ -541,10 +563,14 @@ def debug(
             absolute_offset: int,
             parent_match: Optional[TestResult]
     ) -> Optional[TestResult]:
-        if instrumented_test.original_test is None:
-            result = instrumented_test.test.test(test, data, absolute_offset, parent_match)
-        else:
-            result = instrumented_test.original_test(test, data, absolute_offset, parent_match)
+        profiler = Profiler()
+        with profiler:
+            if instrumented_test.original_test is None:
+                result = instrumented_test.test.test(test, data, absolute_offset, parent_match)
+            else:
+                result = instrumented_test.original_test(test, data, absolute_offset, parent_match)
+        if self.profile.value:
+            self.profile_results[test] = profiler.elapsed_ms
         if self.repl_test is test:
             # this is a one-off test run from the REPL, so do not save its results
             return result
@@ -557,6 +583,7 @@ def debug(
             self.repl()
         return self.last_result
 
+    @unprofiled
     def print_where(
             self,
             test: Optional[MagicTest] = None,
@@ -659,46 +686,59 @@ def print_location():
             # We are already debugging!
             print_location()
             self.write(f"Parsing for submatches using {instrumented_parser.parser!s}.\n")
-            yield from parse(file_stream, match)
+            profiler = Profiler()
+            with profiler:
+                yield from parse(file_stream, match)
+            if self.profile.value:
+                self.profile_results[instrumented_parser.parser] = profiler.elapsed_ms
             return
-        self.print_match(match)
-        print_location()
-        self.write(f"About to parse for submatches using {instrumented_parser.parser!s}.\n")
-        buffer = ANSIWriter(use_ansi=sys.stderr.isatty(), escape_for_readline=True)
-        buffer.write("Debug using PDB? ")
-        buffer.write("(disable this prompt with `", dim=True)
-        buffer.write("set ", color=ANSIColor.BLUE)
-        buffer.write("break_on_parsing ", color=ANSIColor.GREEN)
-        buffer.write("False", color=ANSIColor.CYAN)
-        buffer.write("`)", dim=True)
+        with Unprofiled():
+            self.print_match(match)
+            print_location()
+            self.write(f"About to parse for submatches using {instrumented_parser.parser!s}.\n")
+            buffer = ANSIWriter(use_ansi=sys.stderr.isatty(), escape_for_readline=True)
+            buffer.write("Debug using PDB? ")
+            buffer.write("(disable this prompt with `", dim=True)
+            buffer.write("set ", color=ANSIColor.BLUE)
+            buffer.write("break_on_parsing ", color=ANSIColor.GREEN)
+            buffer.write("False", color=ANSIColor.CYAN)
+            buffer.write("`)", dim=True)
         if not self.prompt(str(buffer), default=False):
-            yield from parse(file_stream, match)
+            with Profiler() as p:
+                yield from parse(file_stream, match)
+                if self.profile.value:
+                    self.profile_results[instrumented_parser.parser] = p.elapsed_ms
             return
         try:
-            self._pdb = Pdb(skip=["polyfile.magic_debugger", "polyfile.magic"])
-            if sys.stderr.isatty():
-                self._pdb.prompt = "\001\u001b[1m\002(polyfile-Pdb)\001\u001b[0m\002 "
-            else:
-                self._pdb.prompt = "(polyfile-Pdb) "
-            generator = parse(file_stream, match)
-            while True:
-                try:
-                    result = self._pdb.runcall(next, generator)
-                    self.write(f"Got a submatch:\n", dim=True)
-                    self.print_match(result)
-                    yield result
-                except StopIteration:
-                    self.write(f"Yielded all submatches from {match.__class__.__name__} at offset {match.offset}.\n")
-                    break
-                print_location()
-                if not self.prompt("Continue debugging the next submatch?", default=True):
-                    if self.prompt("Print the remaining submatches?", default=False):
-                        for result in generator:
-                            self.print_match(result)
-                            yield result
-                    else:
-                        yield from generator
-                    break
+            if self.profile.value:
+                self.write("Warning:", bold=True, color=ANSIColor.RED)
+                self.write(" Profiling will be disabled for this parser while debugging!\n")
+            with Unprofiled():
+                self._pdb = Pdb(skip=["polyfile.magic_debugger", "polyfile.magic"])
+                if sys.stderr.isatty():
+                    self._pdb.prompt = "\001\u001b[1m\002(polyfile-Pdb)\001\u001b[0m\002 "
+                else:
+                    self._pdb.prompt = "(polyfile-Pdb) "
+                generator = parse(file_stream, match)
+                while True:
+                    try:
+                        result = self._pdb.runcall(next, generator)
+                        self.write(f"Got a submatch:\n", dim=True)
+                        self.print_match(result)
+                        yield result
+                    except StopIteration:
+                        self.write(f"Yielded all submatches from {match.__class__.__name__} at offset {match.offset}."
+                                   f"\n")
+                        break
+                    print_location()
+                    if not self.prompt("Continue debugging the next submatch?", default=True):
+                        if self.prompt("Print the remaining submatches?", default=False):
+                            for result in generator:
+                                self.print_match(result)
+                                yield result
+                        else:
+                            yield from generator
+                        break
         finally:
             self._pdb = None
 
@@ -731,6 +771,50 @@ def next(self, arguments: str):
     def where(self, arguments: str):
         self.print_where()
 
+    @command(allows_abbreviation=False, name="profile", help="print current profiling results (to enable profiling, "
+                                                             "use `set profile True`)", )
+    def profile_command(self, arguments: str):
+        if not self.profile_results:
+            if not self.profile.value:
+                self.write("Profiling is disabled.\n", color=ANSIColor.RED)
+                self.write("Enable it by running `set profile True`.\n")
+            else:
+                self.write("No profiling data yet.\n")
+            return
+        self.write("Profile Results:\n", bold=True)
+        tests = sorted([(runtime, test) for test, runtime in self.profile_results.items()], reverse=True,
+                       key=lambda x: x[0])
+        max_text_width = 0
+        for runtime, test in tests:
+            if isinstance(test, MagicTest):
+                if test.source_info is not None and test.source_info.original_line is not None:
+                    max_text_width = max(max_text_width,
+                                         len(test.source_info.path.name) + 1 + len(str(test.source_info.line)))
+                else:
+                    max_text_width = max(max_text_width, test.level + len(str(test.offset)))
+            else:
+                max_text_width = max(max_text_width, len(str(test)))
+        for runtime, test in tests:
+            if isinstance(test, MagicTest):
+                self.write("🪄 ")
+                if test.source_info is not None and test.source_info.original_line is not None:
+                    self.write(test.source_info.path.name, dim=True, color=ANSIColor.CYAN)
+                    self.write(":", dim=True)
+                    self.write(test.source_info.line, dim=True, color=ANSIColor.CYAN)
+                    padding = max_text_width - (len(test.source_info.path.name) + 1 + len(str(test.source_info.line)))
+                else:
+                    self.write(f"{'>' * test.level}{test.offset!s}", color=ANSIColor.BLUE)
+                    padding = max_text_width - test.level - len(str(test.offset))
+            else:
+                self.write("🖥 ")
+                self.write(str(test), color=ANSIColor.BLUE)
+                padding = max_text_width - len(str(test))
+            self.write(" " * padding)
+            if runtime >= 1.0:
+                self.write(f" ⏱  {int(runtime + 0.5)}ms\n")
+            else:
+                self.write(f" ⏱  {runtime:.2f}ms\n")
+
     @command(allows_abbreviation=True, help="test the following libmagic DSL test at the current position")
     def test(self, args: str):
         if args:

polyfile/debugger.py

@@ -5,7 +5,7 @@
 import tempfile as tf
 import shutil
 import sys
-from typing import AnyStr, IO, Iterator, Iterable, List, Optional, Union
+from typing import AnyStr, ContextManager, IO, Iterator, Iterable, List, Optional, TextIO, Union
 
 
 def make_stream(path_or_stream, mode='rb', close_on_exit=None):
@@ -41,7 +41,7 @@ def __init__(self, contents: bytes, name: str):
         super().__init__(contents)
         self._name: str = name
 
-    def __enter__(self):
+    def __enter__(self) -> str:
         tmpdir = Path(tf.mkdtemp())
         file_path = tmpdir / self._name
         with open(file_path, "wb") as f:
@@ -56,16 +56,16 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 class PathOrStdin:
-    def __init__(self, path):
-        self._path = path
-        if self._path == '-':
-            self._tempfile = Tempfile(sys.stdin.buffer.read())
+    def __init__(self, path: str):
+        self.path: str = path
+        if self.path == '-':
+            self._tempfile: Optional[ExactNamedTempfile] = ExactNamedTempfile(sys.stdin.buffer.read(), "STDIN")
         else:
             self._tempfile = None
 
-    def __enter__(self):
+    def __enter__(self) -> str:
         if self._tempfile is None:
-            return self._path
+            return self.path
         else:
             return self._tempfile.__enter__()
 
@@ -74,6 +74,27 @@ def __exit__(self, *args, **kwargs):
             return self._tempfile.__exit__(*args, **kwargs)
 
 
+class PathOrStdout(ContextManager[TextIO]):
+    def __init__(self, path: str):
+        self.path: str = path
+        if self.path == '-':
+            self._tempfile: Optional[TextIO] = sys.stdout
+        else:
+            self._tempfile = None
+
+    def __enter__(self) -> TextIO:
+        if self._tempfile is None:
+            self._tempfile = open(self.path, "w")
+        return self._tempfile
+
+    def __exit__(self, *args, **kwargs) -> None:  # type: ignore
+        if self._tempfile is not None:
+            if self._tempfile is not sys.stdout:
+                self._tempfile.close()
+                self._tempfile = None
+        return None
+
+
 class FileStream(IO):
     def __init__(
             self,

polyfile/fileutils.py

@@ -93,12 +93,12 @@ def __init__(self, _io, _parent=None, _root=None):
             self._debug = collections.defaultdict(dict)
 
         def _read(self):
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32:
                 self._debug['image_base_32']['start'] = self._io.pos()
                 self.image_base_32 = self._io.read_u4le()
                 self._debug['image_base_32']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32_plus:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32_plus:
                 self._debug['image_base_64']['start'] = self._io.pos()
                 self.image_base_64 = self._io.read_u8le()
                 self._debug['image_base_64']['end'] = self._io.pos()
@@ -145,42 +145,42 @@ def _read(self):
             self._debug['dll_characteristics']['start'] = self._io.pos()
             self.dll_characteristics = self._io.read_u2le()
             self._debug['dll_characteristics']['end'] = self._io.pos()
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32:
                 self._debug['size_of_stack_reserve_32']['start'] = self._io.pos()
                 self.size_of_stack_reserve_32 = self._io.read_u4le()
                 self._debug['size_of_stack_reserve_32']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32_plus:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32_plus:
                 self._debug['size_of_stack_reserve_64']['start'] = self._io.pos()
                 self.size_of_stack_reserve_64 = self._io.read_u8le()
                 self._debug['size_of_stack_reserve_64']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32:
                 self._debug['size_of_stack_commit_32']['start'] = self._io.pos()
                 self.size_of_stack_commit_32 = self._io.read_u4le()
                 self._debug['size_of_stack_commit_32']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32_plus:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32_plus:
                 self._debug['size_of_stack_commit_64']['start'] = self._io.pos()
                 self.size_of_stack_commit_64 = self._io.read_u8le()
                 self._debug['size_of_stack_commit_64']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32:
                 self._debug['size_of_heap_reserve_32']['start'] = self._io.pos()
                 self.size_of_heap_reserve_32 = self._io.read_u4le()
                 self._debug['size_of_heap_reserve_32']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32_plus:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32_plus:
                 self._debug['size_of_heap_reserve_64']['start'] = self._io.pos()
                 self.size_of_heap_reserve_64 = self._io.read_u8le()
                 self._debug['size_of_heap_reserve_64']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32:
                 self._debug['size_of_heap_commit_32']['start'] = self._io.pos()
                 self.size_of_heap_commit_32 = self._io.read_u4le()
                 self._debug['size_of_heap_commit_32']['end'] = self._io.pos()
 
-            if self._parent.std.format == MicrosoftPe.PeFormat.pe32_plus:
+            if self._parent.std.output_format == MicrosoftPe.PeFormat.pe32_plus:
                 self._debug['size_of_heap_commit_64']['start'] = self._io.pos()
                 self.size_of_heap_commit_64 = self._io.read_u8le()
                 self._debug['size_of_heap_commit_64']['end'] = self._io.pos()