Skip to content

Commit 9099142

Browse files
CouleeAppsCouleeApps
committed
Dedupe code in example
1 parent 2507284 commit 9099142

File tree

1 file changed

+60
-103
lines changed

1 file changed

+60
-103
lines changed

‎python/examples/args_layer.py‎

Lines changed: 60 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
- Thunks are unhandled
3030
"""
3131

32+
3233
@functools.lru_cache(maxsize=64)
3334
def get_param_sites(mlil: MediumLevelILFunction) -> Mapping[LowLevelILInstruction, List[Tuple[MediumLevelILInstruction, int]]]:
3435
"""
@@ -123,6 +124,63 @@ def get_llil_arg(llil: LowLevelILInstruction) -> Iterator[Tuple[str, MediumLevel
123124
return
124125

125126

127+
def apply_to_lines(lines, get_instr, renderer):
128+
# So we don't process lines twice since we're iterating over a list as we modify it
129+
skip_lines = []
130+
131+
# Tailcalls that don't return incorrectly mark the { Does not return } line as a call
132+
ignore_calls = set()
133+
134+
for i, line in enumerate(lines):
135+
if len(line.tokens) == 0:
136+
continue
137+
if i in skip_lines:
138+
continue
139+
140+
llil_instr = get_instr(line)
141+
if llil_instr is not None:
142+
new_lines = []
143+
for (arg, call) in get_llil_arg(llil_instr):
144+
if call.operation == MediumLevelILOperation.MLIL_TAILCALL_SSA:
145+
comment = f"Argument '{arg}' for tailcall at {call.address:#x}"
146+
else:
147+
comment = f"Argument '{arg}' for call at {call.address:#x}"
148+
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
149+
for j, token in enumerate(line.tokens):
150+
if token.type == InstructionTextTokenType.AddressSeparatorToken:
151+
line.tokens = line.tokens[:j]
152+
break
153+
154+
# Annotate calls too so we can see them easily next to their args
155+
if llil_instr.address == line.address and llil_instr.address not in ignore_calls:
156+
if llil_instr.operation in [
157+
LowLevelILOperation.LLIL_CALL,
158+
LowLevelILOperation.LLIL_CALL_SSA,
159+
LowLevelILOperation.LLIL_TAILCALL,
160+
LowLevelILOperation.LLIL_TAILCALL_SSA
161+
]:
162+
ignore_calls.add(llil_instr.address)
163+
if llil_instr.operation in [
164+
LowLevelILOperation.LLIL_TAILCALL,
165+
LowLevelILOperation.LLIL_TAILCALL_SSA
166+
]:
167+
comment = f"Tailcall at {llil_instr.address:#x}"
168+
else:
169+
comment = f"Call at {llil_instr.address:#x}"
170+
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
171+
for j, token in enumerate(line.tokens):
172+
if token.type == InstructionTextTokenType.AddressSeparatorToken:
173+
line.tokens = line.tokens[:j]
174+
break
175+
176+
# If any of our lines changed, swap out the existing lines with the new ones
177+
if len(new_lines) > 0:
178+
lines.pop(i)
179+
for j, new_line in enumerate(new_lines):
180+
lines.insert(i + j, new_line)
181+
skip_lines.append(i + j)
182+
183+
126184
class ArgumentsRenderLayer(RenderLayer):
127185
name = "Annotate Call Parameters"
128186

@@ -132,116 +190,15 @@ def apply_to_disassembly_block(
132190
lines: List['DisassemblyTextLine']
133191
):
134192
renderer = DisassemblyTextRenderer(block.function)
135-
136-
# So we don't process lines twice since we're iterating over a list as we modify it
137-
skip_lines = []
138-
139-
# Tailcalls that don't return incorrectly mark the { Does not return } line as a call
140-
ignore_calls = set()
141-
142-
for i, line in enumerate(lines):
143-
if len(line.tokens) == 0:
144-
continue
145-
if i in skip_lines:
146-
continue
147-
llil_instr = block.function.get_llil_at(line.address)
148-
if llil_instr is not None:
149-
new_lines = []
150-
for (arg, call) in get_llil_arg(llil_instr):
151-
if call.operation == MediumLevelILOperation.MLIL_TAILCALL_SSA:
152-
comment = f"Argument '{arg}' for tailcall at {call.address:#x}"
153-
else:
154-
comment = f"Argument '{arg}' for call at {call.address:#x}"
155-
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
156-
for j, token in enumerate(line.tokens):
157-
if token.type == InstructionTextTokenType.AddressSeparatorToken:
158-
line.tokens = line.tokens[:j]
159-
break
160-
161-
if llil_instr.address == line.address and llil_instr.address not in ignore_calls:
162-
if llil_instr.operation in [
163-
LowLevelILOperation.LLIL_CALL,
164-
LowLevelILOperation.LLIL_CALL_SSA,
165-
LowLevelILOperation.LLIL_TAILCALL,
166-
LowLevelILOperation.LLIL_TAILCALL_SSA
167-
]:
168-
ignore_calls.add(llil_instr.address)
169-
if llil_instr.operation in [
170-
LowLevelILOperation.LLIL_TAILCALL,
171-
LowLevelILOperation.LLIL_TAILCALL_SSA
172-
]:
173-
comment = f"Tailcall at {llil_instr.address:#x}"
174-
else:
175-
comment = f"Call at {llil_instr.address:#x}"
176-
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
177-
for j, token in enumerate(line.tokens):
178-
if token.type == InstructionTextTokenType.AddressSeparatorToken:
179-
line.tokens = line.tokens[:j]
180-
break
181-
182-
if len(new_lines) > 0:
183-
lines.pop(i)
184-
for j, new_line in enumerate(new_lines):
185-
lines.insert(i + j, new_line)
186-
skip_lines.append(i + j)
193+
apply_to_lines(lines, lambda line: block.function.get_llil_at(line.address), renderer)
187194

188195
def apply_to_low_level_il_block(
189196
self,
190197
block: BasicBlock,
191198
lines: List['DisassemblyTextLine']
192199
):
193200
renderer = DisassemblyTextRenderer(block.function)
194-
195-
# So we don't process lines twice since we're iterating over a list as we modify it
196-
skip_lines = []
197-
198-
# Tailcalls that don't return incorrectly mark the { Does not return } line as a call
199-
ignore_calls = set()
200-
201-
for i, line in enumerate(lines):
202-
if len(line.tokens) == 0:
203-
continue
204-
if i in skip_lines:
205-
continue
206-
if line.il_instruction is not None:
207-
new_lines = []
208-
for (arg, call) in get_llil_arg(line.il_instruction):
209-
if call.operation == MediumLevelILOperation.MLIL_TAILCALL_SSA:
210-
comment = f"Argument '{arg}' for tailcall at {call.address:#x}"
211-
else:
212-
comment = f"Argument '{arg}' for call at {call.address:#x}"
213-
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
214-
for j, token in enumerate(line.tokens):
215-
if token.type == InstructionTextTokenType.AddressSeparatorToken:
216-
line.tokens = line.tokens[:j]
217-
break
218-
219-
if line.il_instruction.address == line.address and line.il_instruction.address not in ignore_calls:
220-
if line.il_instruction.operation in [
221-
LowLevelILOperation.LLIL_CALL,
222-
LowLevelILOperation.LLIL_CALL_SSA,
223-
LowLevelILOperation.LLIL_TAILCALL,
224-
LowLevelILOperation.LLIL_TAILCALL_SSA
225-
]:
226-
ignore_calls.add(line.il_instruction.address)
227-
if line.il_instruction.operation in [
228-
LowLevelILOperation.LLIL_TAILCALL,
229-
LowLevelILOperation.LLIL_TAILCALL_SSA
230-
]:
231-
comment = f"Tailcall at {line.il_instruction.address:#x}"
232-
else:
233-
comment = f"Call at {line.il_instruction.address:#x}"
234-
renderer.wrap_comment(new_lines, line, comment, False, " ", "")
235-
for j, token in enumerate(line.tokens):
236-
if token.type == InstructionTextTokenType.AddressSeparatorToken:
237-
line.tokens = line.tokens[:j]
238-
break
239-
240-
if len(new_lines) > 0:
241-
lines.pop(i)
242-
for j, new_line in enumerate(new_lines):
243-
lines.insert(i + j, new_line)
244-
skip_lines.append(i + j)
201+
apply_to_lines(lines, lambda line: line.il_instruction, renderer)
245202

246203

247204
ArgumentsRenderLayer.register()

0 commit comments

Comments
 (0)