magboot.py

#! /usr/bin/env python
# Copyright (c) 2010-2011 Magnus Olsson (magnus@minimum.se)
# See LICENSE for details

"""magboot - AVR serial bootloader

usage: magboot <port> <device> <[command1] [command2] ... [commandN]>

commands:
-w <file>	Write contents of <file> to beginning of flash (use - for stdin)
-i		Verify device signature
-r		Reset device (will bypass bootloader on next boot)
-z		Wait for device to appear on <port>

supported devices:
atmega328p
atmega168p
"""

import serial
import getopt
import sys
from struct import *
from array import array

global ser
global devicelist
global dev

atmega328p = {
    'signature': "\x1E\x95\x0F",
    'pagesize': 128,
}

atmega168p = {
    'signature': "\x1E\x94\x0B",
    'pagesize': 128,
}

devicelist = {
    'atmega328p': atmega328p,
    'atmega168p': atmega168p,
}


def usage(*args):
    sys.stdout = sys.stderr
    print __doc__
    for msg in args:
        print msg
    sys.exit(2)


def do_cmd(str, expect_reply=True):
    ser.write(str)
    if not expect_reply:
        print "> OK"
        return

    reply = ser.read()
    if (len(reply) == 0):
        print "> FAILED (TIMEOUT)"
        sys.exit(1)

    if (reply == 'Y'):
        print "> OK"
        return

    if (reply == 'N'):
        print "> FAILED"
    else:
        print "> FAILED (UNKNOWN):"
        while (len(reply) == 1):
            print reply
            reply = ser.read()
    sys.exit(1)


def cmd_device_id():
    print "ID"
    do_cmd('I' + dev['signature'])


def cmd_load_addr(addr):
    print "LOAD_ADDR"
    # Little-endian, 16-bit uint load address
    load_addr = pack('<H', addr)
    do_cmd('A' + load_addr)


def checksum(data):
    csum = 0
    words = array('H', data)
    for w in words:
        csum = csum + w

    # Fold overflow bits
    while (csum >> 16):
        csum = (csum & 0xFFFF) + (csum >> 16)

    return pack('<H', csum)


def cmd_write_file(fname):
    print "WRITE_FILE"

    if (fname == "-"):
        f = sys.stdin
    else:
        f = open(fname, "rb")

    eof = False

    cmd_load_addr(0)

    while (not eof):
        buf = array('c')
        bytecount = 0
        while (bytecount < dev['pagesize']):
            data = f.read(1)
            if (len(data) == 0):
                eof = True
                break
            buf.append(data)
            bytecount = bytecount + 1

        # Zerofill remaining part of page
        if (bytecount != dev['pagesize']):
            while (bytecount < dev['pagesize']):
                buf.append(chr(0))
                bytecount = bytecount + 1

        bufstr = buf.tostring()
        csum = checksum(bufstr)
        do_cmd('W' + csum + bufstr)


def cmd_reset():
    print "RESET"
    do_cmd('R', False)


def cmd_wait():
    print "WAIT"
    while True:
        sys.stdout.flush()
        ser.write('I' + dev['signature'])
        data = ser.read()
        if (len(data) == 1 and data == 'Y'):
            break
        sys.stdout.write('.')
    sys.stdout.write('\n')


if __name__ == "__main__":
    if len(sys.argv) < 4:
        usage()

    try:
        dev = devicelist[sys.argv[2]]
    except KeyError:
        print "Unsupported device '" + sys.argv[2] + "'"
        sys.exit(3)

    ser = serial.Serial(port=sys.argv[1], baudrate=19200, timeout=1)

    try:
        opts, args = getopt.getopt(sys.argv[3:], 'a:w:ijrz')
    except getopt.error, msg:
        usage(msg)

    ser.flushInput()
    ser.flushOutput()

    for o, a in opts:
        if o == '-w':
            cmd_write_file(a)
        if o == '-i':
            cmd_device_id()
        if o == '-r':
            cmd_reset()
        if o == '-z':
            cmd_wait()

    ser.close()