greaseweazle-firmware/scripts/greaseweazle/tools/util.py

# greaseweazle/tools/util.py
#
# Greaseweazle control script: Utility functions.
#
# Written & released by Keir Fraser <keir.xen@gmail.com>
#
# This is free and unencumbered software released into the public domain.
# See the file COPYING for more details, or visit <http://unlicense.org>.

import argparse, os, sys, serial, struct, time, re
import importlib
import serial.tools.list_ports

from greaseweazle import version
from greaseweazle import error
from greaseweazle import usb as USB


class CmdlineHelpFormatter(argparse.ArgumentDefaultsHelpFormatter):
    def _get_help_string(self, action):
        help = action.help
        if '%no_default' in help:
            return help.replace('%no_default', '')
        if ('%(default)' in help
            or action.default is None
            or action.default is False
            or action.default is argparse.SUPPRESS):
            return help
        return help + ' (default: %(default)s)'


class ArgumentParser(argparse.ArgumentParser):
    def __init__(self, formatter_class=CmdlineHelpFormatter, *args, **kwargs):
        return super().__init__(formatter_class=formatter_class,
                                *args, **kwargs)

def drive_letter(letter):
    types = {
        'A': (USB.BusType.IBMPC, 0),
        'B': (USB.BusType.IBMPC, 1),
        '0': (USB.BusType.Shugart, 0),
        '1': (USB.BusType.Shugart, 1),
        '2': (USB.BusType.Shugart, 2)
    }
    if not letter.upper() in types:
        raise argparse.ArgumentTypeError("invalid drive letter: '%s'" % letter)
    return types[letter.upper()]

def range_str(s, e):
    str = "%d" % s
    if s != e: str += "-%d" % e
    return str

class TrackSet:

    class TrackIter:
        def __init__(self, ts):
            l = []
            for c in range(ts.cyl[0], ts.cyl[1]+1):
                for s in range(ts.side[0], ts.side[1]+1):
                    pc = (c, c*2)[ts.double_step]
                    pc += ts.side_off[0] if s == ts.side[0] else ts.side_off[1]
                    l.append((pc, s, c))
            l.sort()
            self.l = iter(l)
        def __next__(self):
            self.physical_cyl, self.side, self.cyl = next(self.l)
            return self
    
    def __init__(self):
        self.cyl = (0,79)
        self.side = (0,1)
        self.double_step = False

    def __str__(self):
        a, b = self.cyl
        s = 'c=%d' % a
        if a != b: s += '-%d' % b
        if self.double_step: s += 'x2'
        a, c = self.side
        b, d = self.side_off
        s += ':s=%d' % a
        if b != 0: s += '[%s%d]' % ('+' if b >= 0 else '', b)
        if a != c:
            s += '-%d' % c
            if d != 0: s += '[%s%d]' % ('+' if d >= 0 else '', d)
        return s

    def __iter__(self):
        return self.TrackIter(self)

def trackset(tracks):
    ts = TrackSet()
    for x in tracks.split(':'):
        k,v = x.split('=')
        if k == 'c':
            m = re.match('(\d+)(-(\d+))?(x2)?$', v)
            if m is None: raise ValueError()
            if m.group(3) is None:
                ts.cyl = int(m.group(1)), int(m.group(1))
            else:
                ts.cyl = int(m.group(1)), int(m.group(3))
            if m.group(4) is not None:
                ts.double_step = True
        elif k == 's':
            m = re.match('(?P<s>\d+)(\[(?P<s_off>[-+]\d)\])?'
                         '(-(?P<e>\d+)(\[(?P<e_off>[-+]\d)\])?)?$', v)
            if m is None: raise ValueError()
            s, e = m.group('s'), m.group('e')
            ts.side = int(s), int(e if e is not None else s)
            s, e = m.group('s_off'), m.group('e_off')
            s = int(s) if s is not None else 0
            e = int(e) if e is not None else 0
            ts.side_off = s,e
    return ts


def split_opts(seq):
    """Splits a name from its list of options."""
    parts = seq.split('::')
    name, opts = parts[0], dict()
    for x in map(lambda x: x.split(':'), parts[1:]):
        for y in x:
            try:
                opt, val = y.split('=')
            except ValueError:
                opt, val = y, True
            if opt:
                opts[opt] = val
    return name, opts


def get_image_class(name):
    image_types = { '.adf': 'ADF',
                    '.scp': 'SCP',
                    '.hfe': 'HFE',
                    '.ipf': 'IPF',
                    '.raw': 'KryoFlux' }
    if os.path.isdir(name):
        typename = 'KryoFlux'
    else:
        _, ext = os.path.splitext(name)
        error.check(ext.lower() in image_types,
                    "%s: Unrecognised file suffix '%s'" % (name, ext))
        typename = image_types[ext.lower()]
    mod = importlib.import_module('greaseweazle.image.' + typename.lower())
    return mod.__dict__[typename]


def with_drive_selected(fn, usb, args, *_args, **_kwargs):
    usb.set_bus_type(args.drive[0])
    try:
        usb.drive_select(args.drive[1])
        usb.drive_motor(args.drive[1], _kwargs.pop('motor', True))
        fn(usb, args, *_args, **_kwargs)
    except KeyboardInterrupt:
        print()
        usb.reset()
        usb.ser.close()
        usb.ser.open()
        raise
    finally:
        usb.drive_motor(args.drive[1], False)
        usb.drive_deselect()


def valid_ser_id(ser_id):
    return ser_id and ser_id.upper().startswith("GW")

def score_port(x, old_port=None):
    score = 0
    if x.manufacturer == "Keir Fraser" and x.product == "Greaseweazle":
        score = 20
    elif x.vid == 0x1209 and x.pid == 0x4d69:
        # Our very own properly-assigned PID. Guaranteed to be us.
        score = 20
    elif x.vid == 0x1209 and x.pid == 0x0001:
        # Our old shared Test PID. It's not guaranteed to be us.
        score = 10
    if score > 0 and valid_ser_id(x.serial_number):
        # A valid serial id is a good sign unless this is a reopen, and
        # the serials don't match!
        if not old_port or not valid_ser_id(old_port.serial_number):
            score = 20
        elif x.serial_number == old_port.serial_number:
            score = 30
        else:
            score = 0
    if old_port and old_port.location:
        # If this is a reopen, location field must match. A match is not
        # sufficient in itself however, as Windows may supply the same
        # location for multiple USB ports (this may be an interaction with
        # BitDefender). Hence we do not increase the port's score here.
        if not x.location or x.location != old_port.location:
            score = 0
    return score

def find_port(old_port=None):
    best_score, best_port = 0, None
    for x in serial.tools.list_ports.comports():
        score = score_port(x, old_port)
        if score > best_score:
            best_score, best_port = score, x
    if best_port:
        return best_port.device
    raise serial.SerialException('Cannot find the Greaseweazle device')

def port_info(devname):
    for x in serial.tools.list_ports.comports():
        if x.device == devname:
            return x
    return None

def usb_reopen(usb, is_update):
    mode = { False: 1, True: 0 }
    try:
        usb.switch_fw_mode(mode[is_update])
    except (serial.SerialException, struct.error):
        # Mac and Linux raise SerialException ("... returned no data")
        # Win10 pyserial returns a short read which fails struct.unpack
        pass
    usb.ser.close()
    for i in range(10):
        time.sleep(0.5)
        try:
            devicename = find_port(usb.port_info)
            new_ser = serial.Serial(devicename)
        except serial.SerialException:
            # Device not found
            pass
        else:
            new_usb = USB.Unit(new_ser)
            new_usb.port_info = port_info(devicename)
            return new_usb
    raise serial.SerialException('Could not reopen port after mode switch')


def usb_open(devicename, is_update=False, mode_check=True):

    if devicename is None:
        devicename = find_port()
    
    usb = USB.Unit(serial.Serial(devicename))
    usb.port_info = port_info(devicename)

    if not mode_check:
        return usb

    if usb.update_mode and not is_update:
        if usb.hw_model == 7 and not usb.update_jumpered:
            usb = usb_reopen(usb, is_update)
            if not usb.update_mode:
                return usb
        print("Greaseweazle is in Firmware Update Mode:")
        print(" The only available action is \"update\" of main firmware")
        if usb.update_jumpered:
            print(" Remove the Update Jumper for normal operation")
        else:
            print(" Main firmware is erased: You *must* perform an update!")
        sys.exit(1)

    if is_update and not usb.update_mode:
        if usb.hw_model == 7:
            usb = usb_reopen(usb, is_update)
            error.check(usb.update_mode, """\
Greaseweazle F7 did not change to Firmware Update Mode as requested.
If the problem persists, install the Update Jumper (across RX/TX).""")
            return usb
        print("Greaseweazle is in Normal Mode:")
        print(" To \"update\" you must install the Update Jumper")
        sys.exit(1)

    if not usb.update_mode and usb.update_needed:
        print("Firmware is out of date: Require v%u.%u"
              % (version.major, version.minor))
        if usb.hw_model == 7:
            print("Run \"update <update_file>\"")
        else:
            print("Install the Update Jumper and \"update <update_file>\"")
        sys.exit(1)

    return usb
    


# Local variables:
# python-indent: 4
# End: