from argparse import ArgumentParser
import mmap
from struct import unpack
import binascii

FLASHTOOL_CFG_MAGIC = 0x544F4F4C
FLASHTOOL_CFG_MAGIC_64 = 0x544F4F5C
FLASHTOOL_CFG_SIZE = 76
FLASHTOOL_NON_SLA_FORBID_MAGIC = 0x544F4F4D
FLASHTOOL_NON_SLA_FORBID_MAGIC_64 = 0x544F4F5D
FLASHTOOL_NON_SLA_FORBID_CFG_SIZE = 52
AND_SEC_CTRL_SIZE = 52
AND_SEC_BOOT_CHECK_PART_SIZE = 90


def main():
    parser = ArgumentParser()
    parser.add_argument('file')
    args = parser.parse_args()

    file = open(args.file, 'rb')
    data = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
    rominfo_offset = data.find(b'AND_ROMINFO_v')
    if rominfo_offset == -1:
        print('ROMINFO not found (is this preloader file?)')
        exit(1)

    version = unpack('<I', data[rominfo_offset+16:rominfo_offset+20])[0]
    platform_id = data[rominfo_offset+20:rominfo_offset+36].decode('utf-8')
    project_id = data[rominfo_offset+36:rominfo_offset+52].decode('utf-8')
    secro_exists = unpack('<I', data[rominfo_offset+52:rominfo_offset+56])[0]
    secro_offset = unpack('<I', data[rominfo_offset+56:rominfo_offset+60])[0]
    secro_length = unpack('<I', data[rominfo_offset+60:rominfo_offset+64])[0]
    ac_offset = unpack('<I', data[rominfo_offset+64:rominfo_offset+68])[0]
    ac_length = unpack('<I', data[rominfo_offset+68:rominfo_offset+72])[0]
    seccfg_offset = unpack('<I', data[rominfo_offset+72:rominfo_offset+76])[0]
    seccfg_length = unpack('<I', data[rominfo_offset+76:rominfo_offset+80])[0]
    flash_tool_cfg_magic = unpack(
        '<I', data[rominfo_offset+80:rominfo_offset+84])
    flash_tool_nlsa_forbid_magic = unpack(
        '<I', data[rominfo_offset+80+FLASHTOOL_CFG_SIZE:rominfo_offset+80+FLASHTOOL_CFG_SIZE+4])[0]
    secctrl_magic = data[rominfo_offset+80+FLASHTOOL_CFG_SIZE+FLASHTOOL_NON_SLA_FORBID_CFG_SIZE:rominfo_offset +
                         80+FLASHTOOL_CFG_SIZE+FLASHTOOL_NON_SLA_FORBID_CFG_SIZE+16]
    secboot_check_part_offset = rominfo_offset+80+FLASHTOOL_CFG_SIZE + \
        FLASHTOOL_NON_SLA_FORBID_CFG_SIZE+AND_SEC_CTRL_SIZE+18

    print('version:          %d' % version)
    print('platform id:      %s' % platform_id)
    print('project id:       %s' % project_id)
    print('secro exists:     %d' % secro_exists)
    print('secro offset:     %#08x' % secro_offset)
    print('secro length:     %#08x' % secro_length)
    print('ac offset:        %#08x' % ac_offset)
    print('ac length:        %#08x' % ac_length)
    print('seccfg offset:    %#08x' % seccfg_offset)
    print('seccfg length:    %#08x' % seccfg_length)
    # TODO: print flash tool config
    if flash_tool_cfg_magic == FLASHTOOL_CFG_MAGIC:
        print('flashtool config: 32-bit')
    elif flash_tool_cfg_magic == FLASHTOOL_CFG_MAGIC_64:
        print('flashtool config: 64-bit')
    else:
        print('flashtool config: invalid (magic %#08x)' % flash_tool_cfg_magic)

    if flash_tool_cfg_magic == FLASHTOOL_NON_SLA_FORBID_MAGIC:
        print('flashtool non sla download forbid: 32-bit')
    elif flash_tool_cfg_magic == FLASHTOOL_NON_SLA_FORBID_MAGIC_64:
        print('flashtool non sla download forbid: 64-bit')
    else:
        print('flashtool non sla download forbid: invalid (magic %#08x)' %
              flash_tool_nlsa_forbid_magic)

    if secctrl_magic[0:13] == b'AND_SECCTRL_v':
        print('secctrl:          valid')
    else:
        print('secctrl:          invalid')

    print('Secure Boot protected partitions:')
    for i in range(9):
        partition_name = data[secboot_check_part_offset +
                              i*10:secboot_check_part_offset+i*10+10]
        if partition_name == b'\x00'*10:
            continue
        partition_name = partition_name.decode('utf-8')
        print('   %s' % partition_name)

    seckey_offset = secboot_check_part_offset+AND_SEC_BOOT_CHECK_PART_SIZE
    seckey_magic = data[seckey_offset:seckey_offset+12]
    seckey_version = unpack('<I', data[seckey_offset+16:seckey_offset+20])
    seckey_img_auth_rsa_key_n = data[seckey_offset+20:seckey_offset+276]
    seckey_img_auth_rsa_key_e = data[seckey_offset+276:+seckey_offset+281]
    seckey_sml_aes_key = data[seckey_offset+281:seckey_offset+281+32]
    seckey_crypto_seed = data[seckey_offset+281+32:seckey_offset+281+32+16]

    if seckey_magic == b'AND_SECRO_v\x00' or seckey_magic == b'AND_SECKEY_v':
        print('seckey:           valid')
        print('seckey version:   %d' % seckey_version)
        assert len(seckey_img_auth_rsa_key_n) == 256
        print('img auth rsa N:   ', end='')
        for i in range(len(seckey_img_auth_rsa_key_n)):
            print("\\x%02x" % seckey_img_auth_rsa_key_n[i],end='')
        print('')
        print('img auth rsa E:   ', seckey_img_auth_rsa_key_e)
        print('sml aes key:      ', end='')
        assert len(seckey_sml_aes_key) == 32
        for i in range(len(seckey_sml_aes_key)):
            print("\\x%02x" % seckey_sml_aes_key[i], end='')
        print('')
        print('seed:             ',seckey_crypto_seed)

    else:
        print('seckey:           invalid')

    file.close()


if __name__ == '__main__':
    main()
