import sys
import struct

from unicorn import *
from unicorn.arm_const import *

# address of call to HAL_I2C_Slave_Receive
CALL_ADDR = 0x08003392

def hook_code(mu: Uc, address, size, user_data):
    pin_addr = mu.reg_read(UC_ARM_REG_R1)
    mu.mem_write(pin_addr, user_data)
    mu.reg_write(UC_ARM_REG_PC, 0x08003396 | 1)

with open('code.bin', 'rb') as f:
    code = f.read()

with open('ram.bin', 'rb') as f:
    ram = f.read()

def do_run(pin_test):
    mu = Uc(UC_ARCH_ARM, UC_MODE_THUMB)

    # map code memory
    mu.mem_map(0x08000000, 0x80000)

    # map RAM
    mu.mem_map(0x20000000, 0x24000)

    # Populate memory
    mu.mem_write(0x08000000, code)
    mu.mem_write(0x20000000, ram)

    mu.reg_write(UC_ARM_REG_R0, 0x14)
    mu.reg_write(UC_ARM_REG_R1, 0xF)
    mu.reg_write(UC_ARM_REG_R2, 0x4)
    mu.reg_write(UC_ARM_REG_R3, 0x8003385)
    mu.reg_write(UC_ARM_REG_R4, 0x20023fef)
    mu.reg_write(UC_ARM_REG_R5, 0x17)
    mu.reg_write(UC_ARM_REG_R6, 0xffffffff)
    mu.reg_write(UC_ARM_REG_R7, 0x20023fd0)
    mu.reg_write(UC_ARM_REG_R8, 0xffffffff)
    mu.reg_write(UC_ARM_REG_R9, 0xffffffff)
    mu.reg_write(UC_ARM_REG_R10, 0xffffffff)
    mu.reg_write(UC_ARM_REG_R11, 0xffffffff)
    mu.reg_write(UC_ARM_REG_R12, 0x20023f90)
    mu.reg_write(UC_ARM_REG_SP, 0x20023f90)
    mu.reg_write(UC_ARM_REG_LR, 0x80035f7)
    mu.reg_write(UC_ARM_REG_PC, 0x8003388)
    mu.reg_write(UC_ARM_REG_APSR, 0x1000000)
    mu.reg_write(UC_ARM_REG_MSP, 0x20023fd0)
    mu.reg_write(UC_ARM_REG_PSP, 0xfffffffc)
    mu.reg_write(UC_ARM_REG_CONTROL, 0x0)

    mu.hook_add(UC_HOOK_CODE, hook_code, user_data=pin_test, begin=CALL_ADDR, end=CALL_ADDR)

    mu.emu_start(0x8003388 | 1, 0x80033A8)
    decrypted_addr = mu.reg_read(UC_ARM_REG_R0)
    data = mu.mem_read(decrypted_addr, 1024)
    return data.split(b'\0')[0].decode('ascii', errors='ignore')

for pin in range(0, 1000):
    pin_packed = struct.pack('<H', pin)
    res = do_run(pin_packed)
    if res.startswith('hwsec'):
        print(f"Result with pin {pin_packed!r} ({pin}): {res!r}")
        break
