import argparse
import struct

from capstone import *
from capstone.arm import *


dis = Cs(CS_ARCH_ARM, CS_MODE_LITTLE_ENDIAN + CS_MODE_THUMB + CS_MODE_MCLASS)
dis.detail = True
BASE = 0x10000
RAM = 0x100000c8


def reset_handler_address(bin):
    return struct.unpack("<I", bin[4:8])[0] - 1


def find_data_start(bin, reset_addr):
    # locate copy loop
    for ins in dis.disasm(bin[(reset_addr - BASE):], reset_addr):
        if ins.mnemonic == "b":
            copy_loop = ins.operands[0].value.imm
            break
    print("found RAM copy loop: " + hex(copy_loop))

    # separate condition checking from actual copy operations
    for ins in dis.disasm(bin[(copy_loop - BASE):], copy_loop):
        if ins.mnemonic == "bhs":
            copy_op = ins.address + ins.size
            break
    print("found copy operation: " + hex(copy_op)[:-1])

    # get the destination address!
    for ins in dis.disasm(bin[(copy_op - BASE):], copy_op):
        if ins.mnemonic == "ldr":
            for op in ins.operands:
                if op.type == ARM_OP_MEM:
                    copy_dest_loc = ins.address + ins.size + op.mem.disp + 2
                    copy_dest = struct.unpack("<I", bin[(copy_dest_loc - BASE):copy_dest_loc - BASE + 4])[0]
                    break
            break
    return copy_dest


def find_main(bin, reset):
    pastSystemInit = False
    for ins in dis.disasm(bin[(reset - BASE):], reset):
        if ins.mnemonic == "bl":
            if pastSystemInit == False:
                pastSystemInit = True
            else:
                return ins.operands[0].value.imm


def find_signals_table(bin, main):
    for ins in dis.disasm(bin[(main - BASE):], main):
        if ins.mnemonic == "bl":
            initVehicleInterface = ins.operands[0].value.imm
            break
    print("found initializeVehicleInterface: " + hex(initVehicleInterface))

    init_can_position = 8
    call_position = 0
    for ins in dis.disasm(bin[(initVehicleInterface - BASE):], initVehicleInterface):
        if ins.mnemonic == "bl":
            call_position += 1
        if ins.mnemonic == "bl" and call_position == init_can_position:
            initAllCan = ins.operands[0].value.imm
            break
    print("found initializeAllCAN: " + hex(initAllCan))

    get_signals_position = 3
    call_position = 0
    for ins in dis.disasm(bin[(initAllCan - BASE):], initAllCan):
        if ins.mnemonic == "bl":
            call_position += 1
        if ins.mnemonic == "bl" and call_position == get_signals_position:
            get_signals = ins.operands[0].value.imm
            break
    print("found getSignals: " + hex(get_signals))

    for ins in dis.disasm(bin[(get_signals - BASE):], get_signals):
        if ins.mnemonic == "ldr":
            for op in ins.operands:
                if op.type == ARM_OP_MEM:
                    signals_data_loc = ins.address + ins.size + op.mem.disp
                    signals = struct.unpack("<I", bin[(signals_data_loc - BASE): signals_data_loc - BASE + 4])[0]
                    break
            break
    return signals


def read_structs(bin, data_start, signals):
    ptr = signals - BASE
    structs = []
    while True:
        msg = struct.unpack("<I", bin[(ptr):ptr + 4])[0]
        msg = (msg - RAM) + data_start
        try:
            msg_id = struct.unpack("<I", bin[(msg - BASE + 4):(msg - BASE + 8)])[0]
        except:
            break
        name_off = struct.unpack("<I", bin[(ptr + 4):ptr + 8])[0]
        name = ""
        while True:
            c = struct.unpack("s", bin[(name_off - BASE):(name_off - BASE + 1)])[0]
            if c == "\x00":
                break
            name += c
            name_off += 1
        bitPos = struct.unpack("B", bin[(ptr + 8):ptr + 9])[0]
        bitSize = struct.unpack("B", bin[(ptr + 9):ptr + 10])[0]
        factor = struct.unpack("<f", bin[(ptr + 12):ptr + 16])[0]
        offset = struct.unpack("<f", bin[(ptr + 16):ptr + 20])[0]
        minVal = struct.unpack("<f", bin[(ptr + 20):ptr + 24])[0]
        maxVal = struct.unpack("<f", bin[(ptr + 24):ptr + 28])[0]
        structs.append((msg_id, name, bitPos, bitSize, factor, offset, minVal, maxVal))
        ptr += 0x44
    return structs


def print_structs(structs):
    print("ID|name|bitPos|bitSize|factor|offset|minVal|maxVal")
    for i in structs:
        print(i)


def go(bin):
    reset = reset_handler_address(bin)
    print("found reset_handler: " + hex(reset))
    __data_start__ = find_data_start(bin, reset)
    print("found __data_start__: " + hex(__data_start__))
    main = find_main(bin, reset)
    print("found main: " + hex(main))
    signals = find_signals_table(bin, main)
    print("found SIGNALS: " + hex(signals))
    signals = (signals - RAM) + __data_start__
    print("signals file offset: " + hex(signals))
    structs = read_structs(bin, __data_start__, signals)
    return structs


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="kek")
    parser.add_argument("firmware", help="path to firmware file")
    args = parser.parse_args()
    with open(args.firmware) as f:
        bin = f.read()
        structs = go(bin)
        print_structs(structs)