-
Notifications
You must be signed in to change notification settings - Fork 1
/
__init__.py
145 lines (99 loc) · 3.87 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
from binaryninja.enums import MessageBoxIcon, LowLevelILOperation, MessageBoxButtonSet
from binaryninja.lowlevelil import LowLevelILOperationAndSize
from binaryninja.interaction import show_message_box, get_open_filename_input
from binaryninja.plugin import PluginCommand
from PySide2.QtWidgets import QApplication
from PySide2.QtCore import Qt
from binaryninjaui import DockHandler
from .dbutils import TraceDB, TraceDBError
from .view import XrefsDialog, BBViewerWidget
import sqlite3
# Global instance of the trace database
LOADED_DB = None
BRANCH_IL = [
LowLevelILOperation.LLIL_CALL,
LowLevelILOperation.LLIL_JUMP,
LowLevelILOperation.LLIL_JUMP_TO,
LowLevelILOperation.LLIL_GOTO,
LowLevelILOperation.LLIL_IF # More of a hack for cond jump detection
]
def _print_error(title, message):
show_message_box(title, message, icon=MessageBoxIcon.ErrorIcon)
def _is_branch(ins):
for subins in ins.prefix_operands:
if isinstance(subins, LowLevelILOperationAndSize) and subins.operation in BRANCH_IL:
return True
return False
def _get_llil_at(bv, addr):
""" Gets the low level il instruction at given vaddr """
fns = bv.get_functions_containing(addr)
if len(fns) == 0:
return None
ins = None
for fn in fns:
ins = fn.get_low_level_il_at(addr)
if ins:
return ins
return None
def db_required(func):
""" Decorator marking functions needing a database instance """
global LOADED_DB
def inner(*args, **kwargs):
if not LOADED_DB:
_print_error("Database Error", "No trace was loaded.Please open a trace database.")
else:
func(*args, **kwargs)
return inner
def db_load(bv):
global LOADED_DB
if LOADED_DB:
res = show_message_box(
"Database loading",
"A trace is already loaded. Do you want to overwrite it ?",
buttons=MessageBoxButtonSet.YesNoButtonSet,
icon=MessageBoxIcon.WarningIcon)
if not res:
return
LOADED_DB = None
path = get_open_filename_input("Select trace database")
if not path:
return
path = path.decode("UTF-8")
try:
LOADED_DB = TraceDB(bv, path)
text = ""
text += "Mappings: {:n}\n".format(LOADED_DB.mapping_count)
text += "Branches: {:n}\n".format(LOADED_DB.branch_count)
text += "bb hits : {:n}\n".format(LOADED_DB.hitcount_count)
show_message_box("Trace info", text)
except sqlite3.Error as e:
_print_error("Database error", "sqlite error: {}".format(e))
LOADED_DB = None
except TraceDBError as e:
_print_error("Database error", "Loading error: {}".format(e))
LOADED_DB = None
@db_required
def display_branch_xrefs(bv, addr):
global LOADED_DB
ins = _get_llil_at(bv, addr)
if not _is_branch(ins):
_print_error("Code error", "This instruction is not a branch")
return
dialog = XrefsDialog(addr, bv, LOADED_DB)
dialog.exec_()
@db_required
def display_bb_viewer(bv):
dock_handler = None
for wg in QApplication.allWidgets():
wg_win = wg.window()
dock_handler = wg_win.findChild(DockHandler, "__DockHandler")
if dock_handler:
break
if dock_handler is None:
print("Could not find DockHandler")
return
dock_widget = BBViewerWidget("Basic Block viewer", dock_handler.parent(), bv, LOADED_DB)
dock_handler.addDockWidget(dock_widget, Qt.RightDockWidgetArea, Qt.Vertical, True)
PluginCommand.register("Wakare\\Load trace database", "Loads a trace database", db_load)
PluginCommand.register_for_address("Wakare\\Branch xrefs from", "Displays xrefs from a branch", display_branch_xrefs)
PluginCommand.register("Wakare\\Show basic block viewer", "Displays information about basic blocks", display_bb_viewer)