1#!/usr/bin/env python3
2
3# Microsoft UF2
4#
5# The MIT License (MIT)
6#
7# Copyright (c) Microsoft Corporation
8#
9# All rights reserved.
10#
11# Permission is hereby granted, free of charge, to any person obtaining a copy
12# of this software and associated documentation files (the "Software"), to deal
13# in the Software without restriction, including without limitation the rights
14# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
15# copies of the Software, and to permit persons to whom the Software is
16# furnished to do so, subject to the following conditions:
17#
18# The above copyright notice and this permission notice shall be included in all
19# copies or substantial portions of the Software.
20#
21# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
22# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
23# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
24# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
25# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
26# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
27# SOFTWARE.
28
29import sys
30import struct
31import subprocess
32import re
33import os
34import os.path
35import argparse
36
37
38UF2_MAGIC_START0 = 0x0A324655  # "UF2\n"
39UF2_MAGIC_START1 = 0x9E5D5157  # Randomly selected
40UF2_MAGIC_END = 0x0AB16F30  # Ditto
41
42families = {
43    "SAMD21": 0x68ED2B88,
44    "SAMD51": 0x55114460,
45    "NRF52": 0x1B57745F,
46    "STM32F1": 0x5EE21072,
47    "STM32F4": 0x57755A57,
48    "ATMEGA32": 0x16573617,
49}
50
51INFO_FILE = "/INFO_UF2.TXT"
52
53appstartaddr = 0x2000
54familyid = 0x0
55
56
57def is_uf2(buf):
58    w = struct.unpack("<II", buf[0:8])
59    return w[0] == UF2_MAGIC_START0 and w[1] == UF2_MAGIC_START1
60
61
62def is_hex(buf):
63    try:
64        w = buf[0:30].decode("utf-8")
65    except UnicodeDecodeError:
66        return False
67    if w[0] == ":" and re.match(b"^[:0-9a-fA-F\r\n]+$", buf):
68        return True
69    return False
70
71
72def convert_from_uf2(buf):
73    global appstartaddr
74    numblocks = len(buf) // 512
75    curraddr = None
76    outp = b""
77    for blockno in range(numblocks):
78        ptr = blockno * 512
79        block = buf[ptr : ptr + 512]
80        hd = struct.unpack(b"<IIIIIIII", block[0:32])
81        if hd[0] != UF2_MAGIC_START0 or hd[1] != UF2_MAGIC_START1:
82            print("Skipping block at " + ptr + "; bad magic")
83            continue
84        if hd[2] & 1:
85            # NO-flash flag set; skip block
86            continue
87        datalen = hd[4]
88        if datalen > 476:
89            assert False, "Invalid UF2 data size at " + ptr
90        newaddr = hd[3]
91        if curraddr == None:
92            appstartaddr = newaddr
93            curraddr = newaddr
94        padding = newaddr - curraddr
95        if padding < 0:
96            assert False, "Block out of order at " + ptr
97        if padding > 10 * 1024 * 1024:
98            assert False, "More than 10M of padding needed at " + ptr
99        if padding % 4 != 0:
100            assert False, "Non-word padding size at " + ptr
101        while padding > 0:
102            padding -= 4
103            outp += b"\x00\x00\x00\x00"
104        outp += block[32 : 32 + datalen]
105        curraddr = newaddr + datalen
106    return outp
107
108
109def convert_to_carray(file_content):
110    outp = "const unsigned char bindata[] __attribute__((aligned(16))) = {"
111    for i in range(len(file_content)):
112        if i % 16 == 0:
113            outp += "\n"
114        outp += "0x%02x, " % ord(file_content[i])
115    outp += "\n};\n"
116    return outp
117
118
119def convert_to_uf2(file_content):
120    global familyid
121    datapadding = b""
122    while len(datapadding) < 512 - 256 - 32 - 4:
123        datapadding += b"\x00\x00\x00\x00"
124    numblocks = (len(file_content) + 255) // 256
125    outp = b""
126    for blockno in range(numblocks):
127        ptr = 256 * blockno
128        chunk = file_content[ptr : ptr + 256]
129        flags = 0x0
130        if familyid:
131            flags |= 0x2000
132        hd = struct.pack(
133            b"<IIIIIIII",
134            UF2_MAGIC_START0,
135            UF2_MAGIC_START1,
136            flags,
137            ptr + appstartaddr,
138            256,
139            blockno,
140            numblocks,
141            familyid,
142        )
143        while len(chunk) < 256:
144            chunk += b"\x00"
145        block = hd + chunk + datapadding + struct.pack(b"<I", UF2_MAGIC_END)
146        assert len(block) == 512
147        outp += block
148    return outp
149
150
151class Block:
152    def __init__(self, addr):
153        self.addr = addr
154        self.bytes = bytearray(256)
155
156    def encode(self, blockno, numblocks):
157        global familyid
158        flags = 0x0
159        if familyid:
160            flags |= 0x2000
161        hd = struct.pack(
162            "<IIIIIIII",
163            UF2_MAGIC_START0,
164            UF2_MAGIC_START1,
165            flags,
166            self.addr,
167            256,
168            blockno,
169            numblocks,
170            familyid,
171        )
172        hd += self.bytes[0:256]
173        while len(hd) < 512 - 4:
174            hd += b"\x00"
175        hd += struct.pack("<I", UF2_MAGIC_END)
176        return hd
177
178
179def convert_from_hex_to_uf2(buf):
180    global appstartaddr
181    appstartaddr = None
182    upper = 0
183    currblock = None
184    blocks = []
185    for line in buf.split("\n"):
186        if line[0] != ":":
187            continue
188        i = 1
189        rec = []
190        while i < len(line) - 1:
191            rec.append(int(line[i : i + 2], 16))
192            i += 2
193        tp = rec[3]
194        if tp == 4:
195            upper = ((rec[4] << 8) | rec[5]) << 16
196        elif tp == 2:
197            upper = ((rec[4] << 8) | rec[5]) << 4
198            assert (upper & 0xFFFF) == 0
199        elif tp == 1:
200            break
201        elif tp == 0:
202            addr = upper | (rec[1] << 8) | rec[2]
203            if appstartaddr == None:
204                appstartaddr = addr
205            i = 4
206            while i < len(rec) - 1:
207                if not currblock or currblock.addr & ~0xFF != addr & ~0xFF:
208                    currblock = Block(addr & ~0xFF)
209                    blocks.append(currblock)
210                currblock.bytes[addr & 0xFF] = rec[i]
211                addr += 1
212                i += 1
213    numblocks = len(blocks)
214    resfile = b""
215    for i in range(0, numblocks):
216        resfile += blocks[i].encode(i, numblocks)
217    return resfile
218
219
220def get_drives():
221    drives = []
222    if sys.platform == "win32":
223        r = subprocess.check_output(
224            [
225                "wmic",
226                "PATH",
227                "Win32_LogicalDisk",
228                "get",
229                "DeviceID,",
230                "VolumeName,",
231                "FileSystem,",
232                "DriveType",
233            ]
234        )
235        for line in r.split("\n"):
236            words = re.split("\s+", line)
237            if len(words) >= 3 and words[1] == "2" and words[2] == "FAT":
238                drives.append(words[0])
239    else:
240        rootpath = "/media"
241        if sys.platform == "darwin":
242            rootpath = "/Volumes"
243        elif sys.platform == "linux":
244            tmp = rootpath + "/" + os.environ["USER"]
245            if os.path.isdir(tmp):
246                rootpath = tmp
247        for d in os.listdir(rootpath):
248            drives.append(os.path.join(rootpath, d))
249
250    def has_info(d):
251        try:
252            return os.path.isfile(d + INFO_FILE)
253        except:
254            return False
255
256    return list(filter(has_info, drives))
257
258
259def board_id(path):
260    with open(path + INFO_FILE, mode="r") as file:
261        file_content = file.read()
262    return re.search("Board-ID: ([^\r\n]*)", file_content).group(1)
263
264
265def list_drives():
266    for d in get_drives():
267        print(d, board_id(d))
268
269
270def write_file(name, buf):
271    with open(name, "wb") as f:
272        f.write(buf)
273    print("Wrote %d bytes to %s." % (len(buf), name))
274
275
276def main():
277    global appstartaddr, familyid
278
279    def error(msg):
280        print(msg)
281        sys.exit(1)
282
283    parser = argparse.ArgumentParser(description="Convert to UF2 or flash directly.")
284    parser.add_argument(
285        "input", metavar="INPUT", type=str, nargs="?", help="input file (HEX, BIN or UF2)"
286    )
287    parser.add_argument(
288        "-b",
289        "--base",
290        dest="base",
291        type=str,
292        default="0x2000",
293        help="set base address of application for BIN format (default: 0x2000)",
294    )
295    parser.add_argument(
296        "-o",
297        "--output",
298        metavar="FILE",
299        dest="output",
300        type=str,
301        help='write output to named file; defaults to "flash.uf2" or "flash.bin" where sensible',
302    )
303    parser.add_argument("-d", "--device", dest="device_path", help="select a device path to flash")
304    parser.add_argument("-l", "--list", action="store_true", help="list connected devices")
305    parser.add_argument("-c", "--convert", action="store_true", help="do not flash, just convert")
306    parser.add_argument(
307        "-f",
308        "--family",
309        dest="family",
310        type=str,
311        default="0x0",
312        help="specify familyID - number or name (default: 0x0)",
313    )
314    parser.add_argument(
315        "-C", "--carray", action="store_true", help="convert binary file to a C array, not UF2"
316    )
317    args = parser.parse_args()
318    appstartaddr = int(args.base, 0)
319
320    if args.family.upper() in families:
321        familyid = families[args.family.upper()]
322    else:
323        try:
324            familyid = int(args.family, 0)
325        except ValueError:
326            error("Family ID needs to be a number or one of: " + ", ".join(families.keys()))
327
328    if args.list:
329        list_drives()
330    else:
331        if not args.input:
332            error("Need input file")
333        with open(args.input, mode="rb") as f:
334            inpbuf = f.read()
335        from_uf2 = is_uf2(inpbuf)
336        ext = "uf2"
337        if from_uf2:
338            outbuf = convert_from_uf2(inpbuf)
339            ext = "bin"
340        elif is_hex(inpbuf):
341            outbuf = convert_from_hex_to_uf2(inpbuf.decode("utf-8"))
342        elif args.carray:
343            outbuf = convert_to_carray(inpbuf)
344            ext = "h"
345        else:
346            outbuf = convert_to_uf2(inpbuf)
347        print(
348            "Converting to %s, output size: %d, start address: 0x%x"
349            % (ext, len(outbuf), appstartaddr)
350        )
351        if args.convert:
352            drives = []
353            if args.output == None:
354                args.output = "flash." + ext
355        else:
356            drives = get_drives()
357
358        if args.output:
359            write_file(args.output, outbuf)
360        else:
361            if len(drives) == 0:
362                error("No drive to deploy.")
363        for d in drives:
364            print("Flashing %s (%s)" % (d, board_id(d)))
365            write_file(d + "/NEW.UF2", outbuf)
366
367
368if __name__ == "__main__":
369    main()
370