summaryrefslogtreecommitdiff
path: root/zep2newt.py
blob: 087f2a4263c6bd4a3f22b240d7f74a7754212ae5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
#!/usr/bin/python2
import hashlib
import mmap
import os
import struct
import sys
from argparse import ArgumentParser
from newtimg import *
from ctypes import *

DEBUG = False

################################################################################
def get_args():
    parser = ArgumentParser(description='Script to create images on a format \
                            that Mynewts bootloader expects')

    parser.add_argument('--bin', required=True, dest='binary_file', \
                        help='Name of *.bin file (input)')

    parser.add_argument('--key', required=False, dest='key_file', \
                        help='Name of private key file (*.pem format)')

    parser.add_argument('--out', required=False, dest='image_file', \
                        default='zephyr.img.bin', \
                        help='Name of *.img file (output)')

    parser.add_argument('--sig', required=False, dest='sig_type', \
                        default='SHA256', \
                        help='Type of signature <SHA256|RSA|EC>')

    parser.add_argument('--off', required=False, dest='flash_offs_addr', \
                        default='0x8000', \
                        help='Offset for the binary in flash (at what address \
                        should it be flashed?)')

    parser.add_argument('--vtoff', required=False, dest='vtable_offs', \
                        default=str(hex(OFFSET_VECTOR_TABLE)), \
                        help='Offset to vector table in HEX')

    parser.add_argument('--bit', required=False, \
                        help='Flash boot image trailer magic')

    parser.add_argument('--verbose', required=False, action="store_true", \
                        default=False, \
                        help='Enable verbose mode')

    parser.add_argument('--version', action='version', version='%(prog)s 1.0')

    parser.add_argument('-f', required=False, action="store_true", \
                        default=False, \
                        help='Flash using JLinkExe')

    return parser.parse_args()

################################################################################
def create_header(binary_file, sig_type, vtable_offs):
    """
    Create a header on a format that Mynewt's bootloader expects. Based on
    signature type it updates, TLV size, flags and key ID (if RSA). This
    function also updates/stores the offset to the vector table in the binary.
    For example in Mynewt the offset is 0x20, but in Zephyr it's expected that
    the vector table is at an address % 128 and therefore 0x80 should be used
    there instead (which is the default in this script).
    """
    # The SHA256 hash is always used and therefore we use that as default
    tlv_size = SHA256_DIGEST_SIZE + 4
    flags = IMAGE_F_SHA256
    key_id = 0

    if sig_type == "RSA":
        tlv_size = tlv_size + 4 + RSA_SIZE
        keyId = 0 # FIXME
        flags =  IMAGE_F_PKCS15_RSA2048_SHA256 | IMAGE_F_SHA256
    elif sig_type == "EC":
        tlv_size = tlv_size + 4 + ECDSA_SIZE
        flags = IMAGE_F_ECDSA224_SHA256 | IMAGE_F_SHA256

    image_size = 0
    try:
        # Get the correct size for the image
        image_size = os.path.getsize(binary_file)
        if DEBUG:
            print("[*] Binary size %d (0x%x) of %s" % (image_size, image_size, binary_file))
    except (OSError, IOError):
        print("[ERROR]: Cannot open %s" % binary_file)
        sys.exit(1)

    hdr = bytearray(struct.pack('I', IMAGE_MAGIC) +
                    struct.pack('H', tlv_size) +
                    struct.pack('B', key_id) + # Key ID
                    struct.pack('B', 0) + # PAD 1
                    struct.pack('H', vtable_offs) + # New HDR SIZE
                    struct.pack('H', 0) + # PAD 2
                    struct.pack('I', image_size) + # img size
                    struct.pack('I', flags) + # Flags
                    struct.pack('B', 1) + # Major
                    struct.pack('B', 0) + # Minor
                    struct.pack('H', 0) + # Revision
                    struct.pack('I', 0) + # Build number
                    struct.pack('I', 0)) # PAD3
    if DEBUG:
        try:
            with open(binary_file + ".hdr", "w+b") as f:
                f.write(hdr)
                f.close()
        except (OSError, IOError):
            print("[ERROR]: Cannot write to %s (!) " % (binary_file + ".hdr"))
            sys.exit(1)
    return hdr

################################################################################
def write_partial_img(binary_file, image_file, hdr, vtable_offs):
    try:
        with open(binary_file, "rb") as f:
            image = f.read()
            f.close()
        if DEBUG:
            print("[*] Read %d bytes from %s" % (len(image), binary_file))

    except (OSError, IOError):
        print("[ERROR]: Cannot open %s" % (binary_file))
        sys.exit(1)

    try:
        with open(image_file, "w+b") as f:
            f.write(hdr)
            # Calculate how much to pad before the actual image with the vector
            # table starts.
            f.write('\xFF' * (vtable_offs - IMAGE_HEADER_SIZE))
            f.write(image)
            f.close()
        if DEBUG:
            sz = os.path.getsize(image_file)
            print("[*] Wrote %d (0x%x) bytes to %s" % (sz, sz, image_file))

    except (OSError, IOError):
        print("[ERROR]: Cannot write to %s" % (image_file))
        sys.exit(1)


################################################################################
def calculate_hash(image_file):
    sha256 = hashlib.sha256()
    try:
        with open(image_file, "rb") as f:
            sha256.update(f.read())
            f.close()

    except IOError:
        print("[ERROR]: Cannot open %s" % (image_file))
        sys.exit(1)

    digest = sha256.hexdigest()
    if DEBUG:
        print("[*] Hash of intermediate image: %s" % digest)
    return digest

################################################################################
def append_hash(image_file, digest):
    try:
        with open(image_file, "ab") as f:

            # Start by settings the TLV type
            # https://github.com/apache/incubator-mynewt-newt/blob/master/newt/image/image.go#L109-L116
            tlv_type = struct.pack('b', IMAGE_TLV_SHA256)

            # Next 1 byte padding
            tlv_pad = '\x00'

            # Finally the size of the TLV, for SHA256 that is 32 bytes
            tlv_len = struct.pack('h', SHA256_DIGEST_SIZE)

            f.write(tlv_type)
            f.write(tlv_pad)
            f.write(tlv_len)
            f.write(digest.decode('hex'))
            f.close()

    except IOError:
        print("[ERROR]: Cannot open/append to %s" % (image_file))
        sys.exit(1)

################################################################################
def create_jlink_script(image_file, offset, erase):
    """
    Creates a jlink script to flash the created binary.

    @erase: whether the script first shall erase or not when flashing.
    @offset: where in flash to store the image.
    """
    jlink_file = "flash_zephyr.jlink"
    try:
        with open(jlink_file, "w+") as f:
            f.write("device nrf52\n")
            f.write("power on\n")
            f.write("sleep 10\n")
            f.write("si 1\n")
            f.write("speed auto\n")
            if erase:
                f.write("erase\n")
            f.write("loadfile %s %s\n" % (image_file, hex(int(offset, 16))))
            f.write("q\n")
            f.close()
        if DEBUG:
            print("\n[*] To flash the Image for nrf52 run:")
            print("     JLinkExe -CommanderScript %s" % jlink_file)

    except IOError:
        print("[ERROR]: Cannot create to %s" % (jlink_file))
        sys.exit(1)

################################################################################
def create_jlink_bit_script(bit_file, bitoffset="0x7bff8"):
    jlink_file = "flash_bit.jlink"
    try:
        with open(jlink_file, "w+") as f:
            f.write("device nrf52\n")
            f.write("power on\n")
            f.write("sleep 10\n")
            f.write("si 1\n")
            f.write("speed auto\n")
            f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
            f.write("q\n")
            f.close()
        if DEBUG:
            print("\n[*] To flash Boot Image Trailer for nrf52 run:")
            print("     JLinkExe -CommanderScript %s" % jlink_file)

    except IOError:
        print("[ERROR]: Cannot create to %s" % (jlink_file))
        sys.exit(1)

################################################################################
def create_jlink_clear_bit_script(bit_file, bitoffset="0x7bff8"):
    jlink_file = "flash_clear_bit.jlink"
    try:
        with open(jlink_file, "w+") as f:
            f.write("device nrf52\n")
            f.write("power on\n")
            f.write("sleep 10\n")
            f.write("si 1\n")
            f.write("speed auto\n")
            f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
            f.write("q\n")
            f.close()
        if DEBUG:
            print("\n[*] To clear Boot Image Trailer for nrf52 run:")
            print("     JLinkExe -CommanderScript %s" % jlink_file)

    except IOError:
        print("[ERROR]: Cannot create to %s" % (jlink_file))
        sys.exit(1)

################################################################################
def main(argv):
    args = get_args()
    erase = False

    if args.verbose:
        for a in vars(args):
            print("Arg -> %s: %s" % (a, getattr(args, a)))
        global DEBUG
        DEBUG = True

    # Since it's a hex string, let's convert to an integer instead
    vtable_offs = int(args.vtable_offs, 16)

    # Create the header first
    hdr = create_header(args.binary_file, args.sig_type, vtable_offs)

    # Write the image itself
    write_partial_img(args.binary_file, args.image_file, hdr, vtable_offs)

    # Now we have a header and the binary itself and we should get the hash of
    # those concatenated.
    digest = calculate_hash(args.image_file)
    append_hash(args.image_file, digest)
    print("[*] Successfully created: %s" % args.image_file)

    # Misc function related to flashing
    create_jlink_script(args.image_file, args.flash_offs_addr, erase)
    curdir = os.path.dirname(sys.argv[0])
    create_jlink_bit_script(curdir + "/boot_image_trailer.bin")
    create_jlink_clear_bit_script(curdir + "/empty_boot_image_trailer.bin")

    # Should we try flash?
    if args.f:
        os.system("JLinkExe -CommanderScript flash_zephyr.jlink")

if __name__ == "__main__":
    main(sys.argv)