Improvements to scpair.py (#125)

ReversePuck/scpair.py

@@ -57,21 +57,30 @@ FEAT_LEN = 64  # Valve feature reports are exactly 64 bytes
 
 
 def _sysfs_ids(node):
-    """Return (vid, pid) for a /dev/hidrawN node, or (None, None)."""
+    """Return (vid, pid, serial) for a /dev/hidrawN node, or (None, None, None)."""
     name = os.path.basename(node)
     try:
         # /sys/class/hidraw/hidrawN/device/uevent has HID_ID=0003:000028DE:00001304
         ue = open("/sys/class/hidraw/%s/device/uevent" % name).read()
     except Exception:
-        return (None, None)
+        return (None, None, None)
+    
+    vid = None
+    pid = None 
+    serial = None
+
     for line in ue.splitlines():
         if line.startswith("HID_ID="):
             parts = line.split(":")
             try:
-                return (int(parts[1], 16) & 0xFFFF, int(parts[2], 16) & 0xFFFF)
-            except Exception:
-                return (None, None)
-    return (None, None)
+                vid = int(parts[1], 16) & 0xFFFF
+                pid = int(parts[2], 16) & 0xFFFF
+            except:
+                pass
+        if line.startswith("HID_UNIQ="):
+            serial = line.split('=')[1]
+    
+    return (vid, pid, serial)
 
 
 def _first_usage_page(fd):
@@ -99,10 +108,11 @@ def _first_usage_page(fd):
 
 
 class HidDev:
-    def __init__(self, node, vid, pid):
+    def __init__(self, node, vid, pid, serial):
         self.node = node
         self.vid = vid
         self.pid = pid
+        self.serial = serial
         self.fd = os.open(node, os.O_RDWR)
         self.usage_page = _first_usage_page(self.fd)
 
@@ -136,16 +146,25 @@ class HidDev:
         except Exception:
             return ""
 
+def nodesort(nodeobj):
+    """Sorts the node path properly (numerically)."""
+
+    if isinstance(nodeobj, HidDev):
+        return nodesort(nodeobj.node)
+
+    if isinstance(nodeobj, str):
+        return int(nodeobj.replace("/dev/hidraw", ""))
+
 
 def enumerate_devices():
     """Return dict role -> list[HidDev]: {'puck':[...], 'ctrl':[...]} (control interfaces only)."""
     out = {"puck": [], "ctrl": []}
-    for node in sorted(glob.glob("/dev/hidraw*")):
-        vid, pid = _sysfs_ids(node)
+    for node in sorted(glob.glob("/dev/hidraw*"), key=nodesort):
+        vid, pid, serial = _sysfs_ids(node)
         if vid != VALVE_VID:
             continue
         try:
-            dev = HidDev(node, vid, pid)
+            dev = HidDev(node, vid, pid, serial)
         except PermissionError:
             print("permission denied on %s (run with sudo or install the udev rule)" % node, file=sys.stderr)
             continue
@@ -169,89 +188,151 @@ def _ser16(serial):
     return b + b"\x00" * (16 - len(b))
 
 
-def read_slots(puck):
-    """Read the puck's 4 bond slots (each puck control interface == one slot). Returns list of dicts."""
+def read_slots(puck_list):
+    """Read all bond slots from all attached pucks (each puck control interface == one slot)."""
+    pucks = []
     slots = []
-    for i, dev in enumerate(sorted(puck, key=lambda d: d.node)):
+
+    last_dev = None
+    idx = 0
+    for dev in sorted(puck_list, key=nodesort):
+        if last_dev is not None and dev.serial != last_dev.serial:
+            # This is a new puck, append the old puck and all its slots to the list.
+            pucks.append({"serial": last_dev.serial, "slots": slots})
+            slots = []
+            idx = 0
         try:
             dev.set_feature(2, 0xA3)
             r = dev.get_feature(2)  # [02 A3 18 <8 uuid><16 serial>]
             rec = r[3:3 + 24]
             used = any(rec[0:8])
             serial = rec[8:24].split(b"\x00")[0].decode("latin1", "replace") if used else ""
-            slots.append({"idx": i, "used": used, "serial": serial, "rec": rec})
+            slots.append({"idx": idx, "dev": dev, "used": used, "serial": serial, "rec": rec})
         except Exception as ex:
-            slots.append({"idx": i, "used": False, "serial": "", "rec": b"", "err": str(ex)})
-    return slots
+            slots.append({"idx": idx, "dev": dev, "used": False, "serial": "", "rec": b"", "err": str(ex)})
+        
+        last_dev = dev
+        idx = idx + 1
+
+    pucks.append({"serial": dev.serial, "slots": slots})
+    return pucks
+
+
+def get_puck_slot_for_writing(puck_serial=None, slot=None):
+    devs = enumerate_devices()
+    pucks = sorted(devs["puck"], key=nodesort)
+    if not pucks:
+        raise SystemExit("no puck (28DE:1304) control interface found")
+
+    puck_objects = read_slots(pucks)
+    if puck_serial is None and len(puck_objects) > 1:
+        raise SystemExit("Multiple Pucks connected, you need to specify its serial number")
+    
+    puck_object = None
+
+    if puck_serial is not None:
+        for puck in puck_objects:
+            if puck['serial'] == puck_serial:
+                puck_object = puck
+                break
+        
+        if puck_object is None: 
+            raise SystemExit("Can't find a puck with that serial number!")
+        
+    else: 
+        puck_object = puck_objects[0]
+
+
+    slots = puck_object['slots']
+
+    if slot is None:
+        slot_object = next((s for s in slots if not s["used"]), None)
+    else: 
+        slot_object = next((s for s in slots if s["idx"] == slot), None)
+
+    if slot_object is None:
+        raise SystemExit("This Puck has all slots occupied!")
+
+    
+    return (puck_object['serial'], slot_object)
+    
 
 
-def pair(slot=None, pkey=None):
+def pair(puck_serial=None, puck_slot=None, controller_slot=0, puck_name=None, controller_name=None, pkey=None):
     devs = enumerate_devices()
-    pucks = sorted(devs["puck"], key=lambda d: d.node)
-    ctrls = sorted(devs["ctrl"], key=lambda d: d.node)
+    pucks = sorted(devs["puck"], key=nodesort)
+    ctrls = sorted(devs["ctrl"], key=nodesort)
     if not pucks:
         raise SystemExit("no puck (28DE:1304) control interface found")
     if not ctrls:
         raise SystemExit("no controller (28DE:1302) control interface found — plug the ReversePuck controller in")
     ctrl = ctrls[0]
-    puck_serial = pucks[0].read_serial(2) or "FXB000000000"
-    ctrl_serial = ctrl.read_serial(1) or "FXA000000000"
 
-    if slot is None:
-        slots = read_slots(pucks)
-        slot = next((s["idx"] for s in slots if not s["used"]), 0)
-    if slot >= len(pucks):
-        raise SystemExit("slot %d has no puck control interface" % slot)
+    if controller_name is None: 
+        ctrl_serial_str = ctrl.read_serial(1) or "FXA000000000"
+    else: 
+        ctrl_serial_str = controller_name[:15]
+
+    puck_serial, slot_obj = get_puck_slot_for_writing(puck_serial, puck_slot)  
+
+    if puck_name is None: 
+        puck_serial_str = puck_serial
+    else: 
+        puck_serial_str = puck_name[:15]
 
     if pkey is None: 
         r = os.urandom(8)
     else: 
         r = pkey
 
-    key = b"esb/bond\x00" if slot == 0 else b"esb/bond_2\x00"
-    print("pairing puck %s slot %d <-> controller %s  key=%s" %
-          (puck_serial, slot, ctrl_serial, r.hex()))
+    key = b"esb/bond\x00" if controller_slot == 0 else b"esb/bond_2\x00"
+
+    print("Pairing key: %s" % (r.hex()))
+
+    # Puck name gets written to controller (defaults to puck serial).
+    # Controller name gets written to Puck
+
+    print("Puck: %s slot %d controller %s" % (puck_serial, slot_obj['idx'], ctrl_serial_str))
+    print("Controller: %s slot %d puck %s" % (ctrl.read_serial(1), controller_slot, puck_serial_str))
 
     # puck slot: 0xA2 [r1 r2][ctrl serial]
-    pucks[slot].set_feature(2, 0xA2, r + _ser16(ctrl_serial))
+    slot_obj['dev'].set_feature(2, 0xA2, r + _ser16(ctrl_serial_str))
     # controller esb/bond: 0xEE [key\0][r1 r2][puck serial] + 0xEF [key\0]
-    ctrl.set_feature(1, 0xEE, key + r + _ser16(puck_serial))
+    ctrl.set_feature(1, 0xEE, key + r + _ser16(puck_serial_str))
     ctrl.set_feature(1, 0xEF, key)
     time.sleep(0.1)
     # reboot controller into wireless mode (magic 0xA427AF52)
     ctrl.set_feature(1, 0x95, bytes([0x52, 0xAF, 0x27, 0xA4]))
     print("paired. controller rebooting into wireless — move the puck to your host.")
 
-def write_puck_slot(slot=None, ctrl_serial=None, pkey=None):
-    devs = enumerate_devices()
-    pucks = sorted(devs["puck"], key=lambda d: d.node)
-    if not pucks:
-        raise SystemExit("no puck (28DE:1304) control interface found")
+def write_puck_slot(puck_serial=None, puck_slot=None, ctrl_name=None, pkey=None):
 
-    if slot is None:
-        slots = read_slots(pucks)
-        slot = next((s["idx"] for s in slots if not s["used"]), 0)
-    if slot >= len(pucks):
-        raise SystemExit("slot %d has no puck control interface" % slot)
+    puck_serial, slot_obj = get_puck_slot_for_writing(puck_serial, puck_slot)  
 
-    
-    if ctrl_serial is None and pkey is None: 
-        print("wiping puck slot %d" % (slot))
-        pucks[slot].set_feature(2, 0xA2, b"\x00" * 24)
+    if ctrl_name is None and pkey is None: 
+        print("wiping puck %s slot %d" % (puck_serial, slot_obj['idx']))
+        slot_obj['dev'].set_feature(2, 0xA2, b"\x00" * 24)
     else:
+
+        if ctrl_name is None or len(ctrl_name) > 15: 
+            raise SystemExit("Invalid controller name.")
+
         if pkey is None: 
             raise SystemExit("No key provided!")
         
-        print("writing to puck slot %d: controller %s key=%s" %
-          (slot, ctrl_serial, pkey.hex()))
+        print("writing to puck %s slot %d: controller %s key=%s" %
+          (puck_serial, slot_obj['idx'], ctrl_name, pkey.hex()))
         # puck slot: 0xA2 [r1 r2][ctrl serial]
-        pucks[slot].set_feature(2, 0xA2, pkey + _ser16(ctrl_serial))
+
+        print(slot_obj['dev'].node)
+
+        slot_obj['dev'].set_feature(2, 0xA2, pkey + _ser16(ctrl_name))
 
     print("done")
 
 def write_controller_slot(slot=0, puck_serial=None, pkey=None):
     devs = enumerate_devices()
-    ctrls = sorted(devs["ctrl"], key=lambda d: d.node)
+    ctrls = sorted(devs["ctrl"], key=nodesort)
     if not ctrls:
         raise SystemExit("no controller (28DE:1302) control interface found — plug the ReversePuck controller in")
     
@@ -291,36 +372,26 @@ def write_controller_slot(slot=0, puck_serial=None, pkey=None):
         print("paired. controller rebooting into wireless — move the puck to your host.")
 
 
-def unpair(slot):
-    devs = enumerate_devices()
-    pucks = sorted(devs["puck"], key=lambda d: d.node)
-    if slot >= len(pucks):
-        raise SystemExit("slot %d has no puck control interface" % slot)
-    pucks[slot].set_feature(2, 0xA2, b"\x00" * 24)  # clear puck slot
-    # also clear the controller bond if present
-    for ctrl in devs["ctrl"]:
-        key = b"esb/bond\x00" if slot == 0 else b"esb/bond_2\x00"
-        try:
-            ctrl.set_feature(1, 0xEE, key + b"\x00" * 24)
-            ctrl.set_feature(1, 0xEF, key)
-        except Exception:
-            print("Slot wipe on controller failed.")
-    print("cleared slot %d" % slot)
-
 
 
 def cmd_list():
     devs = enumerate_devices()
-    pucks = sorted(devs["puck"], key=lambda d: d.node)
-    ctrls = sorted(devs["ctrl"], key=lambda d: d.node)
+    pucks = sorted(devs["puck"], key=nodesort)
+    ctrls = sorted(devs["ctrl"], key=nodesort)
     print("controllers:", [(d.node, "%04X" % d.pid, d.read_serial(1)) for d in ctrls] or "(none)")
+
     if pucks:
-        print("puck:", pucks[0].read_serial(2))
-        for s in read_slots(pucks):
-            if s["used"]:
-                print("  slot %d: %s (key %s)" % (s["idx"], s["serial"], s["rec"][:8].hex()))
-            else: 
-                print("  slot %d: Empty" % (s["idx"]))
+        ps = read_slots(pucks)
+        for puck in ps:
+            slots = puck['slots']
+
+            print(f"Puck: {puck['serial']}")
+
+            for s in slots: 
+                if s["used"]:
+                    print("  slot %d: %s (key %s)" % (s["idx"], s["serial"], s["rec"][:8].hex()))
+                else: 
+                    print("  slot %d: Empty" % (s["idx"]))
     else:
         print("puck: (none)")
 
@@ -339,19 +410,26 @@ def main():
     sub = ap.add_subparsers(dest="cmd", required=True)
     sub.add_parser("list")
     p = sub.add_parser("pair")
-    p.add_argument("--slot", type=int, default=None)
+    #def pair(puck_serial=None, puck_slot=None, controller_slot=0, puck_name=None, controller_name=None, pkey=None):
+    p.add_argument("--puck-serial", type=str, default=None)
+    p.add_argument("--puck-slot", type=int, default=None)
+    p.add_argument("--puck-name", type=str, default=None)
+    p.add_argument("--controller-slot", type=int, default=0)
+    p.add_argument("--controller-name", type=str, default=None)
     p.add_argument("--key", type=steam_key_check, default=None)
-    u = sub.add_parser("unpair")
-    u.add_argument("--slot", type=int, required=True)
+
 
     s = sub.add_parser("write-puck")
-    s.add_argument("--slot", type=int, default=None)
-    s.add_argument("--serial", type=str, default=None)
+    #def write_puck_slot(puck_serial=None, puck_slot=None, ctrl_name=None, pkey=None):
+    s.add_argument("--puck-serial", type=str, default=None)
+    s.add_argument("--puck-slot", type=int, default=None)
+    s.add_argument("--controller-name", type=str, default=None)
     s.add_argument("--key", type=steam_key_check, default=None)
 
     c = sub.add_parser("write-controller")
-    c.add_argument("--slot", type=int, default=None)
-    c.add_argument("--serial", type=str, default=None)
+    #def write_controller_slot(slot=0, puck_serial=None, pkey=None):
+    c.add_argument("--controller-slot", type=int, default=None)
+    c.add_argument("--puck-name", type=str, default=None)
     c.add_argument("--key", type=steam_key_check, default=None)
 
 
@@ -360,13 +438,11 @@ def main():
     if args.cmd == "list":
         cmd_list()
     elif args.cmd == "pair":
-        pair(args.slot, args.key)
-    elif args.cmd == "unpair":
-        unpair(args.slot)
+        pair(args.puck_serial, args.puck_slot, args.controller_slot, args.puck_name, args.controller_name, args.key)
     elif args.cmd == "write-puck":
-        write_puck_slot(args.slot, args.serial, args.key)
+        write_puck_slot(args.puck_serial, args.puck_slot, args.controller_name, args.key)
     elif args.cmd == "write-controller":
-        write_controller_slot(args.slot, args.serial, args.key)
+        write_controller_slot(args.controller_slot, args.puck_name, args.key)