Merge pull request #1219 from a3f/write-files

remote/client: implement udisks2-using write-files command

Author: Emantor

CHANGES.rst

@@ -12,6 +12,8 @@ New Features in 24.0
   to the serial console during testing.
 - The `QEMUDriver` now has an additional ``disk_opts`` property which can be
   used to pass additional options for the disk directly to QEMU
+- labgrid-client now has a ``write-files`` subcommand to copy files onto mass
+  storage devices.
 
 Bug fixes in 24.0
 ~~~~~~~~~~~~~~~~~

contrib/completion/labgrid-client.bash

@@ -769,6 +769,40 @@ _labgrid_client_write_image()
     esac
 }
 
+_labgrid_client_write_files()
+{
+    local cur prev words cword
+    _init_completion || return
+
+    case "$prev" in
+    -w|--wait)
+        ;&
+    -p|--partition)
+        ;&
+    -t|--target-directory)
+        ;&
+    -T)
+        ;&
+    -n|--name)
+        _labgrid_complete match-names "$cur"
+        return
+        ;;
+    esac
+
+    case "$cur" in
+    -*)
+        local options="--wait --partition --target-directory --name $_labgrid_shared_options"
+        COMPREPLY=( $(compgen -W "$options" -- "$cur") )
+        ;;
+    *)
+        local args
+        _labgrid_count_args "@(-w|--wait|-p|--partition|-t|--target-directory|-T|-n|--name)" || return
+
+        _filedir
+        ;;
+    esac
+}
+
 _labgrid_client_reserve()
 {
     _labgrid_client_generic_subcommand "--wait --shell --prio"
@@ -888,6 +922,7 @@ _labgrid_client()
                                audio \
                                tmc \
                                write-image \
+                               write-files \
                                reserve \
                                cancel-reservation \
                                wait \

doc/configuration.rst

@@ -610,6 +610,14 @@ device.
      match:
        ID_PATH: 'pci-0000:06:00.0-usb-0:1.3.2:1.0-scsi-0:0:0:3'
 
+Writing images to disk requires installation of ``dd`` or optionally
+``bmaptool`` on the same system as the block device.
+
+For mounting the file system and writing into it,
+`PyGObject <https://pygobject.readthedocs.io/>`_ must be installed.
+For Debian, the necessary packages are `python3-gi` and `gir1.2-udisks-2.0`.
+This is not required for writing images to disks.
+
 Arguments:
   - match (dict): key and value pairs for a udev match, see `udev Matching`_
 
@@ -622,7 +630,7 @@ A :any:`NetworkUSBMassStorage` resource describes a USB memory stick or similar
 device available on a remote computer.
 
 The NetworkUSBMassStorage can be used in test cases by calling the
-``write_image()``, and ``get_size()`` functions.
+``write_files()``, ``write_image()``, and ``get_size()`` functions.
 
 SigrokDevice
 ~~~~~~~~~~~~

labgrid/driver/usbstoragedriver.py

@@ -1,17 +1,20 @@
 import enum
 import os
+import pathlib
 import time
 import subprocess
 
 import attr
 
 from ..factory import target_factory
+from ..resource.remote import RemoteUSBResource
 from ..step import step
 from ..util.managedfile import ManagedFile
 from .common import Driver
 from ..driver.exception import ExecutionError
 
 from ..util.helper import processwrapper
+from ..util.agentwrapper import AgentWrapper
 from ..util import Timeout
 
 
@@ -40,12 +43,69 @@ class USBStorageDriver(Driver):
         default=None,
         validator=attr.validators.optional(attr.validators.instance_of(str))
     )
+    WAIT_FOR_MEDIUM_TIMEOUT = 10.0 # s
+    WAIT_FOR_MEDIUM_SLEEP = 0.5 # s
+
+    def __attrs_post_init__(self):
+        super().__attrs_post_init__()
+        self.wrapper = None
 
     def on_activate(self):
-        pass
+        host = self.storage.host if isinstance(self.storage, RemoteUSBResource) else None
+        self.wrapper = AgentWrapper(host)
+        self.proxy = self.wrapper.load('udisks2')
 
     def on_deactivate(self):
-        pass
+        self.wrapper.close()
+        self.wrapper = None
+        self.proxy = None
+
+    @Driver.check_active
+    @step(args=['sources', 'target', 'partition', 'target_is_directory'])
+    def write_files(self, sources, target, partition, target_is_directory=True):
+        """
+        Write the file(s) specified by filename(s) to the
+        bound USB storage partition.
+
+        Args:
+            sources (List[str]): path(s) to the file(s) to be copied to the bound USB storage
+                partition.
+            target (str): target directory or file to copy to
+            partition (int): mount the specified partition or None to mount the whole disk
+            target_is_directory (bool): Whether target is a directory
+        """
+
+        self.devpath = self._get_devpath(partition)
+        mount_path = self.proxy.mount(self.devpath)
+
+        try:
+            # (pathlib.PurePath(...) / "/") == "/", so we turn absolute paths into relative
+            # paths with respect to the mount point here
+            target_rel = target.relative_to(target.root) if target.root is not None else target
+            target_path = str(pathlib.PurePath(mount_path) / target_rel)
+
+            copied_sources = []
+
+            for f in sources:
+                mf = ManagedFile(f, self.storage)
+                mf.sync_to_resource()
+                copied_sources.append(mf.get_remote_path())
+
+            if target_is_directory:
+                args = ["cp", "-t", target_path] + copied_sources
+            else:
+                if len(sources) != 1:
+                    raise ValueError("single source argument required when target_is_directory=False")
+
+                args = ["cp", "-T", copied_sources[0], target_path]
+
+            processwrapper.check_output(self.storage.command_prefix + args)
+            self.proxy.unmount(self.devpath)
+        except:
+            # We are going to die with an exception anyway, so no point in waiting
+            # to make sure everything has been written before continuing
+            self.proxy.unmount(self.devpath, lazy=True)
+            raise
 
     @Driver.check_active
     @step(args=['filename'])
@@ -68,22 +128,10 @@ def write_image(self, filename=None, mode=Mode.DD, partition=None, skip=0, seek=
         mf = ManagedFile(filename, self.storage)
         mf.sync_to_resource()
 
-        # wait for medium
-        timeout = Timeout(10.0)
-        while not timeout.expired:
-            try:
-                if self.get_size() > 0:
-                    break
-                time.sleep(0.5)
-            except ValueError:
-                # when the medium gets ready the sysfs attribute is empty for a short time span
-                continue
-        else:
-            raise ExecutionError("Timeout while waiting for medium")
+        self._wait_for_medium(partition)
 
-        partition = "" if partition is None else partition
+        target = self._get_devpath(partition)
         remote_path = mf.get_remote_path()
-        target = f"{self.storage.path}{partition}"
 
         if mode == Mode.DD:
             self.logger.info('Writing %s to %s using dd.', remote_path, target)
@@ -139,12 +187,41 @@ def write_image(self, filename=None, mode=Mode.DD, partition=None, skip=0, seek=
             print_on_silent_log=True
         )
 
+    def _get_devpath(self, partition):
+        partition = "" if partition is None else partition
+        # simple concatenation is sufficient for USB mass storage
+        return f"{self.storage.path}{partition}"
+
     @Driver.check_active
-    @step(result=True)
-    def get_size(self):
-        args = ["cat", f"/sys/class/block/{self.storage.path[5:]}/size"]
+    def _wait_for_medium(self, partition):
+        timeout = Timeout(self.WAIT_FOR_MEDIUM_TIMEOUT)
+        while not timeout.expired:
+            if self.get_size(partition) > 0:
+                break
+            time.sleep(self.WAIT_FOR_MEDIUM_SLEEP)
+        else:
+            raise ExecutionError("Timeout while waiting for medium")
+
+    @Driver.check_active
+    @step(args=['partition'], result=True)
+    def get_size(self, partition=None):
+        """
+        Get the size of the bound USB storage root device or partition.
+
+        Args:
+            partition (int or None): optional, get size of the specified partition or None for
+                getting the size of the root device (defaults to None)
+
+        Returns:
+            int: size in bytes
+        """
+        args = ["cat", f"/sys/class/block/{self._get_devpath(partition)[5:]}/size"]
         size = subprocess.check_output(self.storage.command_prefix + args)
-        return int(size)*512
+        try:
+            return int(size) * 512
+        except ValueError:
+            # when the medium gets ready the sysfs attribute is empty for a short time span
+            return 0
 
 
 @target_factory.reg_driver

labgrid/remote/client.py

@@ -5,6 +5,7 @@
 import contextlib
 import enum
 import os
+import pathlib
 import subprocess
 import traceback
 import logging
@@ -1214,6 +1215,32 @@ def tmc_channel(self):
         for k, v in sorted(data.items()):
             print(f"{k:<16s} {str(v):<10s}")
 
+    def write_files(self):
+        place = self.get_acquired_place()
+        target = self._get_target(place)
+        name = self.args.name
+        drv = self._get_driver_or_new(target, "USBStorageDriver", activate=False, name=name)
+        drv.storage.timeout = self.args.wait
+        target.activate(drv)
+
+        try:
+            if self.args.partition == 0:
+                self.args.partition = None
+
+            if self.args.rename:
+                if len(self.args.SOURCE) != 2:
+                    self.args.parser.error("the following arguments are required: SOURCE DEST")
+
+                drv.write_files([self.args.SOURCE[0]], self.args.SOURCE[1],
+                                self.args.partition, target_is_directory=False)
+            else:
+                drv.write_files(self.args.SOURCE, self.args.target_directory,
+                                self.args.partition, target_is_directory=True)
+        except subprocess.CalledProcessError as e:
+            raise UserError(f"could not copy files to network usb storage: {e}")
+        except FileNotFoundError as e:
+            raise UserError(e)
+
     def write_image(self):
         place = self.get_acquired_place()
         target = self._get_target(place)
@@ -1761,6 +1788,27 @@ def main():
     tmc_subparser.add_argument('action', choices=['info', 'values'])
     tmc_subparser.set_defaults(func=ClientSession.tmc_channel)
 
+    subparser = subparsers.add_parser('write-files', help="copy files onto mass storage device",
+                                      usage="%(prog)s [OPTION]... -T SOURCE DEST\n" +
+                                     "       %(prog)s [OPTION]... [-t DIRECTORY] SOURCE...")
+    subparser.add_argument('-w', '--wait', type=float, default=10.0,
+                           help='storage poll timeout in seconds')
+    subparser.add_argument('-p', '--partition', type=int, choices=range(0, 256),
+                           metavar='0-255', default=1,
+                           help='partition number to mount or 0 to mount whole disk (default: %(default)s)')
+    group = subparser.add_mutually_exclusive_group()
+    group.add_argument('-t', '--target-directory', type=pathlib.PurePath, metavar='DIRECTORY',
+                           default=pathlib.PurePath("/"),
+                           help='copy all SOURCE files into DIRECTORY (default: partition root)')
+    group.add_argument('-T', action='store_true', dest='rename',
+                           help='copy SOURCE file and rename to DEST')
+    subparser.add_argument('--name', '-n', help="optional resource name")
+    subparser.add_argument('SOURCE', type=pathlib.PurePath, nargs='+',
+                           help='source file(s) to copy')
+    subparser.add_argument('DEST', type=pathlib.PurePath, nargs='?',
+                           help='destination file name for SOURCE')
+    subparser.set_defaults(func=ClientSession.write_files, parser=subparser)
+
     subparser = subparsers.add_parser('write-image', help="write an image onto mass storage")
     subparser.add_argument('-w', '--wait', type=float, default=10.0)
     subparser.add_argument('-p', '--partition', type=int, help="partition number to write to")