Merge pull request #1794 from volatilityfoundation/feature/volshell-breakpoints

Volshell breakpointing

volatility3/cli/volshell/generic.py

@@ -187,6 +187,9 @@ class Volshell(interfaces.plugins.PluginInterface):
     def construct_locals(self) -> List[Tuple[List[str], Any]]:
         """Returns a listing of the functions to be added to the environment."""
         return [
+            (["bc", "breakpoint_clear"], self.breakpoint_clear),
+            (["bl", "breakpoint_list"], self.breakpoint_list),
+            (["bp", "breakpoint"], self.breakpoint),
             (["dt", "display_type"], self.display_type),
             (["db", "display_bytes"], self.display_bytes),
             (["dw", "display_words"], self.display_words),
@@ -797,6 +800,89 @@ class Volshell(interfaces.plugins.PluginInterface):
 
         return constructed
 
+    def breakpoint(
+        self, address: int, layer_name: Optional[str] = None, lowest: bool = False
+    ) -> None:
+        """Sets a breakpoint on a particular address (within a specific layer)"""
+        if layer_name is None:
+            if self.current_layer is None:
+                raise ValueError("Current layer must be set")
+            layer_name = self.current_layer
+
+        layer: interfaces.layers.DataLayerInterface = self.context.layers[layer_name]
+
+        if lowest:
+            while isinstance(layer, interfaces.layers.TranslationLayerInterface):
+                mapping = layer.mapping(address, 1)
+                if not mapping:
+                    raise ValueError(
+                        "Offset cannot be mapped lower, cannot break at lowest layer"
+                    )
+                _, _, mapped_offset, _, mapped_layer_name = next(mapping)
+                layer = self.context.layers[mapped_layer_name]
+                address = mapped_offset
+
+        # Check if the read value is already overloaded
+        if not hasattr(layer.read, "breakpoints"):
+            # Layer read is not yet wrapped
+            def wrapped_read(offset: int, length: int, pad: bool = False) -> bytes:
+                original_read = getattr(wrapped_read, "original_read")
+                for breakpoint in getattr(wrapped_read, "breakpoints"):
+                    if (offset <= breakpoint) and (breakpoint < offset + length):
+                        print(
+                            "Hit breakpoint, entering python debugger. To continue running without the debugger use the command continue"
+                        )
+                        import pdb
+
+                        pdb.set_trace()
+                        _ = "First statement after the breakpoint, use u(p), d(own) and list to navigate through the execution frames"
+                return original_read(offset, length, pad)
+
+            setattr(wrapped_read, "breakpoints", set())
+            setattr(wrapped_read, "original_read", layer.read)
+            setattr(layer, "read", wrapped_read)
+
+        # Add the new breakpoint
+        print(f"Setting breakpoint {address:#x} on {layer.name}")
+        breakpoints = getattr(layer.read, "breakpoints")
+        breakpoints.add(address)
+        setattr(layer.read, "breakpoints", breakpoints)
+
+    def breakpoint_list(self, layer_names: Optional[List[str]] = None):
+        """List available breakpoints for a set of layers"""
+        if not layer_names:
+            layer_names = [layer_name for layer_name in self.context.layers]
+
+        print("Listing breakpoints:")
+        for layer_name in layer_names:
+            print(f" {layer_name}")
+            layer = self.context.layers.get(layer_name, None)
+            if layer and hasattr(layer.read, "breakpoints"):
+                for breakpoint in layer.read.breakpoints:
+                    print(f"  {breakpoint:#x}")
+
+    def breakpoint_clear(
+        self, offset: Optional[int] = None, layer_name: Optional[str] = None
+    ):
+        """Clears a offset breakpoint on a layer (or all breakpoints if offset or layer not specified)
+
+        Args:
+            offset: Address of the breakpoint to clear (or all if None)
+            layer_name: Layer to clear breakpoints from (or all if None)
+        """
+        print("Clearing breakpoints:")
+        for candidate_layer_name in self.context.layers:
+            candidate_layer = self.context.layers[candidate_layer_name]
+            if layer_name is None or layer_name == candidate_layer_name:
+                print(f" {candidate_layer_name}")
+                if hasattr(candidate_layer.read, "breakpoints"):
+                    breakpoints_to_remove = set()
+                    for breakpoint in candidate_layer.read.breakpoints:
+                        if offset is None or offset == breakpoint:
+                            print(f"  clearing {breakpoint:#x}")
+                            breakpoints_to_remove.add(breakpoint)
+                    candidate_layer.read.breakpoints -= breakpoints_to_remove
+
 
 class NullFileHandler(io.BytesIO, interfaces.plugins.FileHandlerInterface):
     """Null FileHandler that swallows files whole without consuming memory"""