#!/usr/bin/python
#
# Copyright (c) 2016, NORDUnet A/S All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
# - Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# - Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
# - Neither the name of the NORDUnet nor the names of its contributors may
#   be used to endorse or promote products derived from this software
#   without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
# TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
Utility to transfer a new firmware to the bootloader.

Example usage:

   ./bin/dfu path/to/firmware.bin

Then reset the Cryptech Alpha board.
"""
import os
import sys
import time
import struct
import serial
import argparse

from binascii import crc32

CHUNK_SIZE = 4096


def parse_args():
    """
    Parse the command line arguments
    """
    parser = argparse.ArgumentParser(description = "Device firmware upgrader",
                                     add_help = True,
                                     formatter_class = argparse.ArgumentDefaultsHelpFormatter,
                                     )

    parser.add_argument('--verbose',
                        dest='verbose',
                        action='store_true', default=False,
                        help='Verbose operation',
                        )

    parser.add_argument('--device',
                        dest='device',
                        default='/dev/ttyUSB0',
                        help='Name of management port USB serial device',
                        )

    # positional argument(s)
    parser.add_argument('filename')

    return parser.parse_args()


def _write(dst, data):
    for i in range(len(data)):
        dst.write(data[i])
        time.sleep(0.1)
    if len(data) == 4:
        print("Wrote 0x{:02x}{:02x}{:02x}{:02x}".format(ord(data[0]), ord(data[1]), ord(data[2]), ord(data[3])))
    else:
        print("Wrote {!r}".format(data))


def _read(dst, verbose=False):
    res = ''
    while True:
        x = dst.read(1)
        if not x:
            break
        res += x
    if res and verbose:
        print ("Read {!r}".format(res))
    return res


def send_file(filename, args):
    s = os.stat(filename)
    size = s.st_size
    src = open(filename, 'rb')
    chunk_size = CHUNK_SIZE

    print("Waiting for contact with the bootloader on device {!s}, reset the CrypTech Alpha...".format(args.device))

    while True:
        try:
            dst = serial.Serial(args.device, 921600, timeout=0.1)
        except serial.SerialException:
            time.sleep(0.2)
            continue
        dst.write('\r')
        response = _read(dst, args.verbose)
        if 'OK' in response:
            dst.timeout=2
            break

    print('\nUploading firmware\n')

    crc = 0
    counter = 0
    # 1. Write size of file (4 bytes)
    _write(dst, struct.pack('<I', size))
    _read(dst, args.verbose)
    # 2. Write file contents while calculating CRC-32
    while True:
        data = src.read(chunk_size)
        if not data:
            break
        dst.write(data)
        print("Wrote {!s} bytes (chunk {!s}/{!s})".format(len(data), counter, int(size / chunk_size)))
        # read ACK (a counter of number of 4k chunks received)
        while True:
            ack_bytes = dst.read(4)
            if len(ack_bytes) == 4:
                break
            print('ERROR: Did not receive an ACK, got {!r}'.format(ack_bytes))
            dst.write('\r')  # eventually get back to the CLI prompt
        ack = struct.unpack('<I', ack_bytes)[0]
        if ack != counter + 1:
            print('ERROR: Did not receive the expected counter as ACK (got {!r}/{!r}, not {!r})'.format(ack, ack_bytes, counter))
            flush = dst.read(100)
            print('FLUSH data: {!r}'.format(flush))
            return False
        counter += 1

        crc = crc32(data, crc) & 0xffffffff

    _read(dst, args.verbose)

    # 3. Write CRC-32 (4 bytes)
    _write(dst, struct.pack('<I', crc))

    if args.verbose:
        print('\nFinished - this command might hang now depending on the firmware loaded')
        _read(dst, args.verbose)
    else:
        print('\nFinished uploading firmware')

    src.close()
    dst.close()
    return True


def main(args):
    send_file(args.filename, args)
    return True

if __name__ == '__main__':
    try:
        args = parse_args()
        if main(args):
            sys.exit(0)
        sys.exit(1)
    except KeyboardInterrupt:
        pass