#!/usr/bin/python3

#
# USE AT YOUR OWN RISK.
#
# Do whatever you like with this source code.
#
# The constants (and implicitly the protocol) were borrowed from https://github.com/katie-snow/Ultimarc-linux
#
# Simon Funk / 2021-06-04
#

"""Run this script to change the ID of your Ultrastik (2015 firmware).

This worked fine for an ultrastik purchased in May 2021.
"""

import usb1

UM_REQUEST_TYPE = 0x21
UM_REQUEST      = 9
UM_TIMEOUT      = 2000

USTIK_VENDOR             = 0xD209
USTIK_PRODUCT            = 0x0511
USTIK_VALUE              = 0x0200
USTIK_MESG_LENGTH        = 4
USTIK_INTERFACE          = 2
USTIK_CONFIG_BASE        = 0x51

old_id = 1  # Change this if you need to config an Ultrastik that already has a different ID.

context = usb1.USBContext()
device  = context.getByVendorIDAndProductID(USTIK_VENDOR, USTIK_PRODUCT + (old_id-1))

if device is None:
    print("Could not find Ultrastik.")
else:
    print("Found Ulstrastik: %s"%(device,))
    while True:
        while True:
            new_id = input("What ID would you like to set it to (1-4)? ")
            try:
                new_id = int(new_id)
                assert 1 <= new_id <= 4
                assert new_id != old_id
            except:
                print("That's not a valid choice.")
            else:
                break

        if input("Change ID from %d to %d [y/n]? "%(old_id, new_id)) in ['y', 'Y', 'yes']:
            break

    print("Changing ID from %d to %d..."%(old_id, new_id))

    handle = device.open()

    if handle is None:
        print("Unable to open device.  (You may need to run this as root.)")
    else:
        print("Opened device: %s"%(handle,))

        #|  controlWrite(self, request_type, request, value, index, data, timeout=0)
        #|      Synchronous control write.
        #|      request_type: request type bitmask (bmRequestType), see
        #|        constants TYPE_* and RECIPIENT_*.
        #|      request: request id (some values are standard).
        #|      value, index, data: meaning is request-dependent.
        #|      timeout: in milliseconds, how long to wait for device acknowledgement.
        #|        Set to 0 to disable.
        #|
        #|      To avoid memory copies, use an object implementing the writeable buffer
        #|      interface (ex: bytearray) for the "data" parameter.
        #|
        #|      Returns the number of bytes actually sent.
        #|
        data = bytearray([USTIK_CONFIG_BASE + (new_id - 1), 0, 0, 0])
        assert len(data) == USTIK_MESG_LENGTH

        handle.detachKernelDriver(USTIK_INTERFACE)  # Otherwise we can't claim it because the kernel owns it.

        with handle.claimInterface(USTIK_INTERFACE):
            rv = handle.controlWrite(UM_REQUEST_TYPE, UM_REQUEST, USTIK_VALUE, USTIK_INTERFACE, data, UM_TIMEOUT)
            if rv == 4:
                print("Success?")
            else:
                print("Hmm. Appears to have failed, but:")
            print("Uplug the UltraStik and plug it back in to see if it worked (run lsusb before and after).")

        handle.close()

    device.close()