#!/usr/bin/env python
# fsck secneo

from __future__ import print_function
from unicorn import *
from unicorn.arm_const import *
from capstone import *
import binascii

DEBUG = False

#
# Utility functions
#
def info(formatted_string):
    print(formatted_string)

def error(formatted_string):
    print('ERROR - %s' % formatted_string)

def debug(formatted_string):
    if DEBUG:
        print('DEBUG - %s' % formatted_string)

# memory address where emulation starts
ADDRESS    = 0x1000

# Place for writing and playing with memory
SCRATCH_ADDRESS    = 0x1100

def disassemble_arm(arm_code):
    md = Cs(CS_ARCH_ARM, CS_MODE_THUMB)
    for i in md.disasm(arm_code, 0x1000):
        info("0x%x:\t%s\t%s" %(i.address, i.mnemonic, i.op_str))

def emulate_arm(arm_code=None, ro_data=None, ro_data_start=None, ro_offset=None):
    debug(">>> Emulating ARM code")
    try:
        mu = Uc(UC_ARCH_ARM, UC_MODE_THUMB)

        # Memory for actual code
        mu.mem_map(ADDRESS, 2 * 1024 * 1024)
        mu.mem_write(ADDRESS, arm_code)

        # RO Memory
        mu.mem_write(ro_data_start, ro_data)

        # write dummy data to be emulated to memory
        mu.mem_write(SCRATCH_ADDRESS, "\x00"*64)

        # Initialize some registers, automated by script
        mu.reg_write(UC_ARM_REG_R1, ro_offset)
        mu.reg_write(UC_ARM_REG_R3, ro_offset)
        mu.reg_write(UC_ARM_REG_R4, SCRATCH_ADDRESS)
        mu.reg_write(UC_ARM_REG_LR, SCRATCH_ADDRESS)

        mu.reg_write(UC_ARM_REG_R0, 0x00)
        mu.reg_write(UC_ARM_REG_R12, 0x00)

        if DEBUG:
            debug(">>> Before emulation ")
            for i in range(UC_ARM_REG_R0, UC_ARM_REG_R12):
                val = mu.reg_read(i)
                debug("\t %s = 0x%x" % ("R" + str(i-UC_ARM_REG_R0),val))

        mu.emu_start(ADDRESS + 1, ADDRESS + len(arm_code))

        debug(">>> Emulation done.")

        if DEBUG:
            debug(">>> Emulation done. Below is the CPU context")

            sp = mu.reg_read(UC_ARM_REG_SP)
            debug(">>> SP = 0x%x" %sp)
            val = mu.reg_read(UC_ARM_REG_PC)
            debug(">>> PC = 0x%x" %val)
            for i in range(UC_ARM_REG_R0, UC_ARM_REG_R12):
                val = mu.reg_read(i)
                debug(">>> %s = 0x%x" % ("R" + str(i-UC_ARM_REG_R0),val))

            debug("Memory at addr 0x%X %s" % (SCRATCH_ADDRESS, binascii.hexlify(content)))

        content = mu.mem_read(SCRATCH_ADDRESS, 100)
        info(">>> Decrypted text : %s" % content)


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

def get_ro_data():
    rodata_seg = idaapi.get_segm_by_name('.rodata')
    if rodata_seg is not None:
        return rodata_seg.startEA, idc.GetManyBytes(rodata_seg.startEA, rodata_seg.endEA - rodata_seg.startEA, False)
    return None, None

def get_decrypt_functions():
    funcs = []
    text_seg = idaapi.get_segm_by_name('.text')
    if text_seg is None:
        return 0, None

    # This appears to be in all decryption functions
    # TODO : Assert this works across multiple compiles, etc
    # ADD.W R4, LR, R0
    hook_opcodes = '0E EB 00 04'

    start = text_seg.startEA
    next_address = None
    while True:
        next_address = idaapi.find_binary(start, text_seg.endEA, hook_opcodes, 0, SEARCH_DOWN)
        if next_address == idaapi.BADADDR:
            break

        func = idaapi.get_func(next_address)
        # TODO : There may be a case where it isn't a function /yet/, which we may want to cover
        if func is None:
            error("Hit an issue with 0x%x" % next_address)
            break

        start = func.endEA
        funcs.append(func)

        if func.endEA > text_seg.endEA:
            error('Odd case that should not be possible hit!')
            break

    return len(funcs), funcs

# TODO : Beef this up, right now this is simple enough and easy to do
# For tested binaries, this worked for 74/76 functions
# Will likely need to validate new registers aren't used
def check_func(func):
    if func.startEA + 0x38 > func.endEA:
        return False

    if GetMnem(func.startEA + 0x38) == "ADDS":
        return True

    return False

def seek_bound(func):
    if func is None or func.startEA is None or func.startEA + 0x3e > func.endEA:
        return None

    addr = func.startEA + 0x3e

    while True:
        addr = FindCode(addr, SEARCH_DOWN)

        if GetMnem(addr) == "MOVS":
            return addr

        # Unsure if the latter case could even be possible, but whatever - lets check?
        if addr > func.endEA or addr < func.startEA:
            return None

def get_offset_into_ro_data(func):
    # TODO : Better check
    if GetMnem(func.startEA + 0x30) != "LDR":
        return None

    try:
        opnd = GetOpnd(func.startEA + 0x30, 1)
        if opnd.index("byte_") > 0:
            start = opnd.index("byte_") + len("byte_")
            end = opnd.index(" ", start)
            if end > 0:
                return int(opnd[start:end], 16)
    except ValueError as e:
        return None
    return None

def get_arm_code(func):
    lower_bound = func.startEA + 0x3e
    upper_bound = seek_bound(func)

    if upper_bound is None:
        return None

    arm_code = idc.GetManyBytes(lower_bound, upper_bound - lower_bound, False)

    if arm_code is None or arm_code <= 0:
        return None

    return arm_code

def decrypt_strings(func, ro_data, rodata_start):
    arm_code = get_arm_code(func)

    offset = get_offset_into_ro_data(func)
    if offset is not None:
        debug("Using ro_data offset 0x%x for 0x%0x" % (offset, func.startEA))
        emulate_arm(arm_code, ro_data, rodata_start, offset)
        return True

    return False

if __name__ == '__main__':
    rodata_start, ro_data = get_ro_data()
    if ro_data is None or ro_data <= 0:
        error("Issue finding '.rodata', bailing...")
        exit()

    how_many, funcs = get_decrypt_functions()
    info("Found %d decryption functions..." % how_many)

    for func in funcs:
        if check_func(func):
            if DEBUG:
                info("=" * 26)
                disassemble_arm(get_arm_code(func))
                info("=" * 26)
            if not decrypt_strings(func, ro_data, rodata_start):
                error("Decrypt failed..")
        else:
            error("0x%x doesn't seem correct" % func.startEA)

    # TODO : relabel function names and ending pointer