1*817cd834SEugeniy Paltsev#!/usr/bin/env python3
2*817cd834SEugeniy Paltsev
3*817cd834SEugeniy Paltsev#we can use binascii instead of zlib
4*817cd834SEugeniy Paltsevimport os, getopt, sys, zlib
5*817cd834SEugeniy Paltsevfrom elftools.elf.elffile import ELFFile
6*817cd834SEugeniy Paltsev
7*817cd834SEugeniy Paltsev
8*817cd834SEugeniy Paltsevdef usage(exit_code):
9*817cd834SEugeniy Paltsev    print("typical usage:")
10*817cd834SEugeniy Paltsev    print("AXS101:")
11*817cd834SEugeniy Paltsev    print(sys.argv[0] + \
12*817cd834SEugeniy Paltsev        " --header-type v1 --arc-id 0x434 --spi-flash-offset 0x0 --image u-boot.bin --elf u-boot")
13*817cd834SEugeniy Paltsev    print("AXS103:")
14*817cd834SEugeniy Paltsev    print(sys.argv[0] + \
15*817cd834SEugeniy Paltsev        " --header-type v2 --arc-id 0x53 --spi-flash-offset 0x200000 --image u-boot.bin --elf u-boot")
16*817cd834SEugeniy Paltsev    sys.exit(exit_code)
17*817cd834SEugeniy Paltsev
18*817cd834SEugeniy Paltsev
19*817cd834SEugeniy Paltsevdef elf_get_entry(filename):
20*817cd834SEugeniy Paltsev    with open(filename, 'rb') as f:
21*817cd834SEugeniy Paltsev        elffile = ELFFile(f)
22*817cd834SEugeniy Paltsev        return elffile.header['e_entry']
23*817cd834SEugeniy Paltsev
24*817cd834SEugeniy Paltsev
25*817cd834SEugeniy Paltsevdef calc_check_sum(filename):
26*817cd834SEugeniy Paltsev    # Calculate u-boot image check_sum: it is sum of all u-boot binary bytes
27*817cd834SEugeniy Paltsev    with open(filename, "rb") as file:
28*817cd834SEugeniy Paltsev        ba = bytearray(file.read())
29*817cd834SEugeniy Paltsev    return sum(ba) & 0xFF
30*817cd834SEugeniy Paltsev
31*817cd834SEugeniy Paltsev
32*817cd834SEugeniy Paltsevdef arg_verify(uboot_bin_filename, uboot_elf_filename, header_type):
33*817cd834SEugeniy Paltsev    if not os.path.isfile(uboot_bin_filename):
34*817cd834SEugeniy Paltsev        print("uboot bin file not exists: " + uboot_bin_filename)
35*817cd834SEugeniy Paltsev        sys.exit(2)
36*817cd834SEugeniy Paltsev
37*817cd834SEugeniy Paltsev    if not os.path.isfile(uboot_elf_filename):
38*817cd834SEugeniy Paltsev        print("uboot elf file not exists: " + uboot_elf_filename)
39*817cd834SEugeniy Paltsev        sys.exit(2)
40*817cd834SEugeniy Paltsev
41*817cd834SEugeniy Paltsev    if header_type not in ("v1", "v2"):
42*817cd834SEugeniy Paltsev        print("unknown header type: " + header_type)
43*817cd834SEugeniy Paltsev        print("choose between 'v1' (most likely AXS101) and 'v2' (most likely AXS103)")
44*817cd834SEugeniy Paltsev        sys.exit(2)
45*817cd834SEugeniy Paltsev
46*817cd834SEugeniy Paltsev
47*817cd834SEugeniy Paltsevdef main():
48*817cd834SEugeniy Paltsev    try:
49*817cd834SEugeniy Paltsev        opts, args = getopt.getopt(sys.argv[1:],
50*817cd834SEugeniy Paltsev            "ht:a:s:i:l:e:",
51*817cd834SEugeniy Paltsev            ["help", "header-type=", "arc-id=", "spi-flash-offset=", "image=", "elf="])
52*817cd834SEugeniy Paltsev    except getopt.GetoptError as err:
53*817cd834SEugeniy Paltsev        print(err)
54*817cd834SEugeniy Paltsev        usage(2)
55*817cd834SEugeniy Paltsev
56*817cd834SEugeniy Paltsev    # default filenames
57*817cd834SEugeniy Paltsev    uboot_elf_filename  = "u-boot"
58*817cd834SEugeniy Paltsev    uboot_bin_filename  = "u-boot.bin"
59*817cd834SEugeniy Paltsev    headerised_filename = "u-boot.head"
60*817cd834SEugeniy Paltsev    uboot_scrypt_file   = "u-boot-update.txt"
61*817cd834SEugeniy Paltsev
62*817cd834SEugeniy Paltsev    # default values
63*817cd834SEugeniy Paltsev    spi_flash_offset    = 0x200000
64*817cd834SEugeniy Paltsev    header_type         = "v2"
65*817cd834SEugeniy Paltsev    arc_id              = 0x53
66*817cd834SEugeniy Paltsev
67*817cd834SEugeniy Paltsev    # initial header values: place where preloader will store u-boot binary,
68*817cd834SEugeniy Paltsev    # should be equal to CONFIG_SYS_TEXT_BASE
69*817cd834SEugeniy Paltsev    image_copy_adr  = 0x81000000
70*817cd834SEugeniy Paltsev
71*817cd834SEugeniy Paltsev    # initial constant header values, do not change these values
72*817cd834SEugeniy Paltsev    magic1          = 0xdeadbeafaf # big endian byte order
73*817cd834SEugeniy Paltsev    magic2          = [            # big endian byte order
74*817cd834SEugeniy Paltsev    0x20202a2020202020202020202a20202020207c5c2e20202020202e2f7c20202020207c2d,
75*817cd834SEugeniy Paltsev    0x2e5c2020202f2e2d7c20202020205c2020602d2d2d6020202f20202020202f205f202020,
76*817cd834SEugeniy Paltsev    0x205f20205c20202020207c205f60712070205f207c2020202020272e5f3d2f205c3d5f2e,
77*817cd834SEugeniy Paltsev    0x272020202020202020605c202f60202020202020202020202020206f2020202020202020]
78*817cd834SEugeniy Paltsev
79*817cd834SEugeniy Paltsev    for opt, arg in opts:
80*817cd834SEugeniy Paltsev        if opt in ('-h', "--help"):        usage(0)
81*817cd834SEugeniy Paltsev        if opt in ('-t', "--header-type"): header_type           = arg
82*817cd834SEugeniy Paltsev        if opt in ('-a', "--arc-id"):      arc_id                = int(arg, 16)
83*817cd834SEugeniy Paltsev        if opt in ('-s', "--spi-flash-offset"): spi_flash_offset = int(arg, 16)
84*817cd834SEugeniy Paltsev        if opt in ('-i', "--image"):       uboot_bin_filename    = arg
85*817cd834SEugeniy Paltsev        if opt in ('-e', "--elf"):         uboot_elf_filename    = arg
86*817cd834SEugeniy Paltsev
87*817cd834SEugeniy Paltsev    arg_verify(uboot_bin_filename, uboot_elf_filename, header_type)
88*817cd834SEugeniy Paltsev
89*817cd834SEugeniy Paltsev    uboot_img_size = os.path.getsize(uboot_bin_filename)
90*817cd834SEugeniy Paltsev    jump_address = elf_get_entry(uboot_elf_filename)
91*817cd834SEugeniy Paltsev    check_sum = calc_check_sum(uboot_bin_filename)
92*817cd834SEugeniy Paltsev
93*817cd834SEugeniy Paltsev    # Calculate header adresses depend on header type
94*817cd834SEugeniy Paltsev    if header_type == "v2":
95*817cd834SEugeniy Paltsev        image_copy_adr -= 0x4
96*817cd834SEugeniy Paltsev        uboot_img_size += 0x4
97*817cd834SEugeniy Paltsev        # we append image so we need to append checksum
98*817cd834SEugeniy Paltsev        jmpchk_sum = sum(jump_address.to_bytes(4, byteorder='big'))
99*817cd834SEugeniy Paltsev        check_sum = (check_sum + jmpchk_sum) & 0xFF
100*817cd834SEugeniy Paltsev        imade_jump_append = True
101*817cd834SEugeniy Paltsev    else:
102*817cd834SEugeniy Paltsev        imade_jump_append = False
103*817cd834SEugeniy Paltsev
104*817cd834SEugeniy Paltsev    # write header to file
105*817cd834SEugeniy Paltsev    with open(headerised_filename, "wb") as file:
106*817cd834SEugeniy Paltsev        file.write(arc_id.to_bytes(2, byteorder='little'))
107*817cd834SEugeniy Paltsev        file.write(uboot_img_size.to_bytes(4, byteorder='little'))
108*817cd834SEugeniy Paltsev        file.write(check_sum.to_bytes(1, byteorder='little'))
109*817cd834SEugeniy Paltsev        file.write(image_copy_adr.to_bytes(4, byteorder='little'))
110*817cd834SEugeniy Paltsev        file.write(magic1.to_bytes(5, byteorder='big'))
111*817cd834SEugeniy Paltsev        for i in range(16): file.write(0x00.to_bytes(1, byteorder='little'))
112*817cd834SEugeniy Paltsev        for byte in magic2: file.write(byte.to_bytes(36, byteorder='big'))
113*817cd834SEugeniy Paltsev        for i in range(224 - len(magic2) * 36):
114*817cd834SEugeniy Paltsev            file.write(0x00.to_bytes(1, byteorder='little'))
115*817cd834SEugeniy Paltsev        if imade_jump_append:
116*817cd834SEugeniy Paltsev            file.write(jump_address.to_bytes(4, byteorder='little'))
117*817cd834SEugeniy Paltsev
118*817cd834SEugeniy Paltsev    # append u-boot image to header
119*817cd834SEugeniy Paltsev    with open(headerised_filename, "ab") as fo:
120*817cd834SEugeniy Paltsev        with open(uboot_bin_filename,'rb') as fi:
121*817cd834SEugeniy Paltsev            fo.write(fi.read())
122*817cd834SEugeniy Paltsev
123*817cd834SEugeniy Paltsev    # calc u-boot headerised image CRC32 (will be used by uboot update
124*817cd834SEugeniy Paltsev    # command for check)
125*817cd834SEugeniy Paltsev    headerised_image_crc = ""
126*817cd834SEugeniy Paltsev    with open(headerised_filename, "rb") as fi:
127*817cd834SEugeniy Paltsev        headerised_image_crc = hex(zlib.crc32(fi.read()) & 0xffffffff)
128*817cd834SEugeniy Paltsev
129*817cd834SEugeniy Paltsev    load_addr = 0x81000000
130*817cd834SEugeniy Paltsev    crc_store_adr = load_addr - 0x8
131*817cd834SEugeniy Paltsev    crc_calc_adr = crc_store_adr - 0x4
132*817cd834SEugeniy Paltsev    load_size = os.path.getsize(headerised_filename)
133*817cd834SEugeniy Paltsev    crc_calc_cmd = \
134*817cd834SEugeniy Paltsev        "crc32 " + hex(load_addr) + " " + hex(load_size) + " " + hex(crc_calc_adr)
135*817cd834SEugeniy Paltsev    crc_check_cmd = \
136*817cd834SEugeniy Paltsev        "mw.l " + hex(crc_store_adr) + " " + headerised_image_crc + " && " + \
137*817cd834SEugeniy Paltsev        crc_calc_cmd + " && " + \
138*817cd834SEugeniy Paltsev        "cmp.l " + hex(crc_store_adr) + " " + hex(crc_calc_adr) + " 1"
139*817cd834SEugeniy Paltsev
140*817cd834SEugeniy Paltsev    # make errase size to be allighned by 64K
141*817cd834SEugeniy Paltsev    if load_size & 0xFFFF == 0:
142*817cd834SEugeniy Paltsev        errase_size = load_size
143*817cd834SEugeniy Paltsev    else:
144*817cd834SEugeniy Paltsev        errase_size = load_size - (load_size & 0xFFFF) + 0x10000
145*817cd834SEugeniy Paltsev
146*817cd834SEugeniy Paltsev    # Hack to handle n25*** flash protect ops weirdness:
147*817cd834SEugeniy Paltsev    # protect unlock return fail status is region is already unlock (entire or
148*817cd834SEugeniy Paltsev    # partially). Same for lock ops.
149*817cd834SEugeniy Paltsev    # As there is no possibility to check current flash status pretend
150*817cd834SEugeniy Paltsev    # unlock & lock always success.
151*817cd834SEugeniy Paltsev    sf_unlock_cmd = \
152*817cd834SEugeniy Paltsev        "if sf protect unlock 0x0 0x4000000 ; then true ; else true ; fi"
153*817cd834SEugeniy Paltsev    sf_lock_cmd = \
154*817cd834SEugeniy Paltsev        "if sf protect lock 0x0 0x4000000 ; then true ; else true ; fi"
155*817cd834SEugeniy Paltsev
156*817cd834SEugeniy Paltsev    # u-bood CMD to load u-bood with header to SPI flash
157*817cd834SEugeniy Paltsev    sf_load_image_cmd = \
158*817cd834SEugeniy Paltsev        "fatload mmc 0:1 " + hex(load_addr) + " " + headerised_filename + " && " + \
159*817cd834SEugeniy Paltsev        "sf probe 0:0 && " + \
160*817cd834SEugeniy Paltsev        sf_unlock_cmd + " && " + \
161*817cd834SEugeniy Paltsev        "sf erase " + hex(spi_flash_offset) + " " + hex(errase_size) + " && " + \
162*817cd834SEugeniy Paltsev        "sf write " + hex(load_addr) + " " + hex(spi_flash_offset) + " " + hex(load_size) + " && " + \
163*817cd834SEugeniy Paltsev        sf_lock_cmd
164*817cd834SEugeniy Paltsev
165*817cd834SEugeniy Paltsev    update_uboot_cmd = sf_load_image_cmd + " && echo \"u-boot update: OK\""
166*817cd834SEugeniy Paltsev
167*817cd834SEugeniy Paltsev    with open(uboot_scrypt_file, "wb") as fo:
168*817cd834SEugeniy Paltsev        fo.write(update_uboot_cmd.encode('ascii'))
169*817cd834SEugeniy Paltsev
170*817cd834SEugeniy Paltsev
171*817cd834SEugeniy Paltsevif __name__ == "__main__":
172*817cd834SEugeniy Paltsev    try:
173*817cd834SEugeniy Paltsev        main()
174*817cd834SEugeniy Paltsev    except Exception as err:
175*817cd834SEugeniy Paltsev        print(err)
176*817cd834SEugeniy Paltsev        sys.exit(2)
177