add support for rpiboot to enable rpi as msd

Allow for using the rpiboot tool to start up a raspberry pi cm 4 or 5
as a MassStorageDevice so that a USBStorageDriver can be used to bootstrap
a raspberry pi cm 4 or 5.

https://github.com/raspberrypi/usbboot
Signed-off-by: Benjamin B. Frost <[email protected]>

Author: benjamin1313

labgrid/driver/__init__.py

@@ -48,3 +48,4 @@
 from .deditecrelaisdriver import DeditecRelaisDriver
 from .dediprogflashdriver import DediprogFlashDriver
 from .httpdigitaloutput import HttpDigitalOutputDriver
+from .rpibootdriver import RpibootDriver

labgrid/driver/rpibootdriver.py

@@ -0,0 +1,38 @@
+import attr
+
+from ..factory import target_factory
+from ..step import step
+from .common import Driver
+from ..util.helper import processwrapper
+
+@target_factory.reg_driver
+@attr.s(eq=False)
+class RpibootDriver(Driver):
+    bindings = {
+        "rpi": {"RpibootDevice", "NetworkRpibootDevice"},
+        }
+
+    image = attr.ib(default=None)
+
+    def __attrs_post_init__(self):
+        super().__attrs_post_init__()
+        if self.target.env:
+            self.tool = self.target.env.config.get_tool('rpiboot')
+        else:
+            self.tool = 'rpiboot'
+
+    def on_activate(self):
+        pass
+
+    def on_deactivate(self):
+        pass
+
+    @Driver.check_active
+    @step(args=['filename'])
+    def enable(self, filename=None,):
+        # Switch raspberry pi into MassStorageDevice mode using the rpiboot tool
+        args = []
+        processwrapper.check_output(
+            self.rpi.command_prefix + [self.tool] + args,
+            print_on_silent_log=True
+        )

labgrid/remote/exporter.py

@@ -546,6 +546,26 @@ def __attrs_post_init__(self):
         self.data["cls"] = f"Remote{self.cls}".replace("Network", "")
 
 
+@attr.s(eq=False)
+class RpibootExport(USBGenericExport):
+    """ResourceExport for raspberry pi in boot mode"""
+
+    def __attrs_post_init__(self):
+        super().__attrs_post_init__()
+
+    def _get_params(self):
+        """Helper function to return parameters"""
+        return {
+            "host": self.host,
+            "busnum": self.local.busnum,
+            "devnum": self.local.devnum,
+            "path": self.local.path,
+            "vendor_id": self.local.vendor_id,
+            "model_id": self.local.model_id,
+            "serial_id": self.local.serial_id,
+        }
+
+
 exports["AndroidFastboot"] = USBGenericExport
 exports["AndroidUSBFastboot"] = USBGenericRemoteExport
 exports["DFUDevice"] = USBGenericExport
@@ -569,6 +589,7 @@ def __attrs_post_init__(self):
 exports["HIDRelay"] = USBHIDRelayExport
 exports["USBFlashableDevice"] = USBFlashableExport
 exports["LXAUSBMux"] = USBGenericExport
+exports["RpibootDevice"] = RpibootExport
 
 
 @attr.s(eq=False)

labgrid/resource/__init__.py

@@ -32,6 +32,7 @@
     USBSerialPort,
     USBTMC,
     USBVideo,
+    RpibootDevice,
 )
 from .common import Resource, ResourceManager, ManagedResource
 from .ykushpowerport import YKUSHPowerPort, NetworkYKUSHPowerPort

labgrid/resource/remote.py

@@ -414,3 +414,14 @@ class RemoteNFSProvider(NetworkResource):
 @attr.s(eq=False)
 class RemoteHTTPProvider(RemoteBaseProvider):
     pass
+
+@target_factory.reg_resource
+@attr.s(eq=False)
+class NetworkRpibootDevice(RemoteUSBResource):
+    """The NetworkRpibootDevice describes a remotely accessible raspberry pi ready for rpiboot"""
+
+    serial_id = attr.ib(validator=attr.validators.optional(attr.validators.instance_of(str)))
+
+    def __attrs_post_init__(self):
+        self.timeout = 10.0
+        super().__attrs_post_init__()

labgrid/resource/udev.py

@@ -803,3 +803,28 @@ def update(self):
             self.index = int(self.read_attr('base')) + self.pin
         else:
             self.index = None
+
+@target_factory.reg_resource
+@attr.s(eq=False)
+class RpibootDevice(USBResource):
+    """The RpibootDevice describes an attached raspberry pi device in boot mode,
+    it is identified via USB using udev.
+    """
+
+    def filter_match(self, device):
+        match = (device.properties.get('ID_VENDOR_ID'), device.properties.get('ID_MODEL_ID'))
+
+        if match not in [("0a5c", "2711"),  # rpi4
+                         ("0a5c", "2712"),  # rpi5
+                         ]:
+            return False
+
+        return super().filter_match(device)
+
+    @property
+    def serial_id(self):
+        device = self._get_usb_device()
+        if device:
+            return str(device.properties.get('ID_SERIAL_SHORT'))
+
+        return None