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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
|
from typing import List
from PySide6.QtCore import QSettings
from binaryninja import DisassemblyTextLine, LowLevelILOperation, DisassemblyTextRenderer, \
MediumLevelILOperation, \
RenderLayer, BasicBlock, InstructionTextTokenType, RenderLayerDefaultEnableState, \
PluginCommand, BinaryView, interaction, InstructionTextToken, RegisterValueType
from binaryninjaui import UIContext, UIActionHandler, UIAction, UIActionContext
"""
Render layer that adds annotations to follow the value of a register, for Disasm/LLIL.
- Adds "<reg> after: <value>" comments to every line that modifies your register
- Adds "<reg> at block entry: <value>" comments to every block
- Adds "<reg> at block exit: <value>" comments to every block
You can switch registers with the plugin command "Set Followed Register"
"""
class FollowRegRenderLayer(RenderLayer):
name = "Follow Register"
default_enable_state = RenderLayerDefaultEnableState.EnabledByDefaultRenderLayerDefaultEnableState
def __init__(self, handle=None):
super().__init__(handle)
self.followed_reg = QSettings().value("layer/followed_reg", None)
def apply_to_lines(
self,
block: BasicBlock,
lines: List['DisassemblyTextLine'],
):
if self.followed_reg is None:
return lines
if self.followed_reg not in block.arch.regs:
return lines
func = block.function
arch = block.arch
# Give our tokens a unique value so they don't highlight all the other comments
token_value = arch.get_reg_index(self.followed_reg) ^ 0x12345678
# ===============================================================================
# Render the "<reg> after" comments
for i, line in enumerate(lines):
# Ignore non-disassembly lines
if not any(token.type == InstructionTextTokenType.AddressSeparatorToken for token in line.tokens):
continue
# If the register has changed, we should render the update
written = False
for w in func.get_regs_written_by(line.address, arch):
# In case this instruction only modifies a sub register of our target
# Or in the case that we wanted to follow a sub register
if arch.regs[w].full_width_reg == arch.regs[self.followed_reg].full_width_reg:
written = True
break
if written:
# Add comment tokens to existing line
line.tokens.extend([
InstructionTextToken(InstructionTextTokenType.CommentToken, ' // ', token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, ' after: ', token_value),
])
after = func.get_reg_value_after(line.address, self.followed_reg, arch)
if after.type == RegisterValueType.UndeterminedValue:
if line.il_instruction is not None:
instr = line.il_instruction
else:
instr_start = func.low_level_il.get_instruction_start(line.address, arch)
instr = func.low_level_il[instr_start]
if instr is not None:
after_possible = instr.get_possible_reg_values_after(self.followed_reg)
line.tokens.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(after_possible), token_value)
)
else:
line.tokens.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(after), token_value)
)
else:
line.tokens.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(after), token_value)
)
block_start = block.start
block_end = block.end
if block.is_low_level_il:
block_start = func.llil[block.start].address
block_end = func.llil[block.end - 1].address
# ===============================================================================
# Render the before-block "<reg> at block entry" comment
start_before = func.get_reg_value_at(block_start, self.followed_reg, arch)
line = [
InstructionTextToken(InstructionTextTokenType.CommentToken, '// ', token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, ' at block entry: ', token_value),
]
# Sometimes the first line is blank and we want to insert after it
first_line = 0
if len(lines[0].tokens) == 0:
first_line = 1
if start_before.type == RegisterValueType.UndeterminedValue:
if lines[first_line].il_instruction is not None:
instr = lines[first_line].il_instruction
else:
instr_start = func.low_level_il.get_instruction_start(block_start, arch)
instr = func.low_level_il[instr_start]
if instr is not None:
start_before_possible = instr.get_possible_reg_values(self.followed_reg)
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before_possible), token_value)
)
else:
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before), token_value)
)
else:
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(start_before), token_value)
)
lines.insert(first_line, DisassemblyTextLine(line, block_start))
# ===============================================================================
# Render the after-block "<reg> at block exit" comment
end_after = func.get_reg_value_after(block_end, self.followed_reg, arch)
line = [
InstructionTextToken(InstructionTextTokenType.CommentToken, '// ', token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, self.followed_reg, token_value),
InstructionTextToken(InstructionTextTokenType.CommentToken, ' at block exit: ', token_value),
]
# Sometimes the last line is blank and we want to insert before it
if len(lines) > 1 and len(lines[-1].tokens) == 0:
last_line = len(lines) - 1
else:
last_line = len(lines)
if end_after.type == RegisterValueType.UndeterminedValue:
if lines[last_line - 1].il_instruction is not None:
instr = lines[last_line - 1].il_instruction
else:
instr_start = func.low_level_il.get_instruction_start(block_end, arch)
instr = func.low_level_il[instr_start]
if instr is not None:
end_after_possible = instr.get_possible_reg_values_after(self.followed_reg)
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after_possible), token_value)
)
else:
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after), token_value)
)
else:
line.append(
InstructionTextToken(InstructionTextTokenType.CommentToken, str(end_after), token_value)
)
lines.insert(last_line, DisassemblyTextLine(line, block_end))
return lines
def apply_to_disassembly_block(
self,
block: BasicBlock,
lines: List['DisassemblyTextLine']
):
# Break this out into a helper so we don't have to write it twice
return self.apply_to_lines(block, lines)
def apply_to_low_level_il_block(
self,
block: BasicBlock,
lines: List['DisassemblyTextLine']
):
# Break this out into a helper so we don't have to write it twice
return self.apply_to_lines(block, lines)
def set_follow_reg(ctx: UIActionContext):
bv = ctx.binaryView
if bv is not None:
if bv.platform is not None:
regs = list(bv.platform.arch.regs.keys())
idx = interaction.get_large_choice_input("Choose", "Choose Followed Register", regs)
# Save choice both in QSettings and on the RenderLayer object instance
layer = RenderLayer[FollowRegRenderLayer.name]
if idx is None:
layer.followed_reg = None
QSettings().remove("layer/followed_reg")
else:
layer.followed_reg = regs[idx]
QSettings().setValue("layer/followed_reg", regs[idx])
# Trigger view refresh
ctx.context.refreshCurrentViewContents()
def can_set_follow_reg(ctx: UIActionContext) -> bool:
return ctx.binaryView is not None
FollowRegRenderLayer.register()
# Using a command for this is kinda janky but currently the only option
UIAction.registerAction("Set Followed Register")
UIActionHandler.globalActions().bindAction("Set Followed Register", UIAction(set_follow_reg, can_set_follow_reg))
|
