qmk_firmware

QMK firmware for my keyboards (Corne, Sweep Ferris) and trackball (Ploopy Adept)
Log | Files | Refs | Submodules | LICENSE

uf2conv.py (13112B)


      1 #!/usr/bin/env python3
      2 # yapf: disable
      3 import sys
      4 import struct
      5 import subprocess
      6 import re
      7 import os
      8 import os.path
      9 import argparse
     10 import json
     11 from time import sleep
     12 
     13 
     14 UF2_MAGIC_START0 = 0x0A324655 # "UF2\n"
     15 UF2_MAGIC_START1 = 0x9E5D5157 # Randomly selected
     16 UF2_MAGIC_END    = 0x0AB16F30 # Ditto
     17 
     18 INFO_FILE = "/INFO_UF2.TXT"
     19 
     20 appstartaddr = 0x2000
     21 familyid = 0x0
     22 
     23 
     24 def is_uf2(buf):
     25     w = struct.unpack("<II", buf[0:8])
     26     return w[0] == UF2_MAGIC_START0 and w[1] == UF2_MAGIC_START1
     27 
     28 def is_hex(buf):
     29     try:
     30         w = buf[0:30].decode("utf-8")
     31     except UnicodeDecodeError:
     32         return False
     33     if w[0] == ':' and re.match(rb"^[:0-9a-fA-F\r\n]+$", buf):
     34         return True
     35     return False
     36 
     37 def convert_from_uf2(buf):
     38     global appstartaddr
     39     global familyid
     40     numblocks = len(buf) // 512
     41     curraddr = None
     42     currfamilyid = None
     43     families_found = {}
     44     prev_flag = None
     45     all_flags_same = True
     46     outp = []
     47     for blockno in range(numblocks):
     48         ptr = blockno * 512
     49         block = buf[ptr:ptr + 512]
     50         hd = struct.unpack(b"<IIIIIIII", block[0:32])
     51         if hd[0] != UF2_MAGIC_START0 or hd[1] != UF2_MAGIC_START1:
     52             print("Skipping block at " + ptr + "; bad magic")
     53             continue
     54         if hd[2] & 1:
     55             # NO-flash flag set; skip block
     56             continue
     57         datalen = hd[4]
     58         if datalen > 476:
     59             assert False, "Invalid UF2 data size at " + ptr
     60         newaddr = hd[3]
     61         if (hd[2] & 0x2000) and (currfamilyid == None):
     62             currfamilyid = hd[7]
     63         if curraddr == None or ((hd[2] & 0x2000) and hd[7] != currfamilyid):
     64             currfamilyid = hd[7]
     65             curraddr = newaddr
     66             if familyid == 0x0 or familyid == hd[7]:
     67                 appstartaddr = newaddr
     68         padding = newaddr - curraddr
     69         if padding < 0:
     70             assert False, "Block out of order at " + ptr
     71         if padding > 10*1024*1024:
     72             assert False, "More than 10M of padding needed at " + ptr
     73         if padding % 4 != 0:
     74             assert False, "Non-word padding size at " + ptr
     75         while padding > 0:
     76             padding -= 4
     77             outp.append(b"\x00\x00\x00\x00")
     78         if familyid == 0x0 or ((hd[2] & 0x2000) and familyid == hd[7]):
     79             outp.append(block[32 : 32 + datalen])
     80         curraddr = newaddr + datalen
     81         if hd[2] & 0x2000:
     82             if hd[7] in families_found.keys():
     83                 if families_found[hd[7]] > newaddr:
     84                     families_found[hd[7]] = newaddr
     85             else:
     86                 families_found[hd[7]] = newaddr
     87         if prev_flag == None:
     88             prev_flag = hd[2]
     89         if prev_flag != hd[2]:
     90             all_flags_same = False
     91         if blockno == (numblocks - 1):
     92             print("--- UF2 File Header Info ---")
     93             families = load_families()
     94             for family_hex in families_found.keys():
     95                 family_short_name = ""
     96                 for name, value in families.items():
     97                     if value == family_hex:
     98                         family_short_name = name
     99                 print("Family ID is {:s}, hex value is 0x{:08x}".format(family_short_name,family_hex))
    100                 print("Target Address is 0x{:08x}".format(families_found[family_hex]))
    101             if all_flags_same:
    102                 print("All block flag values consistent, 0x{:04x}".format(hd[2]))
    103             else:
    104                 print("Flags were not all the same")
    105             print("----------------------------")
    106             if len(families_found) > 1 and familyid == 0x0:
    107                 outp = []
    108                 appstartaddr = 0x0
    109     return b"".join(outp)
    110 
    111 def convert_to_carray(file_content):
    112     outp = "const unsigned long bindata_len = %d;\n" % len(file_content)
    113     outp += "const unsigned char bindata[] __attribute__((aligned(16))) = {"
    114     for i in range(len(file_content)):
    115         if i % 16 == 0:
    116             outp += "\n"
    117         outp += "0x%02x, " % file_content[i]
    118     outp += "\n};\n"
    119     return bytes(outp, "utf-8")
    120 
    121 def convert_to_uf2(file_content):
    122     global familyid
    123     datapadding = b""
    124     while len(datapadding) < 512 - 256 - 32 - 4:
    125         datapadding += b"\x00\x00\x00\x00"
    126     numblocks = (len(file_content) + 255) // 256
    127     outp = []
    128     for blockno in range(numblocks):
    129         ptr = 256 * blockno
    130         chunk = file_content[ptr:ptr + 256]
    131         flags = 0x0
    132         if familyid:
    133             flags |= 0x2000
    134         hd = struct.pack(b"<IIIIIIII",
    135             UF2_MAGIC_START0, UF2_MAGIC_START1,
    136             flags, ptr + appstartaddr, 256, blockno, numblocks, familyid)
    137         while len(chunk) < 256:
    138             chunk += b"\x00"
    139         block = hd + chunk + datapadding + struct.pack(b"<I", UF2_MAGIC_END)
    140         assert len(block) == 512
    141         outp.append(block)
    142     return b"".join(outp)
    143 
    144 class Block:
    145     def __init__(self, addr, default_data=0xFF):
    146         self.addr = addr
    147         self.bytes = bytearray([default_data] * 256)
    148 
    149     def encode(self, blockno, numblocks):
    150         global familyid
    151         flags = 0x0
    152         if familyid:
    153             flags |= 0x2000
    154         if devicetype:
    155             flags |= 0x8000
    156         hd = struct.pack("<IIIIIIII",
    157             UF2_MAGIC_START0, UF2_MAGIC_START1,
    158             flags, self.addr, 256, blockno, numblocks, familyid)
    159         hd += self.bytes[0:256]
    160         if devicetype:
    161             hd += bytearray(b'\x08\x29\xa7\xc8')
    162             hd += bytearray(devicetype.to_bytes(4, 'little'))
    163         while len(hd) < 512 - 4:
    164             hd += b"\x00"
    165         hd += struct.pack("<I", UF2_MAGIC_END)
    166         return hd
    167 
    168 def convert_from_hex_to_uf2(buf):
    169     global appstartaddr
    170     appstartaddr = None
    171     upper = 0
    172     currblock = None
    173     blocks = []
    174     for line in buf.split('\n'):
    175         if line[0] != ":":
    176             continue
    177         i = 1
    178         rec = []
    179         while i < len(line) - 1:
    180             rec.append(int(line[i:i+2], 16))
    181             i += 2
    182         tp = rec[3]
    183         if tp == 4:
    184             upper = ((rec[4] << 8) | rec[5]) << 16
    185         elif tp == 2:
    186             upper = ((rec[4] << 8) | rec[5]) << 4
    187         elif tp == 1:
    188             break
    189         elif tp == 0:
    190             addr = upper + ((rec[1] << 8) | rec[2])
    191             if appstartaddr == None:
    192                 appstartaddr = addr
    193             i = 4
    194             while i < len(rec) - 1:
    195                 if not currblock or currblock.addr & ~0xff != addr & ~0xff:
    196                     currblock = Block(addr & ~0xff)
    197                     blocks.append(currblock)
    198                 currblock.bytes[addr & 0xff] = rec[i]
    199                 addr += 1
    200                 i += 1
    201     numblocks = len(blocks)
    202     resfile = b""
    203     for i in range(0, numblocks):
    204         resfile += blocks[i].encode(i, numblocks)
    205     return resfile
    206 
    207 def to_str(b):
    208     return b.decode("utf-8")
    209 
    210 def get_drives():
    211     drives = []
    212     if sys.platform == "win32":
    213         r = subprocess.check_output([
    214             "powershell",
    215             "-Command",
    216             '(Get-WmiObject Win32_LogicalDisk -Filter "FileSystem=\'FAT\'").DeviceID'
    217             ])
    218         drives = [drive.strip() for drive in to_str(r).splitlines()]
    219     else:
    220         searchpaths = ["/mnt", "/media"]
    221         if sys.platform == "darwin":
    222             searchpaths = ["/Volumes"]
    223         elif sys.platform == "linux":
    224             searchpaths += ["/media/" + os.environ["USER"], "/run/media/" + os.environ["USER"]]
    225             if "SUDO_USER" in os.environ.keys():
    226                 searchpaths += ["/media/" + os.environ["SUDO_USER"]]
    227                 searchpaths += ["/run/media/" + os.environ["SUDO_USER"]]
    228 
    229         for rootpath in searchpaths:
    230             if os.path.isdir(rootpath):
    231                 for d in os.listdir(rootpath):
    232                     if os.path.isdir(os.path.join(rootpath, d)):
    233                         drives.append(os.path.join(rootpath, d))
    234 
    235 
    236     def has_info(d):
    237         try:
    238             return os.path.isfile(d + INFO_FILE)
    239         except:
    240             return False
    241 
    242     return list(filter(has_info, drives))
    243 
    244 
    245 def board_id(path):
    246     with open(path + INFO_FILE, mode='r') as file:
    247         file_content = file.read()
    248     return re.search(r"Board-ID: ([^\r\n]*)", file_content).group(1)
    249 
    250 
    251 def list_drives():
    252     for d in get_drives():
    253         print(d, board_id(d))
    254 
    255 
    256 def write_file(name, buf):
    257     with open(name, "wb") as f:
    258         f.write(buf)
    259     print("Wrote %d bytes to %s" % (len(buf), name))
    260 
    261 
    262 def load_families():
    263     # The expectation is that the `uf2families.json` file is in the same
    264     # directory as this script. Make a path that works using `__file__`
    265     # which contains the full path to this script.
    266     filename = "uf2families.json"
    267     pathname = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
    268     with open(pathname) as f:
    269         raw_families = json.load(f)
    270 
    271     families = {}
    272     for family in raw_families:
    273         families[family["short_name"]] = int(family["id"], 0)
    274 
    275     return families
    276 
    277 
    278 def main():
    279     global appstartaddr, familyid
    280     def error(msg):
    281         print(msg, file=sys.stderr)
    282         sys.exit(1)
    283     parser = argparse.ArgumentParser(description='Convert to UF2 or flash directly.')
    284     parser.add_argument('input', metavar='INPUT', type=str, nargs='?',
    285                         help='input file (HEX, BIN or UF2)')
    286     parser.add_argument('-b', '--base', dest='base', type=str,
    287                         default="0x2000",
    288                         help='set base address of application for BIN format (default: 0x2000)')
    289     parser.add_argument('-f', '--family', dest='family', type=str,
    290                         default="0x0",
    291                         help='specify familyID - number or name (default: 0x0)')
    292     parser.add_argument('-t' , '--device-type', dest='devicetype', type=str,
    293                         help='specify deviceTypeID extension tag - number')
    294     parser.add_argument('-o', '--output', metavar="FILE", dest='output', type=str,
    295                         help='write output to named file; defaults to "flash.uf2" or "flash.bin" where sensible')
    296     parser.add_argument('-d', '--device', dest="device_path",
    297                         help='select a device path to flash')
    298     parser.add_argument('-l', '--list', action='store_true',
    299                         help='list connected devices')
    300     parser.add_argument('-c', '--convert', action='store_true',
    301                         help='do not flash, just convert')
    302     parser.add_argument('-D', '--deploy', action='store_true',
    303                         help='just flash, do not convert')
    304     parser.add_argument('-w', '--wait', action='store_true',
    305                         help='wait for device to flash')
    306     parser.add_argument('-C', '--carray', action='store_true',
    307                         help='convert binary file to a C array, not UF2')
    308     parser.add_argument('-i', '--info', action='store_true',
    309                         help='display header information from UF2, do not convert')
    310     args = parser.parse_args()
    311     appstartaddr = int(args.base, 0)
    312 
    313     families = load_families()
    314 
    315     if args.family.upper() in families:
    316         familyid = families[args.family.upper()]
    317     else:
    318         try:
    319             familyid = int(args.family, 0)
    320         except ValueError:
    321             error("Family ID needs to be a number or one of: " + ", ".join(families.keys()))
    322 
    323     global devicetype
    324     devicetype = int(args.devicetype, 0) if args.devicetype else None
    325 
    326     if args.list:
    327         list_drives()
    328     else:
    329         if not args.input:
    330             error("Need input file")
    331         with open(args.input, mode='rb') as f:
    332             inpbuf = f.read()
    333         from_uf2 = is_uf2(inpbuf)
    334         ext = "uf2"
    335         if args.deploy:
    336             outbuf = inpbuf
    337         elif from_uf2 and not args.info:
    338             outbuf = convert_from_uf2(inpbuf)
    339             ext = "bin"
    340         elif from_uf2 and args.info:
    341             outbuf = ""
    342             convert_from_uf2(inpbuf)
    343         elif is_hex(inpbuf):
    344             outbuf = convert_from_hex_to_uf2(inpbuf.decode("utf-8"))
    345         elif args.carray:
    346             outbuf = convert_to_carray(inpbuf)
    347             ext = "h"
    348         else:
    349             outbuf = convert_to_uf2(inpbuf)
    350         if not args.deploy and not args.info:
    351             print("Converted to %s, output size: %d, start address: 0x%x" %
    352                   (ext, len(outbuf), appstartaddr))
    353         if args.convert or ext != "uf2":
    354             if args.output == None:
    355                 args.output = "flash." + ext
    356         if args.output:
    357             write_file(args.output, outbuf)
    358         if ext == "uf2" and not args.convert and not args.info:
    359             drives = get_drives()
    360             if len(drives) == 0:
    361                 if args.wait:
    362                     print("Waiting for drive to deploy...")
    363                     while len(drives) == 0:
    364                         sleep(0.1)
    365                         drives = get_drives()
    366                 elif not args.output:
    367                     error("No drive to deploy.")
    368             for d in drives:
    369                 print("Flashing %s (%s)" % (d, board_id(d)))
    370                 write_file(d + "/NEW.UF2", outbuf)
    371 
    372 
    373 if __name__ == "__main__":
    374     main()