DragonProbe/host/dpctl/connection.py

277 lines
8.5 KiB
Python
Raw Normal View History

2021-07-24 02:24:55 +00:00
from __future__ import annotations
2021-07-25 03:06:39 +00:00
import glob
2021-07-24 02:24:55 +00:00
import os
2021-07-25 01:13:20 +00:00
import re
import struct
2021-07-25 01:43:38 +00:00
import sys
2021-07-24 02:24:55 +00:00
import abc
from typing import *
# TODO: EVERYTHING
# * implement connecting to a device
# * implement lower-level device commands
# * implement device commands
2021-07-25 01:13:20 +00:00
class DevConn:
2021-07-25 01:43:38 +00:00
_VER_MIN = 0x0010
_VER_MAX = 0x00ff # ???
2021-07-25 01:13:20 +00:00
pass
2021-07-24 02:24:55 +00:00
2021-07-25 01:43:38 +00:00
2021-07-25 01:13:20 +00:00
class UsbConn(DevConn):
_USB_DEFAULT_VID = 0xcafe
_USB_DEFAULT_PID = 0x1312
_SUBCLASS = ord('D')
_PROTOCOL = ord('P')
2021-07-25 01:13:20 +00:00
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)
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.)"
2021-07-25 01:43:38 +00:00
if len(resp) < 4 or 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 < DevConn._VER_MIN:
return "Version of device (%04x) too old, must be at least %04x" \
% (hex(verno, DevConn._VER_MIN))
if verno > DevConn._VER_MAX:
return "Version of device (%04x) too new, must be max. %04x" \
% (hex(verno, DevConn._VER_MAX))
2021-07-25 16:16:55 +00:00
return UsbConn(dev, cfg, itf, epin, epout)
2021-07-24 02:24:55 +00:00
2021-07-25 01:13:20 +00:00
def try_find() -> Optional[UsbConn]:
import usb.core
2021-07-24 02:24:55 +00:00
2021-07-25 01:13:20 +00:00
dev = list(usb.core.find(find_all=True, idVendor=UsbConn._USB_DEFAULT_VID,
idProduct=UsbConn._USB_DEFAULT_PID))
2021-07-24 02:24:55 +00:00
2021-07-25 01:13:20 +00:00
if dev is None or len(dev) != 1:
return None
2021-07-24 02:24:55 +00:00
2021-07-25 01:43:38 +00:00
rv = UsbConn._open_dev(dev[0])
return None if isinstance(rv, str) else rv
2021-07-24 02:24:55 +00:00
2021-07-25 02:35:39 +00:00
def is_path(conn: str) -> bool:
2021-07-25 01:13:20 +00:00
# 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
2021-07-24 02:24:55 +00:00
def try_open(conn: str) -> Union[UsbConn, str]:
2021-07-25 01:13:20 +00:00
import usb.core
2021-07-24 02:24:55 +00:00
2021-07-25 01:13:20 +00:00
conn_busdev = False
conntup = None
2021-07-24 02:24:55 +00:00
2021-07-25 01:13:20 +00:00
# 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
2021-07-24 02:24:55 +00:00
def __enter__(self):
return self
def __exit__(self, type, value, tb):
2021-07-25 01:13:20 +00:00
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)
2021-07-24 02:24:55 +00:00
class ChardevConn(DevConn):
_DEVCLASSNAME = "dragonprobe"
2021-07-25 01:43:38 +00:00
2021-07-24 02:24:55 +00:00
def try_find() -> Optional[ChardevConn]:
2021-07-25 01:43:38 +00:00
if sys.platform != 'linux':
return None
opts = glob.glob("/dev/%s-*" % ChardevConn._DEVCLASSNAME)
if len(opts) != 1:
return None
rv = ChardevConn.try_open(opts[0])
return None if isinstance(rv, str) else rv
2021-07-24 02:24:55 +00:00
2021-07-25 02:35:39 +00:00
def is_path(conn: str) -> bool:
2021-07-25 01:43:38 +00:00
if sys.platform != 'linux':
return False
return re.match('^/dev/%s-[0-9]+$' % ChardevConn._DEVCLASSNAME, conn) is not None
2021-07-24 02:24:55 +00:00
def try_open(conn: str) -> Union[ChardevConn, str]:
2021-07-25 01:43:38 +00:00
if sys.platform != 'linux':
return "Chardev connections not available on %s, as these require "+\
"a Linux kernel module" % sys.platform
try:
fd = os.open(conn, os.O_RDWR)
if fd < 0:
raise OSError("Negative file descriptor returned")
except OSError as e:
return "Could not open character device '%s': %s" % \
(conn, e.message if hasattr(e, 'message') else e.strerror)
os.write(fd, b'\x00')
resp = os.read(fd, 4)
2021-07-24 02:24:55 +00:00
2021-07-25 01:43:38 +00:00
if len(resp) < 4 or 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 < DevConn._VER_MIN:
return "Version of device (%04x) too old, must be at least %04x" \
% (hex(verno, DevConn._VER_MIN))
if verno > DevConn._VER_MAX:
return "Version of device (%04x) too new, must be max. %04x" \
% (hex(verno, DevConn._VER_MAX))
return ChardevConn(fd)
def read_raw(self, arr) -> int:
blob = os.read(self._fd, len(arr))
for i in range(len(blob)): # TODO: memcpy?
arr[i] = blob[i]
return len(blob)
def write_raw(self, b: bytes) -> int:
return os.write(self._fd, b)
2021-07-24 02:24:55 +00:00
2021-07-25 01:43:38 +00:00
def __init__(self, fd: int):
2021-07-24 02:24:55 +00:00
self._fd = fd
def __enter__(self):
return self
def __exit__(self, type, value, tb):
os.close(self._fd)
self._fd = -1
2021-07-25 02:35:39 +00:00
_BACKENDS = [ChardevConn, UsbConn]
2021-07-25 01:13:20 +00:00
2021-07-25 02:35:39 +00:00
def connect(conn: Optional[str]) -> Union[DevConn, str]:
global _BACKENDS
2021-07-25 01:13:20 +00:00
2021-07-24 02:24:55 +00:00
if conn is None:
2021-07-25 02:35:39 +00:00
for backend in _BACKENDS:
attempt = backend.try_find()
if attempt is not None:
return attempt
2021-07-24 02:24:55 +00:00
return "no device specified, and none could be found"
2021-07-25 02:35:39 +00:00
for backend in _BACKENDS:
if backend.is_path(conn):
return backend.try_open(conn)
2021-07-24 02:24:55 +00:00
return "connection string '%s' not recognised" % conn
2021-07-25 02:35:39 +00:00
def register_backend(backend):
global _BACKENDS
_BACKENDS.append(backend)