#!/usr/bin/python

# Firmware Goblin
# version 0.1
#
# Copyright (C) 2006 by Daniel Lenski <lenski@umd.edu>
# Time-stamp: <2006-11-13 09:57:50 dlenski>
#
# Released under the terms of the
# GNU General Public License version 2 or later

helptext = \
'''This program makes it easier to extract the
firmware of an EZ-USB peripheral using the Windows driver.

How to use it:
  (1) Get the USB Snoop software from
      <http://benoit.papillault.free.fr/usbsnoop/>
  (2) Install USB Snoop and the Windows driver for the
      device.
  (3) Use USB Snoop to produce a log file of the EZ-USB
      peripheral being plugged into the bus.
  (4) Run this program on the log file:
      %prog usbsnoop.log -o usbsnoop.hex
  (5) Now usbsnoop.hex contains the device firmware
      in Intel HEX8 format
  (6) Reboot into Linux or BSD and write a driver :-)

The output of this program is compatible with the fxload
program used to load EZ-USB firmware under Linux.'''

version = '0.1'

headline = \
'''FirmwareGoblin, version %s
Copyright (C) 2006 by Dan Lenski.
This is free software; see the source code for copying conditions.
''' % version

hexheader = \
'''# Playback of EZ-USB firmware download
#   USB Snoop log file: %(logfile)s
#   Time stamp: %(stamp)s
# Produced by FirmwareGoblin (C) 2006 by Daniel Lenski
#'''

badinput = '''
ERROR: The file %(logfile)s does not appear to contain a
transcript of an EZ-USB firmware download.  Check that:
(a) %(logfile)s really is a log file from USB Snoop:
    http://benoit.papillault.free.fr/usbsnoop
(b) the device being studied contains an Anchor/Cypress EZ-USB chip'''

######################################################################

import re
from optparse import OptionParser
from array import array
from datetime import datetime
from sys import stdin, stderr

class IH8Chunk(object):
    '''IH8Chunk(buffer, address, rectype=0) -> Intel HEX8 file format line

    Holds a chunk or line in an Intel HEX8 file.  Its string representation
    is the matching HEX8 line.'''

    addr = 0
    buf = array('B')
    rectype = 0

    def __init__(c, buf, addr, rectype=0):
        c.addr = addr;
        c.buf = array('B', buf)
        c.rectype = rectype

    def __len__(c):
        return len(c.buf)
    
    def __str__(c):
        line = array('B', [len(c), (c.addr>>8) & 0xff, c.addr & 0xff, c.rectype])
        line.extend(c.buf)
        line.append(-sum(line) & 0xff)
        return ":" + ''.join(( "%02X"%b for b in line ))

def hexstring(s):
    return [int(b,16) for b in s.split()]

######################################################################

# greetings

print>>stderr, headline

# parse command line

parser = OptionParser(usage="usage: %%prog [options] [logfile]\n\n%s" % helptext)
parser.add_option('-o', '--output', dest='output', help='write output to FILE', metavar='FILE')
parser.add_option('-m', '--memory-size', dest='memsize', help='set the size of the EZ-USB device RAM',
          type='int', default=0x1b40)
parser.add_option("--no-header", action="store_false", dest="header",
                  help="don't include a header comment in the output", default=1)

options, args = parser.parse_args()

if len(args): stdin = open(args[0], 'r')

# regular expressions to match from the "USB Snoop" log

outgoing = re.compile(r'\[\d+ ms\]\s+>>>\s+URB (\d+)')              # start of an outgoing URB
vendor_block = re.compile(r'-- URB_FUNCTION_VENDOR_DEVICE:')        # start of a vendor-request control URB
key_value = re.compile(r'  (\w+)\s*=\s*([a-f0-9]*)(.*)', re.I)      # a key-value pair in the URB block
buf_line = re.compile(r'    ([a-f0-9]+):((?: [a-f0-9]{2})+)', re.I) # a line of transfer buffer contents

# play back "USB Snoop" log

firmware = []
outgoing_flag = False
vendor_flag = False
reset_state = False

print>>stderr, "Playing back USB Snoop log:"
for line in stdin:
    start_outgoing = outgoing.match(line)
    start_vendor_block = vendor_block.match(line)
    
    if start_outgoing:
        outgoing_flag = True
    elif start_vendor_block:
        vendor_flag = True
        urb_dict = {}
        urb_buf = array('B')
    elif outgoing_flag and vendor_flag:
        kv = key_value.match(line)
        bl = buf_line.match(line)
        
        if kv:
            # add a k-v pair to the dict
            k = kv.group(1)
            v = int(kv.group(2), 16)
            urb_dict[ k ] = v
        elif bl:
            # read a line of hex data
            addr = int(bl.group(1), 16)
            data = hexstring(bl.group(2))
            assert addr == len(urb_buf)
            urb_buf.extend(data)
        else:
            # done with block
            outgoing_flag = vendor_flag = False
            if urb_dict.get('Request') == 0xa0:
                assert urb_dict['TransferBufferLength'] == len(urb_buf)
                addr = urb_dict['Value']

                # device must be put into reset mode while downloading firmware
                if addr == 0x7f92:
                    reset_state = (urb_buf[0] & 1)
                    if reset_state: print>>stderr, "  Entering EZ-USB device reset."
                    else: print>>stderr, "  Exiting EZ-USB device reset."
                else:
                    if not reset_state:
                        print>>stderr, "  WARNING: firmware load outside of reset state."
                    if addr+len(urb_buf)>options.memsize:
                        print>>stderr, "  WARNING: firmware load beyond address 0x%x" % options.memsize
                    firmware.append( IH8Chunk(urb_buf, addr) )

# write the firmware download in Intel HEX8 file format

if firmware:

    print>>stderr, "Converting firmware download to Intel HEX format:"
    
    if options.output: stdout = open(options.output, 'w')
    if options.header: print hexheader % {'logfile': stdin.name, 'stamp' : datetime.today()}
    
    for chunk in firmware: print chunk
    print IH8Chunk([], 0, 1) # IH8 eof

    print>>stderr, "  Wrote %d bytes in %d chunks." % (sum((len(c) for c in firmware)), len(firmware))

else:

    print>>stderr, badinput % {'logfile': stdin.name}