59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
from typing import Optional
|
|
|
|
from angr.analyses.decompiler.structured_codegen.c import CVariable
|
|
from angrmanagement.plugins import BasePlugin
|
|
from angrmanagement.ui.views import CodeView
|
|
|
|
from .analysis import TypeTapperAnalysis
|
|
from .data import Atom
|
|
from .knowledge import TypeTapperManager
|
|
from .hierarchy_graph_view import HierarchicalGraphView
|
|
|
|
|
|
class TypeTapper(BasePlugin):
|
|
def __init__(self, workspace):
|
|
super().__init__(workspace)
|
|
self.kp: Optional[TypeTapperManager] = None
|
|
|
|
MENU_BUTTONS = ["TypeTapper initial analysis"]
|
|
|
|
def handle_click_menu(self, idx):
|
|
assert idx == 0
|
|
cfg = self.workspace.main_instance.cfg
|
|
tt = self.workspace.main_instance.project.analyses[TypeTapperAnalysis](cfg)
|
|
self.kp = tt.manager
|
|
self._start(self.kp.lookup_reg(0x4144a5, "rdi"))
|
|
|
|
def build_context_menu_node(self, node):
|
|
# this is bad lol
|
|
code_view = self.workspace.view_manager.current_tab
|
|
if not isinstance(code_view, CodeView):
|
|
return
|
|
addr = code_view._textedit.get_closest_insaddr(node)
|
|
|
|
if self.kp is not None and isinstance(node, CVariable):
|
|
try:
|
|
atom = self.kp.lookup_variable(addr, node.variable)
|
|
except:
|
|
pass
|
|
else:
|
|
yield 'Start TypeTapper', lambda: self._start(atom)
|
|
|
|
def _start(self, node: Atom):
|
|
hg = self.kp.session(node)
|
|
start_relatoms = list(hg.frontier)
|
|
func = self.kp.atom_to_function(node)
|
|
func_group = hg.create_group(start_relatoms, hg.root_group)
|
|
hg.expand_while(start_relatoms, func_group, lambda atom: self.kp.atom_to_function(atom) == func)
|
|
newgroups = hg.expand_by_key(func_group, hg.root_group, self.kp.atom_to_function)
|
|
|
|
view = HierarchicalGraphView(
|
|
self.workspace.main_instance,
|
|
"center",
|
|
hg,
|
|
set(g for k, g in newgroups.items() if not self.kp.kb.functions[k].alignment),
|
|
set(g for k, g in newgroups.items() if self.kp.kb.functions[k].alignment),
|
|
)
|
|
self.workspace.add_view(view)
|
|
self.workspace.raise_view(view)
|