264 lines
8.1 KiB
Python
264 lines
8.1 KiB
Python
|
|
from __future__ import annotations
|
|
|
|
import array
|
|
import os
|
|
import re
|
|
import struct
|
|
|
|
import abc
|
|
from typing import *
|
|
|
|
# TODO: EVERYTHING
|
|
# * implement connecting to a device
|
|
# * implement lower-level device commands
|
|
# * implement device commands
|
|
|
|
|
|
class DevConn:
|
|
pass
|
|
|
|
class UsbConn(DevConn):
|
|
_USB_DEFAULT_VID = 0xcafe
|
|
_USB_DEFAULT_PID = 0x1312
|
|
_SUBCLASS = 42
|
|
_PROTOCOL = 69
|
|
_VER_MIN = 0x0010
|
|
_VER_MAX = 0x00ff # ???
|
|
|
|
def _open_dev(dev) -> Union[UsbConn, str]:
|
|
import usb, usb.core
|
|
|
|
cfg = dev.get_active_configuration()
|
|
|
|
if cfg is None: # should be configured already, but eh
|
|
dev.set_configuration()
|
|
cfg = dev.get_active_configuration()
|
|
|
|
if cfg is None:
|
|
return "Couldn't get or set device configuration, aaaaa"
|
|
|
|
|
|
itf = [i for i in cfg.interfaces()
|
|
if i.bInterfaceClass == usb.CLASS_VENDOR_SPEC and
|
|
i.bInterfaceSubClass == UsbConn._SUBCLASS and
|
|
i.bInterfaceProtocol == UsbConn._PROTOCOL]
|
|
|
|
if len(itf) == 0:
|
|
return "No vendor control interface found for device"
|
|
if len(itf) != 1:
|
|
return "Multiple vendor control interfaces found for device, wtf?"
|
|
|
|
itf = itf[0]
|
|
|
|
epout = usb.util.find_descriptor(itf, custom_match =
|
|
lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_OUT)
|
|
epin = usb.util.find_descriptor(itf, custom_match =
|
|
lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_IN)
|
|
|
|
try:
|
|
# try to read the version number. if it throws, it means the usbdev
|
|
# is in use by something else
|
|
epout.write(b'\x00')
|
|
resp = epin.read(4)
|
|
if resp[0] != 0 or resp[1] != 2:
|
|
return "Device does not recognise the 'get protocol version' command"
|
|
|
|
verno = struct.unpack('<H', resp[2:])[0]
|
|
if verno < UsbConn._VER_MIN:
|
|
return "Version of device (%04x) too old, must be at least %04x" \
|
|
% (hex(verno, USbConn._VER_MIN))
|
|
if verno > UsbConn._VER_MAX:
|
|
return "Version of device (%04x) too new, must be max. %04x" \
|
|
% (hex(verno, USbConn._VER_MAX))
|
|
except usb.core.USBError:
|
|
return "Device is busy, already used by something else? (If you use "+\
|
|
"the kernel module, use a character device from /dev instead.)"
|
|
|
|
return UsbConn(dev, cfg, itf, epout, epin)
|
|
|
|
def try_find() -> Optional[UsbConn]:
|
|
import usb.core
|
|
|
|
dev = list(usb.core.find(find_all=True, idVendor=UsbConn._USB_DEFAULT_VID,
|
|
idProduct=UsbConn._USB_DEFAULT_PID))
|
|
|
|
if dev is None or len(dev) != 1:
|
|
return None
|
|
|
|
return UsbConn._open_dev(dev[0])
|
|
|
|
def is_usbdev_path(conn: str) -> bool:
|
|
# eg. cafe:1312
|
|
match_vidpid = re.match('^[0-9a-fA-F]{4}:[0-9a-fA-F]{4}$', conn)
|
|
if match_vidpid is not None:
|
|
return True
|
|
|
|
# eg. 1.123 or 1.123.1
|
|
match_busdev = re.match('^[0-9]{1}\.[0-9]{1,3}(\.[0-9]{1,3})?$', conn)
|
|
return match_busdev is not None
|
|
|
|
def try_open(conn: str) -> Union[UsbConn, str]:
|
|
import usb.core
|
|
|
|
conn_busdev = False
|
|
conntup = None
|
|
|
|
# eg. cafe:1312
|
|
match_vidpid = re.match('^([0-9a-fA-F]{4}):([0-9a-fA-F]{4})$', conn)
|
|
if match_vidpid is not None:
|
|
conntup = tuple(int(x,16) for x in match_vidpid.groups())
|
|
else:
|
|
# eg. 1.123
|
|
match_busdev = re.match('^([0-9]{1,3})\.([0-9]{1,3})(\.([0-9]{1,3}))?$', conn)
|
|
if match_busdev is not None:
|
|
conn_busdev = True
|
|
conntup = match_busdev.groups()
|
|
if conntup is not None:
|
|
if conntup[3] is None:
|
|
conntup = tuple(int(x) for x in conntup[0:2])
|
|
else:
|
|
conntup = tuple(int(x) for x in (conntup[0:2] + (conntup[3],)))
|
|
|
|
if conntup is None:
|
|
return "Could not open USB device '%s': not recognised" % conn
|
|
|
|
dev = None
|
|
if conn_busdev:
|
|
if len(conntup) == 2:
|
|
dev = list(usb.core.find(find_all=True, bus=conntup[0], address=conntup[1]))
|
|
elif len(conntup) == 3:
|
|
dev = list(usb.core.find(find_all=True, bus=conntup[0], address=conntup[1], port=conntup[2]))
|
|
else:
|
|
assert False, ("huh? conntup=%s"%repr(conntup))
|
|
else:
|
|
dev = list(usb.core.find(find_all=True, idVendor=conntup[0], idProduct=conntup[1]))
|
|
|
|
if len(dev) == 0:
|
|
return "Connect to '%s' (%s): no such device found" % \
|
|
(conn, "bus.address(.port)" if conn_busdev else "VID:PID")
|
|
if len(dev) != 1:
|
|
# TODO: nicer usb device list?
|
|
return "Connection string '%s' ambiguous, found more than one device: %s" % (conn, str(dev))
|
|
|
|
return UsbConn._open_dev(dev[0])
|
|
|
|
def read_raw(self, arr) -> int:
|
|
return self._epin.read(arr)
|
|
|
|
def write_raw(self, b: bytes) -> int:
|
|
return self._epout.write(b)
|
|
|
|
def __init__(self, dev, cfg, itf, epin, epout):
|
|
self._dev = dev
|
|
self._cfg = cfg
|
|
self._itf = itf
|
|
self._epin = epin
|
|
self._epout = epout
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, type, value, tb):
|
|
import usb.util
|
|
|
|
usb.util.release_interface(self._dev, self._itf)
|
|
usb.util.dispose_resources(self._dev)
|
|
self._epout = None
|
|
self._epin = None
|
|
self._itf = None
|
|
self._cfg = None
|
|
self._dev = None
|
|
|
|
def __str__(self):
|
|
return str(self._dev)
|
|
|
|
|
|
class ChardevConn(DevConn):
|
|
def try_find() -> Optional[ChardevConn]:
|
|
return None
|
|
|
|
def is_chardev_path(conn: str) -> bool:
|
|
return None
|
|
|
|
def try_open(conn: str) -> Union[ChardevConn, str]:
|
|
return "Could not open character device '%s'" % conn
|
|
|
|
|
|
def __init__(self, fd):
|
|
self._fd = fd
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, type, value, tb):
|
|
os.close(self._fd)
|
|
self._fd = -1
|
|
|
|
|
|
class DmjDevice:
|
|
def __init__(self, conn: DevConn):
|
|
self._conn = conn
|
|
self._buf = array.array('B')
|
|
self._buf.fromlist([0]*64)
|
|
self._bufpos = 64
|
|
|
|
def read(self, nb: int) -> bytes:
|
|
if len(self._buf) - self._bufpos > nb:
|
|
rv = bytes(self._buf[self._bufpos:self._bufpos+nb])
|
|
self._bufpos += nb
|
|
return rv
|
|
|
|
rv = list(self._buf[self._bufpos:])
|
|
|
|
while True: # TODO: timeout?
|
|
nrd = self.conn.read_raw(self._buf)
|
|
if len(rv) + nrd >= nb: # last read, will have enough now
|
|
rv = rv + list(self._buf[nb - len(rv):])
|
|
self._bufpos = nb - len(rv)
|
|
return bytes(rv)
|
|
else:
|
|
rv += list(self._buf)
|
|
|
|
# TODO: buffer(/retry) writes as well?
|
|
def write(self, b: bytes):
|
|
return self.conn.write_raw(b)
|
|
|
|
def __enter__(self):
|
|
self._conn.__enter__()
|
|
return self
|
|
|
|
def __exit__(self, type, value, tb):
|
|
self._conn.__exit__(type, value, tb)
|
|
|
|
|
|
def connect(conn: Optional[str]) -> Union[DmjDevice, str]:
|
|
# TODO: some kind of connection backend registration? shrug
|
|
if conn is None:
|
|
attempt = ChardevConn.try_find()
|
|
if attempt is not None:
|
|
return DmjDevice(attempt)
|
|
|
|
attempt = UsbConn.try_find()
|
|
if attempt is not None:
|
|
return DmjDevice(attempt)
|
|
|
|
return "no device specified, and none could be found"
|
|
|
|
if ChardevConn.is_chardev_path(conn):
|
|
attempt = ChardevConn.try_open(conn)
|
|
if isinstance(attempt, str):
|
|
return attempt
|
|
else:
|
|
return DmjDevice(attempt)
|
|
|
|
if UsbConn.is_usbdev_path(conn):
|
|
attempt = UsbConn.try_open(conn)
|
|
if isinstance(attempt, str):
|
|
return attempt
|
|
else:
|
|
return DmjDevice(attempt)
|
|
|
|
return "connection string '%s' not recognised" % conn
|
|
|