lambdasoc/tools/flterm.py

   1 #!/usr/bin/env python3.5
   2
   3 # This file is Copyright (c) 2015-2020 M-Labs Limited.
   4 # License: BSD
   5
   6 import sys
   7 import os
   8 import time
   9 import asyncio
  10 import asyncserial
  11 import serial
  12 import argparse
  13
  14
  15 if sys.platform == "win32":
  16     import msvcrt
  17     import threading
  18
  19     def init_getkey(callback):
  20         loop = asyncio.get_event_loop()
  21         def getkey_thread():
  22             while True:
  23                 c = msvcrt.getch()
  24                 # HACK: This may still attempt to use the loop
  25                 # after it is closed - see comment below.
  26                 loop.call_soon_threadsafe(callback, c)
  27         threading.Thread(target=getkey_thread, daemon=True).start()
  28
  29     def deinit_getkey():
  30         # Python threads suck.
  31         pass
  32 else:
  33     import termios
  34
  35     def init_getkey(callback):
  36         global old_termios
  37
  38         fd = sys.stdin.fileno()
  39         old_termios = termios.tcgetattr(fd)
  40         new = old_termios.copy()
  41         new[3] = new[3] & ~termios.ICANON & ~termios.ECHO
  42         termios.tcsetattr(fd, termios.TCSANOW, new)
  43
  44         loop = asyncio.get_event_loop()
  45         def callback_wrapper():
  46             callback(os.read(sys.stdin.fileno(), 1))
  47         loop.add_reader(sys.stdin.fileno(), callback_wrapper)
  48
  49     def deinit_getkey():
  50         termios.tcsetattr(sys.stdin.fileno(), termios.TCSANOW, old_termios)
  51
  52
  53
  54 sfl_magic_req = b"sL5DdSMmkekro\n"
  55 sfl_magic_ack = b"z6IHG7cYDID6o\n"
  56
  57 # General commands
  58 sfl_cmd_abort = b"\x00"
  59 sfl_cmd_load  = b"\x01"
  60 sfl_cmd_jump  = b"\x02"
  61
  62 # Replies
  63 sfl_ack_success  = b"K"
  64 sfl_ack_crcerror = b"C"
  65 sfl_ack_unknown  = b"U"
  66 sfl_ack_error    = b"E"
  67
  68
  69 crc16_table = [
  70     0x0000, 0x1021, 0x2042, 0x3063, 0x4084, 0x50A5, 0x60C6, 0x70E7,
  71     0x8108, 0x9129, 0xA14A, 0xB16B, 0xC18C, 0xD1AD, 0xE1CE, 0xF1EF,
  72     0x1231, 0x0210, 0x3273, 0x2252, 0x52B5, 0x4294, 0x72F7, 0x62D6,
  73     0x9339, 0x8318, 0xB37B, 0xA35A, 0xD3BD, 0xC39C, 0xF3FF, 0xE3DE,
  74     0x2462, 0x3443, 0x0420, 0x1401, 0x64E6, 0x74C7, 0x44A4, 0x5485,
  75     0xA56A, 0xB54B, 0x8528, 0x9509, 0xE5EE, 0xF5CF, 0xC5AC, 0xD58D,
  76     0x3653, 0x2672, 0x1611, 0x0630, 0x76D7, 0x66F6, 0x5695, 0x46B4,
  77     0xB75B, 0xA77A, 0x9719, 0x8738, 0xF7DF, 0xE7FE, 0xD79D, 0xC7BC,
  78     0x48C4, 0x58E5, 0x6886, 0x78A7, 0x0840, 0x1861, 0x2802, 0x3823,
  79     0xC9CC, 0xD9ED, 0xE98E, 0xF9AF, 0x8948, 0x9969, 0xA90A, 0xB92B,
  80     0x5AF5, 0x4AD4, 0x7AB7, 0x6A96, 0x1A71, 0x0A50, 0x3A33, 0x2A12,
  81     0xDBFD, 0xCBDC, 0xFBBF, 0xEB9E, 0x9B79, 0x8B58, 0xBB3B, 0xAB1A,
  82     0x6CA6, 0x7C87, 0x4CE4, 0x5CC5, 0x2C22, 0x3C03, 0x0C60, 0x1C41,
  83     0xEDAE, 0xFD8F, 0xCDEC, 0xDDCD, 0xAD2A, 0xBD0B, 0x8D68, 0x9D49,
  84     0x7E97, 0x6EB6, 0x5ED5, 0x4EF4, 0x3E13, 0x2E32, 0x1E51, 0x0E70,
  85     0xFF9F, 0xEFBE, 0xDFDD, 0xCFFC, 0xBF1B, 0xAF3A, 0x9F59, 0x8F78,
  86     0x9188, 0x81A9, 0xB1CA, 0xA1EB, 0xD10C, 0xC12D, 0xF14E, 0xE16F,
  87     0x1080, 0x00A1, 0x30C2, 0x20E3, 0x5004, 0x4025, 0x7046, 0x6067,
  88     0x83B9, 0x9398, 0xA3FB, 0xB3DA, 0xC33D, 0xD31C, 0xE37F, 0xF35E,
  89     0x02B1, 0x1290, 0x22F3, 0x32D2, 0x4235, 0x5214, 0x6277, 0x7256,
  90     0xB5EA, 0xA5CB, 0x95A8, 0x8589, 0xF56E, 0xE54F, 0xD52C, 0xC50D,
  91     0x34E2, 0x24C3, 0x14A0, 0x0481, 0x7466, 0x6447, 0x5424, 0x4405,
  92     0xA7DB, 0xB7FA, 0x8799, 0x97B8, 0xE75F, 0xF77E, 0xC71D, 0xD73C,
  93     0x26D3, 0x36F2, 0x0691, 0x16B0, 0x6657, 0x7676, 0x4615, 0x5634,
  94     0xD94C, 0xC96D, 0xF90E, 0xE92F, 0x99C8, 0x89E9, 0xB98A, 0xA9AB,
  95     0x5844, 0x4865, 0x7806, 0x6827, 0x18C0, 0x08E1, 0x3882, 0x28A3,
  96     0xCB7D, 0xDB5C, 0xEB3F, 0xFB1E, 0x8BF9, 0x9BD8, 0xABBB, 0xBB9A,
  97     0x4A75, 0x5A54, 0x6A37, 0x7A16, 0x0AF1, 0x1AD0, 0x2AB3, 0x3A92,
  98     0xFD2E, 0xED0F, 0xDD6C, 0xCD4D, 0xBDAA, 0xAD8B, 0x9DE8, 0x8DC9,
  99     0x7C26, 0x6C07, 0x5C64, 0x4C45, 0x3CA2, 0x2C83, 0x1CE0, 0x0CC1,
 100     0xEF1F, 0xFF3E, 0xCF5D, 0xDF7C, 0xAF9B, 0xBFBA, 0x8FD9, 0x9FF8,
 101     0x6E17, 0x7E36, 0x4E55, 0x5E74, 0x2E93, 0x3EB2, 0x0ED1, 0x1EF0
 102 ]
 103
 104
 105 def crc16(l):
 106     crc = 0
 107     for d in l:
 108         crc = crc16_table[((crc >> 8) ^ d) & 0xff] ^ (crc << 8)
 109     return crc & 0xffff
 110
 111
 112 class SFLFrame:
 113     def __init__(self):
 114         self.cmd = bytes()
 115         self.payload = bytes()
 116
 117     def compute_crc(self):
 118         return crc16(self.cmd + self.payload)
 119
 120     def encode(self):
 121         packet = bytes([len(self.payload)])
 122         packet += self.compute_crc().to_bytes(2, "big")
 123         packet += self.cmd
 124         packet += self.payload
 125         return packet
 126
 127
 128 class Flterm:
 129     def __init__(self, port, speed, kernel_image, kernel_address,
 130                  upload_only, output_only):
 131         self.kernel_image = kernel_image
 132         self.kernel_address = kernel_address
 133         self.upload_only = upload_only
 134         self.output_only = output_only
 135
 136         self.port = asyncserial.AsyncSerial(port, baudrate=speed)
 137         if serial.__version__[0] == "2":
 138             self.port.ser.setRTS(False)
 139         else:
 140             self.port.ser.rts = False
 141
 142     def init(self):
 143         if not (self.upload_only or self.output_only):
 144             self.keyqueue = asyncio.Queue(100)
 145             def getkey_callback(c):
 146                 self.keyqueue.put_nowait(c)
 147             init_getkey(getkey_callback)
 148
 149         if self.upload_only:
 150             self.main_task = asyncio.ensure_future(self.upload_only_coro())
 151         else:
 152             self.main_task = asyncio.ensure_future(self.main_coro())
 153
 154     async def send_frame(self, frame):
 155         while True:
 156             await self.port.write_exactly(frame.encode())
 157             reply = await self.port.read(1)
 158             if reply == sfl_ack_success:
 159                 return
 160             elif reply == sfl_ack_crcerror:
 161                 pass  # retry
 162             else:
 163                 print("[FLTERM] Got unknown reply '{}' from the device, aborting.".format(reply))
 164                 raise ValueError
 165
 166     async def upload(self, filename, address):
 167         with open(filename, "rb") as f:
 168             data = f.read()
 169         print("[FLTERM] Uploading {} ({} bytes)...".format(filename, len(data)))
 170         current_address = address
 171         position = 0
 172         length = len(data)
 173         start = time.time()
 174         while len(data):
 175             sys.stdout.write("|{}>{}| {}%\r".format('=' * (20*position//length),
 176                                                     ' ' * (20-20*position//length),
 177                                                     100*position//length))
 178             sys.stdout.flush()
 179             frame = SFLFrame()
 180             frame_data = data[:251]
 181             frame.cmd = sfl_cmd_load
 182             frame.payload = current_address.to_bytes(4, "big")
 183             frame.payload += frame_data
 184             try:
 185                 await self.send_frame(frame)
 186             except ValueError:
 187                 return
 188             current_address += len(frame_data)
 189             position += len(frame_data)
 190             try:
 191                 data = data[251:]
 192             except:
 193                 data = []
 194         end = time.time()
 195         elapsed = end - start
 196         print("[FLTERM] Upload complete ({0:.1f}KB/s).".format(length/(elapsed*1024)))
 197         return length
 198
 199     async def boot(self):
 200         print("[FLTERM] Booting the device.")
 201         frame = SFLFrame()
 202         frame.cmd = sfl_cmd_jump
 203         frame.payload = self.kernel_address.to_bytes(4, "big")
 204         await self.send_frame(frame)
 205
 206     async def answer_magic(self):
 207         print("[FLTERM] Received firmware download request from the device.")
 208         await self.port.write_exactly(sfl_magic_ack)
 209         try:
 210             await self.upload(self.kernel_image, self.kernel_address)
 211         except FileNotFoundError:
 212             print("[FLTERM] File not found")
 213         else:
 214             await self.boot()
 215         print("[FLTERM] Done.");
 216
 217     async def main_coro(self):
 218         magic_detect_buffer = b"\x00"*len(sfl_magic_req)
 219         port_reader = None
 220         key_getter = None
 221         while True:
 222             if port_reader is None:
 223                 port_reader = asyncio.ensure_future(self.port.read(1024))
 224             fs = [port_reader]
 225             if not self.output_only:
 226                 if key_getter is None:
 227                     key_getter = asyncio.ensure_future(self.keyqueue.get())
 228                 fs += [key_getter]
 229             try:
 230                 done, pending = await asyncio.wait(
 231                     fs, return_when=asyncio.FIRST_COMPLETED)
 232             except asyncio.CancelledError:
 233                 for f in fs:
 234                     f.cancel()
 235                     try:
 236                         await f
 237                     except asyncio.CancelledError:
 238                         pass
 239                 raise
 240             if port_reader in done:
 241                 data = port_reader.result()
 242                 port_reader = None
 243                 sys.stdout.buffer.write(data)
 244                 sys.stdout.flush()
 245
 246                 if self.kernel_image is not None:
 247                     for c in data:
 248                         magic_detect_buffer = magic_detect_buffer[1:] + bytes([c])
 249                         if magic_detect_buffer == sfl_magic_req:
 250                             await self.answer_magic()
 251                             break
 252
 253             if key_getter in done:
 254                 await self.port.write(key_getter.result())
 255                 key_getter = None
 256
 257     async def upload_only_coro(self):
 258         magic_detect_buffer = b"\x00"*len(sfl_magic_req)
 259         while True:
 260             data = await self.port.read(1024)
 261             sys.stdout.buffer.write(data)
 262             sys.stdout.flush()
 263
 264             for c in data:
 265                 magic_detect_buffer = magic_detect_buffer[1:] + bytes([c])
 266                 if magic_detect_buffer == sfl_magic_req:
 267                     await self.answer_magic()
 268                     return
 269
 270     async def close(self):
 271         if not (self.upload_only or self.output_only):
 272             deinit_getkey()
 273         self.main_task.cancel()
 274         try:
 275             await self.main_task
 276         except asyncio.CancelledError:
 277             pass
 278         finally:
 279             self.port.close()
 280
 281
 282 def _get_args():
 283     parser = argparse.ArgumentParser()
 284     parser.add_argument("port", help="serial port")
 285     parser.add_argument("--speed", default=115200, help="serial baudrate")
 286     parser.add_argument("--kernel", default=None, help="kernel image")
 287     parser.add_argument("--kernel-addr", type=lambda a: int(a, 0),
 288                         default=0x40000000, help="kernel address")
 289     parser.add_argument("--upload-only", default=False, action="store_true",
 290                         help="only upload kernel")
 291     parser.add_argument("--output-only", default=False, action="store_true",
 292                         help="do not receive keyboard input or require a pty")
 293     return parser.parse_args()
 294
 295
 296 def main():
 297     if os.name == "nt":
 298         loop = asyncio.ProactorEventLoop()
 299         asyncio.set_event_loop(loop)
 300     else:
 301         loop = asyncio.get_event_loop()
 302     args = _get_args()
 303     flterm = Flterm(args.port, args.speed, args.kernel, args.kernel_addr,
 304                     args.upload_only, args.output_only)
 305     try:
 306         flterm.init()
 307         loop.run_until_complete(flterm.main_task)
 308     except KeyboardInterrupt:
 309         pass
 310     finally:
 311         loop.run_until_complete(flterm.close())
 312         loop.close()
 313
 314
 315 if __name__ == "__main__":
 316     main()