os.close(tmpfd)
with zipfile.ZipFile(jar, 'r') as z_in, zipfile.ZipFile(tmpname,
'w') as z_out:
readFunc = functools.partial(readArchive, z_in)
targets_dsm = disassembleSub(readFunc,
out_dsm,
targets=targets,
roundtrip=True)
for i, cls in enumerate(patches):
with open('dsm_classes/' + str(cls['class']) + '.j',
'r+') as file:
patcher = Patcher(cls, file)
patch_res = patcher.patch()
patch_id = str(cls['id'])
if 'depends' in cls:
patch_id += '/d' + str(cls['depends'])
patches_finished.append([patch_id, patch_res])
file.close()
gauage += iter_cnt
d.gauge_update(percent=int(gauage))
for item in z_in.infolist():
buffer = z_in.read(item.filename)
def patch_code(self, enable_trampoline=False):
patcher = Patcher(self.filename)
assert len(self.paths.items()) > 1
self.logger.info("-" * 30 + 'begin patch' + "-" * 30)
# nop irrelevant nodes
for node in self.deflat_analyzer.irrelevant_nodes:
if self.arch_type == QL_ARCH.ARM64:
start_addr = node.addr # irrelevant nodes keeps its zero base
self.logger.debug(
f"nop at {hex(start_addr)}, size {node.size}")
assert node.size % 4 == 0
nop_count = int(node.size / 4)
instruction_value = arm64_util.assemble_nop_instruction(
) * nop_count
patcher.patch(start_addr, node.size, instruction_value)
else:
raise Exception("Unsupported Arch")
max_trampoline_pool_cnt = 2 * len(self.paths.keys())
trampoline_pool: List[int] = [] # only for arm64.
used_trampoline_pool: List[int] = [] # block addr, target_addr
if enable_trampoline:
if self.arch_type == QL_ARCH.ARM64:
for node in self.deflat_analyzer.irrelevant_nodes:
if node.size >= 4:
for cur_addr_offset in range(0, node.size, 4):
trampoline_pool.append(node.addr + cur_addr_offset)
if len(trampoline_pool) >= max_trampoline_pool_cnt:
break
if len(trampoline_pool) >= max_trampoline_pool_cnt:
break
trampoline_pool.append(
*self.deflat_analyzer.manual_trampoline_addr_list)
# handle control flow
for block_id, successors in self.paths.items():
block = self.block_container.get_block_from_id(block_id)
start_addr = block.start_addr - self.base_addr
self.logger.debug(
f"patch working on {hex(start_addr)}, {successors}")
instructions = [
ins for ins in self.md.disasm(block.code, start_addr)
]
if self.arch_type == QL_ARCH.ARM64:
# ARM64 patch
if len(successors) == 2:
# real branch
true_branch = self.block_container.get_block_from_id(
successors[0]).start_addr - self.base_addr
false_branch = self.block_container.get_block_from_id(
successors[1]).start_addr - self.base_addr
self.logger.debug(
f"true {hex(true_branch)}, false {hex(false_branch)}")
should_trampoline = False
current_trampoline_addr = -1
if true_branch < start_addr and enable_trampoline:
should_trampoline = True
for i in range(len(trampoline_pool)):
if trampoline_pool[i] > start_addr:
current_trampoline_addr = trampoline_pool[i]
trampoline_pool.pop(i)
break
if current_trampoline_addr == -1:
self.logger.error(
f"Fail to find the suitable trampoline at {hex(start_addr)} with branch {hex(true_branch)}"
)
should_trampoline = False
if instructions[
-2].id in arm64_util.get_branch_instruction_types(
):
self.logger.debug(f"at -2")
instruction_address = instructions[-2].address
if should_trampoline:
self.logger.debug(
f"trampoline to {hex(current_trampoline_addr)}"
)
assert current_trampoline_addr != -1
patched_code = arm64_util.assemble_branch_instruction(
instruction_address,
current_trampoline_addr - instruction_address,
false_branch,
instructions[-2].op_str.split(',')[-1].strip())
assert len(patched_code) == 8
# patch in trampoline
used_trampoline_pool.append(
current_trampoline_addr)
b_instruction = arm64_util.assemble_no_branch_instruction(
current_trampoline_addr, true_branch)
patcher.patch(current_trampoline_addr, 4,
b_instruction)
else:
# recalculate the offset due to the keystone bug, maybe.
patched_code = arm64_util.assemble_branch_instruction(
instruction_address,
true_branch - instruction_address,
false_branch,
instructions[-2].op_str.split(',')[-1].strip())
assert len(patched_code) == 8
patcher.patch(instruction_address, 8, patched_code)
elif instructions[
-3].id in arm64_util.get_branch_instruction_types(
):
self.logger.debug(f"at -3")
instruction_address = instructions[-3].address
branch_offset = instruction_address + 4
if should_trampoline:
self.logger.debug(
f"trampoline to {hex(current_trampoline_addr)}"
)
assert current_trampoline_addr != -1
patched_code = arm64_util.assemble_branch_instruction(
branch_offset,
current_trampoline_addr - branch_offset,
false_branch,
instructions[-3].op_str.split(',')[-1].strip())
assert len(patched_code) == 8
# patch in trampoline
used_trampoline_pool.append(
current_trampoline_addr)
b_instruction = arm64_util.assemble_no_branch_instruction(
current_trampoline_addr, true_branch)
patcher.patch(current_trampoline_addr, 4,
b_instruction)
else:
patched_code = arm64_util.assemble_branch_instruction(
branch_offset, true_branch - branch_offset,
false_branch,
instructions[-3].op_str.split(',')[-1].strip())
assert len(patched_code) == 8
patcher.copy_to_patch(instruction_address,
branch_offset, 4)
patcher.patch(branch_offset, 8, patched_code)
else:
assert len(instructions) > 4
self.logger.warning(
"may encounter special csel instruction with larger than 2 offset, this is an experimental patch."
)
target_branch_instruction = None
# instructions[-3].id in arm64_util.get_branch_instruction_types()
for cursor_offset in range(
len(instructions) - 1, -1, -1):
if instructions[
cursor_offset].id in arm64_util.get_branch_instruction_types(
):
target_branch_instruction = instructions[
cursor_offset]
break
if target_branch_instruction is None:
raise Exception("Unhandled Branch Block")
target_branch_instruction_address = target_branch_instruction.address
self.logger.debug(
f"target_branch_instruction at {hex(target_branch_instruction_address)}, {target_branch_instruction.op_str}"
)
# Disable trampoline here!
branch_offset = start_addr + block.size - 8 # instructions[-2].addresss # we assume the last two instructions to be branch
self.logger.debug(
f"last two instruction at {hex(branch_offset)}")
assert isinstance(branch_offset, int)
patched_code = arm64_util.assemble_branch_instruction(
branch_offset, true_branch - branch_offset,
false_branch,
target_branch_instruction.op_str.split(
',')[-1].strip())
assert len(patched_code) == 8
cnt = branch_offset - target_branch_instruction_address
patcher.copy_to_patch(
target_branch_instruction_address,
target_branch_instruction_address + 4, cnt)
patcher.patch(branch_offset, 8, patched_code)
elif len(successors) == 1:
# force jump
instruction_address = instructions[-1].address
next_block = self.block_container.get_block_from_id(
successors[0]).start_addr - self.base_addr
self.logger.debug(f"next_block {hex(next_block)}")
patched_code = arm64_util.assemble_no_branch_instruction(
instruction_address, next_block)
patcher.patch(instruction_address, 4, patched_code)
else:
assert len(successors) == 0
# return block
continue
else:
raise Exception("Unsupported Arch")
patcher.write_patch_to_file()
self.logger.info("Patch code finish.")