flash

#!/usr/bin/python3
'''
---- UoS3 OBC Firmware Flash Script ----

This script will flash the given binary onto the TM4C.

Author: Duncan Hamill (dh2g16@soton.ac.uk/duncanrhamill@googlemail.com)
Version: 0.1
Date: 2020-11-05
Copyright (c) UoS3 2020
'''

import sys
import argparse
import subprocess
from pathlib import Path

class TermCols:
    '''Terminal colour codes'''
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def error(msg, code=1):
    '''Print error message and exit with an error code'''
    print(f'[{TermCols.FAIL}ERR{TermCols.ENDC}] {msg}')
    sys.exit(code)

def info(msg):
    '''Print an information message to the screen'''
    print(f'[---] {msg}')

def ok(msg):
    '''Print a success message to the screen'''
    print(f'[{TermCols.OKGREEN} OK{TermCols.ENDC}] {msg}')

def main():
    '''
    Run the main flash script.
    '''

    # Get the command line args
    args = parse_args()

    # Get the path and check it exists
    elf_path = Path(args.elf)
    if not elf_path.exists():
        error(f'{elf_path} does not exist.')

    # Check it's got an .elf extension
    if not elf_path.suffix == '.elf':
        error(f'Expected a .elf file, found {elf_path}')

    # Flash to the correct board using either GDB for the TOBC or lm4flash for 
    # the launchpad
    flash_ok = False
    if args.target == 'launchpad':
        flash_ok = flash_launchpad(elf_path)
    elif args.target == 'tobc':
        flash_ok = flash_tobc(elf_path, args.debug)

    if not flash_ok:
        error('Could not flash binary to target')
        return 1
    
    ok('Binary successfully flashed to target')
    return 0

def flash_launchpad(elf_path):
    '''
    Flash the launchpad using lm4flash with the given ELF file path.
    '''

    # Get the binary path by replacing .elf with .bin
    bin_path = elf_path.with_suffix('.bin')

    info('Converting ELF file to binary format')

    # Convert the file from elf to binary format using the right objcopy.
    objcopy = subprocess.run([
        'arm-none-eabi-objcopy',
        '-O',
        'binary',
        elf_path,
        bin_path
    ])

    if objcopy.returncode != 0:
        error('Could not convert from ELF to binary format')
        return False
    else:
        ok('ELF converted to binary format')

    info(f'Flashing {bin_path} using lm4flash')
    
    # Flash the binary
    lm4flash = subprocess.run([
        'lm4flash',
        bin_path
    ])

    return lm4flash.returncode == 0

def flash_tobc(elf_path, debug):
    '''
    Use gdb to flash the TOBC with the given elf file, which should be 
    connected to the host via the BlackMagic probe
    '''

    info('Launching arm-none-eabi-gdb and flashing')

    # Use gdb to flash the file, in debug mode if needed
    if debug:
        info('Starting debug mode')
        gdb = subprocess.run([
            'arm-none-eabi-gdb',
            '-nx',
            '-ex', 'target extended-remote /dev/serial/by-id/usb-Black_Sphere_Technologies_Black_Magic_Probe_79A779DC-if00',
            '-x', 'blackmagic_debug.gdb',
            elf_path
        ])
    else:
        gdb = subprocess.run([
            'arm-none-eabi-gdb',
            '-nx',
            '--batch',
            '-ex', 'target extended-remote /dev/serial/by-id/usb-Black_Sphere_Technologies_Black_Magic_Probe_79A779DC-if00',
            '-x', 'blackmagic.gdb',
            elf_path
        ])

    return gdb.returncode == 0

def parse_args():
    '''
    Parse command line arguments
    '''

    # Create argparser
    parser = argparse.ArgumentParser(description='Build the OBC firmware.')
    parser.add_argument(
        'target',
        help='The target board to flash to, either `launchpad` or `tobc`'        
    )
    parser.add_argument(
        'elf',
        help='The path to the ELF image to flash'
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        help='If flashing the TOBC start a GDB debugging session'
    )

    # Parse the arguments
    return parser.parse_args()

if __name__ == '__main__':
    main()