emusca.py

#!/usr/bin/env python2

from __future__ import print_function

import binascii as ba
import os
import struct
import sys
from multiprocessing import Pool

from unicorn import *
from unicorn.arm_const import *

from capstone import *
from capstone.arm import *
from xprint import to_hex, to_x_32

import lief

import intervaltree as itree

import numpy as np

import gmpy2

from deco import concurrent, synchronized

cs_arm = Cs(CS_ARCH_ARM, CS_MODE_ARM)
cs_arm.detail = True
cs_thumb = Cs(CS_ARCH_ARM, CS_MODE_THUMB)
cs_thumb.detail = True

IRAM_START = 0x40000000
IRAM_SIZE = 256 * 1024
RCM_PAYLOAD_ADDR    = 0x4000A000

USB_SEND_ADDR = 0xfff05092
NvOsWaitUS_ADDR = 0xFFF00B9A

PMC_BASE = 0x7000E400
PMC_SCRATCH0 = 0x50
PMC_SCRATCH1 = 0x54

FOURK_ALIGN_MASK = 0xFFFFF000
EVEN_MASK = 0xFFFFFFFE

def roundup(x, m):
	return x if x % m == 0 else x + m - x % m

def rounddown(x, m):
	return x if x % m == 0 else x - x % m

reg_ids = (UC_ARM_REG_APSR, UC_ARM_REG_APSR_NZCV, UC_ARM_REG_CPSR, UC_ARM_REG_LR, UC_ARM_REG_PC, UC_ARM_REG_SP, UC_ARM_REG_SPSR, UC_ARM_REG_R0, UC_ARM_REG_R1, UC_ARM_REG_R2, UC_ARM_REG_R3, UC_ARM_REG_R4, UC_ARM_REG_R5, UC_ARM_REG_R6, UC_ARM_REG_R7, UC_ARM_REG_R8, UC_ARM_REG_R9, UC_ARM_REG_R10, UC_ARM_REG_R11, UC_ARM_REG_R12)

reg_names = {
	UC_ARM_REG_APSR: 'APSR',
	UC_ARM_REG_APSR_NZCV: 'APSR_NZCV',
	UC_ARM_REG_CPSR: 'CPSR',
	UC_ARM_REG_LR: 'LR',
	UC_ARM_REG_PC: 'PC',
	UC_ARM_REG_SP: 'SP',
	UC_ARM_REG_SPSR: 'SPSR',
	UC_ARM_REG_R0: 'R0',
	UC_ARM_REG_R1: 'R1',
	UC_ARM_REG_R2: 'R2',
	UC_ARM_REG_R3: 'R3',
	UC_ARM_REG_R4: 'R4',
	UC_ARM_REG_R5: 'R5',
	UC_ARM_REG_R6: 'R6',
	UC_ARM_REG_R7: 'R7',
	UC_ARM_REG_R8: 'R8',
	UC_ARM_REG_R9: 'R9',
	UC_ARM_REG_R10: 'R10',
	UC_ARM_REG_R11: 'R11',
	UC_ARM_REG_R12: 'R12',
}

def print_keystuff(uc):
	key = uc.mem_read(uc.key_sym.value, uc.key_sym.size)
	iv = uc.mem_read(uc.iv_sym.value, uc.iv_sym.size)
	ct = uc.mem_read(uc.ct_sym.value, uc.ct_sym.size)
	pt = uc.mem_read(uc.pt_sym.value, uc.pt_sym.size)
	# init_key = uc.mem_read(uc.init_key_sym.value, uc.init_key_sym.size)
	# init_iv = uc.mem_read(uc.init_iv_sym.value, uc.init_iv_sym.size)
	# init_ct = uc.mem_read(uc.init_ct_sym.value, uc.init_ct_sym.size)
	# init_pt = uc.mem_read(uc.init_pt_sym.value, uc.init_pt_sym.size)
	# aes_init_done = uc.mem_read(uc.aes_init_done_sym.value, uc.aes_init_done_sym.size)
	print("\tkey: {}".format(ba.hexlify(key)))
	print("\tiv: {}".format(ba.hexlify(iv)))
	print("\tct: {}".format(ba.hexlify(ct)))
	print("\tpt: {}".format(ba.hexlify(pt)))
	# print("\tinit_key: {}".format(ba.hexlify(init_key)))
	# print("\tinit_iv: {}".format(ba.hexlify(init_iv)))
	# print("\tinit_ct: {}".format(ba.hexlify(init_ct)))
	# print("\tinit_pt: {}".format(ba.hexlify(init_pt)))
	# print("\taes_init_done: {}".format(ba.hexlify(aes_init_done)))

def dump_regs(regs):
	for k in regs.keys():
		print("\t%s:\t0x%08x" % (reg_names[k], regs[k]))

def dump_regs_changed(regs):
	regs.pop(UC_ARM_REG_PC, None)
	for k in regs.keys():
		print("\t%s:\t0x%08x -> 0x%08x" % (reg_names[k], regs[k]['old'], regs[k]['new']))

def all_regs(uc):
	reg_vals = uc.reg_read_batch(reg_ids)
	regs = {}
	for i, reg_val in enumerate(reg_vals):
		regs[reg_ids[i]] = reg_val
	return regs

def changed_regs(old_regs, new_regs):
	regs = {}
	for k in old_regs.keys():
		if old_regs[k] != new_regs[k]:
			regs[k] = {'old': old_regs[k], 'new': new_regs[k]}
			# regs[k] = new_regs[k]
	return regs

def hamming_distance_pair(old_regs, new_regs):
	pass

def hamming_distance_changed(changed_regs):
	keys = changed_regs.keys()
	old = map(lambda r: changed_regs[r]['old'], keys)
	new = map(lambda r: changed_regs[r]['new'], keys)
	# print("old: {} new: {}".format(old, new))
	# xor = np.bitwise_xor(old, new)
	# print("xor: {}".format(xor))
	# popcnt = np.bincount(xor.transpose())
	# print("popcnt: {}".format(popcnt))
	hd = sum(map(lambda p: gmpy2.hamdist(p[0], p[1]), zip(old, new)))
	# print("hd: {}".format(hd))
	return hd


def print_insn_detail(insn, insn_bytes):
    # print address, insn bytes, mnemonic and operands
    insn_bytes_str = ba.hexlify(insn_bytes)
    insn_bytes_str = ' '.join(insn_bytes_str[i:i+2] for i in range(0, len(insn_bytes_str), 2))
    print("0x%x: %s\t%s\t%s" % (insn.address, insn_bytes_str, insn.mnemonic, insn.op_str))

    return

    # "data" instruction generated by SKIPDATA option has no detail
    if insn.id == 0:
        return

    if len(insn.operands) > 0:
        print("\top_count: %u" % len(insn.operands))
        c = 0
        for i in insn.operands:
            if i.type == ARM_OP_REG:
                print("\t\toperands[%u].type: REG = %s" % (c, insn.reg_name(i.reg)))
            if i.type == ARM_OP_IMM:
                print("\t\toperands[%u].type: IMM = 0x%s" % (c, to_x_32(i.imm)))
            if i.type == ARM_OP_PIMM:
                print("\t\toperands[%u].type: P-IMM = %u" % (c, i.imm))
            if i.type == ARM_OP_CIMM:
                print("\t\toperands[%u].type: C-IMM = %u" % (c, i.imm))
            if i.type == ARM_OP_FP:
                print("\t\toperands[%u].type: FP = %f" % (c, i.fp))
            if i.type == ARM_OP_SYSREG:
                print("\t\toperands[%u].type: SYSREG = %u" % (c, i.reg))
            if i.type == ARM_OP_SETEND:
                if i.setend == ARM_SETEND_BE:
                    print("\t\toperands[%u].type: SETEND = be" % c)
                else:
                    print("\t\toperands[%u].type: SETEND = le" % c)
            if i.type == ARM_OP_MEM:
                print("\t\toperands[%u].type: MEM" % c)
                if i.mem.base != 0:
                    print("\t\t\toperands[%u].mem.base: REG = %s" \
                        % (c, insn.reg_name(i.mem.base)))
                if i.mem.index != 0:
                    print("\t\t\toperands[%u].mem.index: REG = %s" \
                        % (c, insn.reg_name(i.mem.index)))
                if i.mem.scale != 1:
                    print("\t\t\toperands[%u].mem.scale: %u" \
                        % (c, i.mem.scale))
                if i.mem.disp != 0:
                    print("\t\t\toperands[%u].mem.disp: 0x%s" \
                        % (c, to_x_32(i.mem.disp)))
                if i.mem.lshift != 0:
                    print("\t\t\toperands[%u].mem.lshift: 0x%s" \
                        % (c, to_x_32(i.mem.lshift)))

            if i.neon_lane != -1:
                print("\t\toperands[%u].neon_lane = %u" % (c, i.neon_lane))

            if i.access == CS_AC_READ:
                print("\t\toperands[%u].access: READ\n" % (c))
            elif i.access == CS_AC_WRITE:
                print("\t\toperands[%u].access: WRITE\n" % (c))
            elif i.access == CS_AC_READ | CS_AC_WRITE:
                print("\t\toperands[%u].access: READ | WRITE\n" % (c))

            if i.shift.type != ARM_SFT_INVALID and i.shift.value:
                print("\t\t\tShift: %u = %u" \
                    % (i.shift.type, i.shift.value))
            if i.vector_index != -1:
                print("\t\t\toperands[%u].vector_index = %u" %(c, i.vector_index))
            if i.subtracted:
                print("\t\t\toperands[%u].subtracted = True" %c)

            c += 1

    if insn.update_flags:
        print("\tUpdate-flags: True")
    if insn.writeback:
        print("\tWrite-back: True")
    if not insn.cc in [ARM_CC_AL, ARM_CC_INVALID]:
        print("\tCode condition: %u" % insn.cc)
    if insn.cps_mode:
        print("\tCPSI-mode: %u" %(insn.cps_mode))
    if insn.cps_flag:
        print("\tCPSI-flag: %u" %(insn.cps_flag))
    if insn.vector_data:
        print("\tVector-data: %u" %(insn.vector_data))
    if insn.vector_size:
        print("\tVector-size: %u" %(insn.vector_size))
    if insn.usermode:
        print("\tUser-mode: True")
    if insn.mem_barrier:
        print("\tMemory-barrier: %u" %(insn.mem_barrier))

    (regs_read, regs_write) = insn.regs_access()

    if len(regs_read) > 0:
        print("\tRegisters read:", end="")
        for r in regs_read:
            print(" %s" %(insn.reg_name(r)), end="")
        print("")

    if len(regs_write) > 0:
        print("\tRegisters modified:", end="")
        for r in regs_write:
            print(" %s" %(insn.reg_name(r)), end="")
        print("")

def findsym(elf, name):
	return next(s for s in elf.symbols if s.name == name)

# callback for tracing basic blocks
def hook_block(uc, address, size, user_data):
	print(">>> Tracing basic block at 0x%x, block size = 0x%x" %(address, size))

# callback for tracing instructions
def hook_code(uc, address, size, user_data):
	new_regs = all_regs(uc)
	ch_regs = changed_regs(uc.saved_regs, new_regs)
	# ch_regs.pop(UC_ARM_REG_PC, None)
	hamming_distance_changed(ch_regs)
	dump_regs_changed(ch_regs)
	is_thumb = new_regs[UC_ARM_REG_CPSR] & (1 << 5) != 0
	mode_str = None
	if is_thumb:
		mode_str = "Thumb"
	else:
		mode_str = "ARM"
	print(">>> Tracing instruction at 0x%x, instruction size = 0x%x, mode: %s" % (address, size, mode_str))
	pc = new_regs[UC_ARM_REG_PC]
	insn_bytes = uc.mem_read(pc, size)
	insn = None
	if is_thumb:
		insn = list(cs_thumb.disasm(insn_bytes, pc))[0]
	else:
		insn = list(cs_arm.disasm(insn_bytes, pc))[0]
	print_insn_detail(insn, insn_bytes)
	uc.saved_regs = new_regs

def hook_code_hamming_distance(uc, address, size, user_data):
	new_regs = all_regs(uc)
	ch_regs = changed_regs(uc.saved_regs, new_regs)
	# ch_regs.pop(UC_ARM_REG_PC, None)
	hd = hamming_distance_changed(ch_regs)
	if address == uc.aes_round_marker_addr:
		uc.aes_round_marks.append(uc.insn_count)
	uc.hd_list.append(hd)
	uc.saved_regs = new_regs
	uc.insn_count += 1

def hook_exit(uc, address, size, user_data):
	# print(">>> Hooking _exit")
	(retval, lr, pc) = uc.reg_read_batch((UC_ARM_REG_R0, UC_ARM_REG_LR, UC_ARM_REG_PC))
	# print(">>> \tretval: %d lr: 0x%08x address: 0x%08X pc: 0x%08x size: %d" % (retval, lr, address, pc, size))
	# print_keystuff(uc)
	uc.emu_stop()

# callback for tracing invalid memory access (READ or WRITE)
def hook_mem_unmapped(uc, access, address, size, value, user_data):
	if access == UC_MEM_WRITE_UNMAPPED:
		print(">>> Missing memory is being WRITE at 0x%08x, data size = %u, data value = 0x%08x" \
				%(address, size, value))
	else:
		print(">>> Missing memory is being READ at 0x%08x, data size = %u, data value = 0x%08x" \
				%(address, size, value))
	# return False to indicate we want to stop emulation
	return False

# callback for tracing memory access (READ or WRITE)
def hook_mem_access(uc, access, address, size, value, user_data):
	old_val_bytes = uc.mem_read(address, size)
	old_val_bytes = uc.mem_read(address, size)
	struct_fmt = {1: '<B', 2: '<H', 4: '<I'}[size]
	old_val = struct.unpack(struct_fmt, old_val_bytes)[0]
	if access == UC_MEM_WRITE:
		print(">>> Memory is being WRITE at 0x%08x, data size = %u, data value = 0x%08x, old value = 0x%08x" \
				%(address, size, value, old_val))
		if address == PMC_BASE and value & (1 << 4) != 0:
			print(">>> PMC reset issued, stopping emu")
			uc.emu_stop()
		if address == PMC_BASE + PMC_SCRATCH0:
			print(">>> PMC_SCRATCH0 written")
		if address == PMC_BASE + PMC_SCRATCH1:
			print(">>> PMC_SCRATCH1 written")
	else:   # READ
		print(">>> Memory is being READ at 0x%08x, data size = %u, data value = 0x%08x" \
				%(address, size, old_val))


uc_inited_ctx = None
uc_elf = None
uc_stack_top = None
uc_maxaddr = None
def get_inited_uc(elf_path, key, iv, ct):
	global uc_inited_ctx, uc_elf, uc_stack_top, uc_maxaddr
	if uc_inited_ctx is None:
		elf = lief.parse(elf_path)
		uc_elf = elf
		tree = itree.IntervalTree()
		for seg in elf.segments:
			rstart = rounddown(seg.virtual_address, 4096)
			rend = roundup(seg.virtual_address + seg.virtual_size, 4096)
			rsize_content = roundup(seg.virtual_size, 4096)
			rsize_range = rend - rstart
			zpad_front = b'\x00' * (seg.virtual_address - rstart)
			zpad_back = b'\x00' * (rend - (seg.virtual_address + seg.virtual_size))
			# print("rstart: 0x{:08x} rend: 0x{:08x} rsize_content: 0x{:08x} rsize_content: 0x{:08x}".format(rstart, rend, rsize_content, rsize_range))
			assert(not tree.overlaps(rstart, rend))
			tree[rstart:rend] = zpad_front + bytes(bytearray(seg.content)) + zpad_back
		minaddr = tree.begin()
		maxaddr = tree.end()
		uc_maxaddr = maxaddr
		# print("minaddr: 0x{:08x} maxaddr: 0x{:08x}".format(minaddr, maxaddr))
		stack_gap = 1024*1024
		stack_size = 1024*1024
		stack_begin = tree.end() + stack_gap
		stack_end = stack_begin + stack_size
		tree[stack_begin:stack_end] = b'\x00' * stack_size
		stack_top = stack_end
		uc_stack_top = stack_top
		# print("tree: {}".format(tree))
		minaddr = tree.begin()
		maxaddr = tree.end()
		# print("minaddr: 0x{:08x} maxaddr: 0x{:08x}".format(minaddr, maxaddr))

		# print("Emulate thumb code")
		try:
			# Initialize emulator in thumb mode
			uc = Uc(UC_ARCH_ARM, UC_MODE_THUMB)
			uc.elf = elf
			uc.maxaddr = maxaddr

			for seg in tree:
				uc.mem_map(seg.begin, seg.end - seg.begin)
				uc.mem_write(seg.begin, seg.data)

			uc.reg_write(UC_ARM_REG_SP, stack_top)
			uc.reg_write(UC_ARM_REG_APSR, 0xFFFFFFFF) #All application flags turned on

			def hook_aes_init_marker(uc, address, size, user_data):
				uc.emu_stop()

			uc.key_sym = findsym(uc.elf, 'key')
			uc.iv_sym = findsym(uc.elf, 'iv')
			uc.ct_sym = findsym(uc.elf, 'ct')
			uc.pt_sym = findsym(uc.elf, 'pt')
			# uc.init_key_sym = findsym(uc.elf, 'init_key')
			# uc.init_iv_sym = findsym(uc.elf, 'init_iv')
			# uc.init_ct_sym = findsym(uc.elf, 'init_ct')
			# uc.init_pt_sym = findsym(uc.elf, 'init_pt')
			# uc.aes_init_done_sym = findsym(uc.elf, 'aes_init_done')
			uc.exit_sym = findsym(uc.elf, '_exit')

			aes_init_marker_sym = findsym(uc.elf, 'aes_init_marker')
			aes_init_marker_addr = aes_init_marker_sym.value & EVEN_MASK

			uc.hook_add(UC_HOOK_CODE, hook_aes_init_marker, begin=aes_init_marker_addr & EVEN_MASK, end=aes_init_marker_addr + 2)

			# print_keystuff(uc)
			uc.emu_start(uc.elf.entrypoint, uc.maxaddr)
			# print_keystuff(uc)
			mem_ctx = []
			for region in uc.mem_regions():
				mem_ctx.append({'addr': region[0], 'sz': region[1] - region[0] + 1, 'data': bytes(uc.mem_read(region[0], region[1] - region[0] + 1))})
			uc_inited_ctx = mem_ctx
		except UcError as e:
			print("ERROR: %s" % e)

	elf = uc_elf
	uc = Uc(UC_ARCH_ARM, UC_MODE_THUMB)
	uc.elf = elf
	uc.insn_count = 0
	uc.maxaddr = uc_maxaddr

	for seg in uc_inited_ctx:
		uc.mem_map(seg['addr'], seg['sz'])
		uc.mem_write(seg['addr'], seg['data'])

	uc.reg_write(UC_ARM_REG_SP, uc_stack_top)
	uc.reg_write(UC_ARM_REG_APSR, 0xFFFFFFFF) #All application flags turned on

	uc.saved_regs = all_regs(uc)

	uc.hook_add(UC_HOOK_CODE, hook_code_hamming_distance)

	uc.key_sym = findsym(uc.elf, 'key')
	uc.iv_sym = findsym(uc.elf, 'iv')
	uc.ct_sym = findsym(uc.elf, 'ct')
	uc.pt_sym = findsym(uc.elf, 'pt')
	# uc.init_key_sym = findsym(uc.elf, 'init_key')
	# uc.init_iv_sym = findsym(uc.elf, 'init_iv')
	# uc.init_ct_sym = findsym(uc.elf, 'init_ct')
	# uc.init_pt_sym = findsym(uc.elf, 'init_pt')
	# uc.aes_init_done_sym = findsym(uc.elf, 'aes_init_done')
	uc.exit_sym = findsym(uc.elf, '_exit')
	uc.aes_round_marker_sym = findsym(uc.elf, 'aes_round_marker')
	uc.aes_round_marker_addr = uc.aes_round_marker_sym.value & EVEN_MASK
	uc.aes_round_marks = []
	uc.hook_add(UC_HOOK_CODE, hook_exit, begin=uc.exit_sym.value & EVEN_MASK, end=(uc.exit_sym.value & EVEN_MASK) + 2)

	assert(len(key) == 16)
	assert(len(iv) == 16)
	assert(len(ct) == 16)
	# print_keystuff(uc)
	uc.mem_write(uc.key_sym.value, key)
	uc.mem_write(uc.iv_sym.value, iv)
	uc.mem_write(uc.ct_sym.value, ct)
	# print_keystuff(uc)

	uc.hd_list = []
	return uc

@concurrent
def gen_traces(elf_path, key, iv, ct):
	try:
		uc = get_inited_uc(elf_path, key, iv, ct)
		uc.emu_start(uc.elf.entrypoint, uc.maxaddr)
		pt = uc.mem_read(uc.pt_sym.value, uc.pt_sym.size)
		uc.hd_list = np.asarray(uc.hd_list, dtype=np.dtype(np.uint16))
		# print_keystuff(uc)
		# print("pt: {}".format(ba.hexlify(pt)))
		# print("len(hd_list): {}".format(len(uc.hd_list)))
		return (pt, uc.aes_round_marks, uc.hd_list)

	except UcError as e:
		print("ERROR: %s" % e)

def write_daredevil(key, cts, pts, hd_lists):
	prefix = 'foo'
	with open(prefix + '_in_daredevil.bin', 'wb') as in_f, open(prefix + '_out_daredevil.bin', 'wb') as out_f, open(prefix + '_trace_daredevil.bin', 'wb') as trace_f:
		for i in range(len(cts)):
			indata = cts[i]
			outdata = pts[i]
			tracedata = struct.pack("={}f".format(len(hd_lists[i])), *hd_lists[i])
			in_f.write(indata)
			out_f.write(outdata)
			trace_f.write(tracedata)
	with open(prefix + '_daredevil.cfg', 'w') as cfgf:
		cfgf.write(
	"""
[Traces]
files=1
trace_type={format}
transpose=true
index=0
nsamples={nsamples}
trace={traces_filename} {ntraces} {nsamples}

[Guesses]
files=1
guess_type=u
transpose=true
guess={input_filename} {ntraces} 16

[General]
threads=8
order=1
return_type=double
algorithm=AES
position=LUT/AES_AFTER_SBOXINV
round=0
bitnum=none
bytenum=all
correct_key=0x{key}
memory=4G
top=20
	""".format(format='f', ntraces=len(hd_lists), nsamples=len(hd_lists[0]), traces_filename=os.getcwd() + '/' + prefix + '_trace_daredevil.bin', input_filename=os.getcwd() + '/' + prefix + '_in_daredevil.bin', key=ba.hexlify(key)))

def main(argv):
	get_inited_uc(argv[1], b'\x00' * 16, b'\x00' * 16, b'\x00' * 16) # cache init
	key = ba.unhexlify('00112233445566778899aabbccddeeff')
	# key = b'\x00' * 16
	iv = b'\x00' * 16
	res = []
	for i in range(4096):
		ct = os.urandom(16)
		# ct = b'\x00' * 16
		# print("main ct: {}".format(ba.hexlify(ct)))
		res.append({'ct': ct, 'res': gen_traces(argv[1], key, iv, ct)})
	gen_traces.wait()
	cts = []
	pts = []
	hd_lists = []
	round_marks = []
	for r in res:
		ct = r['ct']
		(pt, aes_round_marks, hd_list) = r['res'].result()
		cts.append(ct)
		pts.append(pt)
		round_marks.append(aes_round_marks)
		hd_lists.append(hd_list)
		# np.save('foo.npy', hd_list)
	write_daredevil(key, cts, pts, hd_lists)
	print("round_marks: {}".format(round_marks[0]))
	last = 0
	diffs = []
	for n in round_marks[0]:
		diffs.append(n - last)
		last = n
	print("diffs: {}".format(diffs))

if __name__ == "__main__":
	sys.exit(main(sys.argv))