diff --git a/__init__.py b/__init__.py index 4516b8b..534e1c8 100644 --- a/__init__.py +++ b/__init__.py @@ -25,9 +25,49 @@ from .menus import * from .helpers import * from .slice_types import * +from .api import * from binaryninja import Settings from binaryninjaui import ViewType +from binaryninja.scriptingprovider import PythonScriptingProvider, PythonScriptingInstance + Settings().register_group("tanto", "Tanto Settings") ViewType.registerViewType(tanto.tanto_view.TantoViewType()) + + +def _get_current_tanto_view(instance: PythonScriptingInstance): + view_frame = instance.interpreter.locals["current_ui_view_frame"] + current_view = instance.interpreter.locals["current_view"] + if view_frame != None and current_view != None: + view = view_frame.getViewForType(f"Tanto:{current_view.view_type}") + if view != None: + return TantoApiView(view) + return None + + +PythonScriptingProvider.register_magic_variable( + "current_tanto_view", + _get_current_tanto_view, + depends_on=[ + "current_view", + "current_ui_view", + ] +) + + +def _get_current_tanto_slice(instance: PythonScriptingInstance): + tv = instance.interpreter.locals["current_tanto_view"] + ts = tv.parent.current_slice + if ts != None: + return TantoApiSlice(ts) + + +PythonScriptingProvider.register_magic_variable( + "current_tanto_slice", + _get_current_tanto_slice, + depends_on=[ + "current_tanto_view", + ] +) + diff --git a/api.py b/api.py new file mode 100644 index 0000000..38c4bd3 --- /dev/null +++ b/api.py @@ -0,0 +1,50 @@ +import tanto + +import binaryninja +from binaryninja import BinaryView, Variable + +from typing import Callable, Optional, Tuple + + +class TantoApiSliceAction: + + def __init__(self, name: str, action: Callable[..., None], validate: Optional[Callable[..., bool]] = None): + self._name = name + self._action = action + self._validate = validate + + def __call__(self, *args) -> bool: + if self._validate and not self._validate(*args): + return False + self._action(*args) + return True + + +class TantoApiSlice: + + def __init__(self, parent: 'tanto.slices.Slice'): + self.parent = parent + self._actions = {n: TantoApiSliceAction(n, *self.parent.actions[n]) for n in self.parent.actions} + + def __call__(self, action: str, *args): + return self._actions[action](*args) + + @property + def actions(self) -> list: + return list(self._actions.keys()) + + @property + def flowgraph(self) -> 'binaryninja.flowgraph.FlowGraph': + return self.parent.get_flowgraph() + + +class TantoApiView: + + def __init__(self, parent: 'tanto.tanto_view.TantoView'): + self.parent = parent + + def get_slice_by_name(self, name: str) -> Optional[TantoApiSlice]: + for _slicer_name, slice_name, slicer in self.parent.slices: + if slice_name == name: + return TantoApiSlice(slicer) + return None diff --git a/helpers.py b/helpers.py index c80056c..5df5ec4 100644 --- a/helpers.py +++ b/helpers.py @@ -20,14 +20,14 @@ from binaryninjaui import UIContext -from binaryninja import BinaryView +from binaryninja import BinaryView, Variable from binaryninja import Function, LowLevelILFunction, MediumLevelILFunction, HighLevelILFunction from binaryninja import BasicBlock, LowLevelILBasicBlock, MediumLevelILBasicBlock, HighLevelILBasicBlock from binaryninja import LowLevelILInstruction, MediumLevelILInstruction, HighLevelILInstruction from binaryninja.log import log_error from binaryninja.enums import FunctionGraphType -from typing import Union, Optional +from typing import Union, Optional, Callable BN_INVALID_EXPR = 0xffffffffffffffff @@ -38,6 +38,17 @@ ILBasicBlock = Union[LowLevelILBasicBlock, MediumLevelILBasicBlock, HighLevelILBasicBlock] ILInstruction = Union[LowLevelILInstruction, MediumLevelILInstruction, HighLevelILInstruction] +BinaryViewAction = Callable[[BinaryView], None] +BinaryViewValidator = Callable[[BinaryView], None] +FunctionAction = Callable[[BinaryView, AnyFunction], None] +FunctionValidator = Callable[[BinaryView, AnyFunction], bool] +BasicBlockAction = Callable[[BinaryView, AnyBasicBlock], None] +BasicBlockValidator = Callable[[BinaryView, AnyBasicBlock], bool] +VariableAction = Callable[[BinaryView, Variable], None] +VariableValidator = Callable[[BinaryView, Variable], bool] +AddressAction = Callable[[BinaryView, int], None] +AddressValidator = Callable[[BinaryView, int], bool] + def get_disassembly_settings(): view_context = UIContext.activeContext() diff --git a/slice_types/basic_block_slice.py b/slice_types/basic_block_slice.py index 8245958..e1031d6 100644 --- a/slice_types/basic_block_slice.py +++ b/slice_types/basic_block_slice.py @@ -36,6 +36,7 @@ class BasicBlockSlice(Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.parent = parent self.bv = parent.bv self.flowgraph_widget = parent.flowgraph_widget @@ -46,10 +47,10 @@ def __init__(self, parent: 'tanto.tanto_view.TantoView'): self.included_blocks = [] self.func = None - parent.register_for_basic_block("Include Block in Slice", self.include_block, lambda bv, bb: bb in self.excluded_blocks or bb not in self.included_blocks, "TantoGroup0", 0) - parent.register_for_basic_block("Exclude Block from Slice", self.exclude_block, lambda bv, bb: bb in self.included_blocks or bb not in self.excluded_blocks, "TantoGroup0", 1) - parent.register_for_basic_block("Reset Block State", self.reset_block, lambda bv, bb: bb in self.included_blocks or bb in self.excluded_blocks, "TantoGroup1", 2) - parent.register_for_binary_view("Clear All Block States", self.clear, lambda bv: len(self.included_blocks) + len(self.excluded_blocks) > 0, "TantoGroup2", 3) + self.register_for_basic_block("Include Block in Slice", self.include_block, lambda bv, bb: bb in self.excluded_blocks or bb not in self.included_blocks, "TantoGroup0", 0) + self.register_for_basic_block("Exclude Block from Slice", self.exclude_block, lambda bv, bb: bb in self.included_blocks or bb not in self.excluded_blocks, "TantoGroup0", 1) + self.register_for_basic_block("Reset Block State", self.reset_block, lambda bv, bb: bb in self.included_blocks or bb in self.excluded_blocks, "TantoGroup1", 2) + self.register_for_binary_view("Clear All Block States", self.clear, lambda bv: len(self.included_blocks) + len(self.excluded_blocks) > 0, "TantoGroup2", 3) def helperPaintEvent(self, event): p = QPainter(self.flowgraph_widget.viewport()) diff --git a/slice_types/butterfly_slice.py b/slice_types/butterfly_slice.py index 7fbc6b9..8448de6 100644 --- a/slice_types/butterfly_slice.py +++ b/slice_types/butterfly_slice.py @@ -28,6 +28,7 @@ class ButterflySlice(tanto.slices.Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.navigation_style = tanto.slices.NavigationStyle.FUNCTION_START self.update_style = tanto.slices.UpdateStyle.ON_NAVIGATE diff --git a/slice_types/dynamic_call_graph.py b/slice_types/dynamic_call_graph.py index 9fba622..8a5508f 100644 --- a/slice_types/dynamic_call_graph.py +++ b/slice_types/dynamic_call_graph.py @@ -31,6 +31,7 @@ class DynamicCallGraph(tanto.slices.Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.parent = parent self.bv = parent.bv self.graph_funcs = set() @@ -45,29 +46,32 @@ def __init__(self, parent: 'tanto.tanto_view.TantoView'): parent.flowgraph_widget.dragEnterEvent = self.dragEnterEvent parent.flowgraph_widget.dropEvent = self.dropEvent - parent.register_for_function("Include Function in Graph", + self.register_for_function("Include Function in Graph", self.include_function, lambda bv, func: self._to_source(func) not in self.graph_funcs, menu_group="TantoGroup0", menu_order=0) - parent.register_for_function("Exclude Function from Graph", + self.register_for_function("Exclude Function from Graph", self.exclude_function, lambda bv, func: self._to_source(func) in self.graph_funcs and func not in self.excluded_funcs, menu_group="TantoGroup0", menu_order=1) - parent.register_for_function("Include Callers", + self.register_for_function("Include Callers", self.include_callers, lambda bv, func: self._to_source(func) in self.graph_funcs, menu_group="TantoGroup1", menu_order=0) - parent.register_for_function("Include Callees", + + self.register_for_function("Include Callees", self.include_callees, lambda bv, func: self._to_source(func) in self.graph_funcs, menu_group="TantoGroup1", menu_order=1) - parent.register_for_function("Include All Callers", + + self.register_for_function("Include All Callers", self.include_all_callers, lambda bv, func: self._to_source(func) in self.graph_funcs, menu_group="TantoGroup2", menu_order=0) - parent.register_for_function("Include All Callees", + + self.register_for_function("Include All Callees", self.include_all_callees, lambda bv, func: self._to_source(func) in self.graph_funcs, menu_group="TantoGroup2", menu_order=1) diff --git a/slice_types/full_call_graph.py b/slice_types/full_call_graph.py index b372def..3856f11 100644 --- a/slice_types/full_call_graph.py +++ b/slice_types/full_call_graph.py @@ -34,6 +34,7 @@ class FullCallGraph(tanto.slices.Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.bv = parent.bv self.navigation_style = tanto.slices.NavigationStyle.FUNCTION_START diff --git a/slice_types/scatter_slice.py b/slice_types/scatter_slice.py index b87e521..9723bbe 100644 --- a/slice_types/scatter_slice.py +++ b/slice_types/scatter_slice.py @@ -28,6 +28,7 @@ class ScatterSlice(tanto.slices.Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.navigation_style = tanto.slices.NavigationStyle.FUNCTION_START self.update_style = tanto.slices.UpdateStyle.ON_NAVIGATE diff --git a/slice_types/source_to_sink.py b/slice_types/source_to_sink.py index 9b7f52c..38aaf8b 100644 --- a/slice_types/source_to_sink.py +++ b/slice_types/source_to_sink.py @@ -35,10 +35,10 @@ def __init__(self, parent: 'tanto.tanto_view.TantoView'): self.navigation_style = tanto.slices.NavigationStyle.FUNCTION_START - parent.register_for_function("Add Function as Source", self.add_function_as_source, menu_group="TantoGroup0", menu_order=0) - parent.register_for_function("Add Function as Sink", self.add_function_as_sink, menu_group="TantoGroup0", menu_order=1) - parent.register_for_function("Reset Function State", self.reset_function_state, menu_group="TantoGroup1", menu_order=2) - parent.register_for_binary_view("Clear All", self.clear_all, menu_group="TantoGroup2", menu_order=3) + self.register_for_function("Add Function as Source", self.add_function_as_source, menu_group="TantoGroup0", menu_order=0) + self.register_for_function("Add Function as Sink", self.add_function_as_sink, menu_group="TantoGroup0", menu_order=1) + self.register_for_function("Reset Function State", self.reset_function_state, menu_group="TantoGroup1", menu_order=2) + self.register_for_binary_view("Clear All", self.clear_all, menu_group="TantoGroup2", menu_order=3) def get_il_view_type(self) -> FunctionViewType: return FunctionViewType(FunctionGraphType.NormalFunctionGraph) diff --git a/slice_types/variable_slice.py b/slice_types/variable_slice.py index 3432633..d6ee15d 100644 --- a/slice_types/variable_slice.py +++ b/slice_types/variable_slice.py @@ -34,6 +34,7 @@ class VariableBlockSlice(Slice): def __init__(self, parent: 'tanto.tanto_view.TantoView'): + super().__init__() self.parent = parent self.bv = parent.bv self.flowgraph_widget = self.parent.flowgraph_widget @@ -43,9 +44,9 @@ def __init__(self, parent: 'tanto.tanto_view.TantoView'): self.variables = [] self.func = None - parent.register_for_variable("Include Variable in Slice", self.include_variable, menu_group="TantoGroup0", menu_order=0) - parent.register_for_variable("Remove Variable from Slice", self.remove_variable, menu_group="TantoGroup0", menu_order=1) - parent.register_for_binary_view("Remove All Variables", self.clear, lambda bv: len(self.variables) > 0, "TantoGroup1", 2) + self.register_for_variable("Include Variable in Slice", self.include_variable, menu_group="TantoGroup0", menu_order=0) + self.register_for_variable("Remove Variable from Slice", self.remove_variable, menu_group="TantoGroup0", menu_order=1) + self.register_for_binary_view("Remove All Variables", self.clear, lambda bv: len(self.variables) > 0, "TantoGroup1", 2) def helperPaintEvent(self, event): p = QPainter(self.flowgraph_widget.viewport()) diff --git a/slices.py b/slices.py index 250dcf1..8f58820 100644 --- a/slices.py +++ b/slices.py @@ -18,10 +18,16 @@ # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS # IN THE SOFTWARE. +from typing import Callable, Optional, Tuple from enum import Enum from binaryninja import FlowGraph, FunctionViewType +from binaryninja.binaryview import BinaryView +from binaryninja.variable import Variable from binaryninja.enums import FunctionGraphType +from binaryninja.log import log_warn + +import tanto ######### @@ -53,6 +59,9 @@ class UpdateStyle(Enum): class Slice(): + def __init__(self): + self.actions: dict[str, Tuple[Callable[..., None], Optional[Callable[..., bool]]]] = {} + # This is really the only function you need to implement def get_flowgraph(self) -> FlowGraph: raise NotImplementedError @@ -82,3 +91,46 @@ def update_style(self): @update_style.setter def update_style(self, value): self._update_style = value + + def register_action(self, name: str, action: Callable[..., None], is_valid: Optional[Callable[..., bool]] = None): + self.actions[name] = (action, is_valid) + + def register_for_binary_view(self, name: str, + action: Callable[['BinaryView'], None], + is_valid: Optional[Callable[['BinaryView'], bool]] = None, + menu_group: str = "", menu_order: int = 0, api_accessible: bool = True): + if api_accessible: + self.register_action(name, action, is_valid) + return self.parent.register_for_binary_view(name, action, is_valid, menu_group, menu_order) + + def register_for_function(self, name: str, + action: Callable[['BinaryView', 'tanto.helpers.AnyFunction'], None], + is_valid: Optional[Callable[['BinaryView', 'tanto.helpers.AnyFunction'], bool]] = None, + menu_group: str = "", menu_order: int = 0, api_accessible: bool = True): + if api_accessible: + self.register_action(name, action, is_valid) + return self.parent.register_for_function(name, action, is_valid, menu_group, menu_order) + + def register_for_basic_block(self, name: str, + action: Callable[['BinaryView', 'tanto.helpers.AnyBasicBlock'], None], + is_valid: Optional[Callable[['BinaryView', 'tanto.helpers.AnyBasicBlock'], bool]] = None, + menu_group: str = "", menu_order: int = 0, api_accessible: bool = True): + if api_accessible: + self.register_action(name, action, is_valid) + return self.parent.register_for_basic_block(name, action, is_valid, menu_group, menu_order) + + def register_for_variable(self, name: str, + action: Callable[['BinaryView', Variable], None], + is_valid: Optional[Callable[['BinaryView', Variable], bool]] = None, + menu_group: str = "", menu_order: int = 0, api_accessible: bool = True): + if api_accessible: + self.register_action(name, action, is_valid) + return self.parent.register_for_variable(name, action, is_valid, menu_group, menu_order) + + def register_for_address(self, name: str, + action: Callable[['BinaryView', int], None], + is_valid: Optional[Callable[['BinaryView', int], bool]] = None, + menu_group: str = "", menu_order: int = 0, api_accessible: bool = True): + if api_accessible: + self.register_action(name, action, is_valid) + return self.parent.register_for_address(name, action, is_valid, menu_group, menu_order)