[CERTTF-316] Use attachments for provisioning (muxpi) (#265)

* refactor(muxpi): introduce methods to provide common functionality
* feat(muxpi): first download image to agent, then pipe to controller for flashing
* refactor(muxpi): add `check=True` to `subprocess.run` to raise exception on error
* feat(muxpi): Introduce `use_attachment` to control provisioning source
* docs: add `use_attachment` to documentation example with attachments

device-connectors/src/testflinger_device_connectors/devices/muxpi/muxpi.py

@@ -14,14 +14,16 @@
 
 """Ubuntu Raspberry PI muxpi support code."""
 
+from contextlib import contextmanager
 import json
 import logging
+from pathlib import Path
+import requests
 import subprocess
-import shlex
+import tempfile
 import time
-import urllib.request
-from contextlib import contextmanager
-from pathlib import Path
+from typing import Optional, Union
+import urllib
 
 import yaml
 
@@ -33,6 +35,11 @@
 logger = logging.getLogger(__name__)
 
 
+# should mirror `testflinger_agent.config.ATTACHMENTS_DIR`
+# [TODO] Merge both constants into testflinger.common
+ATTACHMENTS_DIR = "attachments"
+
+
 class MuxPi:
     """Device Connector for MuxPi."""
 
@@ -57,6 +64,20 @@ def __init__(self, config=None, job_data=None):
         self.agent_name = self.config.get("agent_name")
         self.mount_point = Path("/mnt") / self.agent_name
 
+    def get_ssh_options(self):
+        return (
+            "-o",
+            "StrictHostKeyChecking=no",
+            "-o",
+            "UserKnownHostsFile=/dev/null",
+        )
+
+    def get_credentials(self):
+        return (
+            self.config.get("control_user", "ubuntu"),
+            self.config.get("control_host"),
+        )
+
     def _run_control(self, cmd, timeout=60):
         """
         Run a command on the control host over ssh
@@ -68,15 +89,11 @@ def _run_control(self, cmd, timeout=60):
         :returns:
             Return output from the command, if any
         """
-        control_host = self.config.get("control_host")
-        control_user = self.config.get("control_user", "ubuntu")
+        control_user, control_host = self.get_credentials()
         ssh_cmd = [
             "ssh",
-            "-o",
-            "StrictHostKeyChecking=no",
-            "-o",
-            "UserKnownHostsFile=/dev/null",
-            "{}@{}".format(control_user, control_host),
+            *self.get_ssh_options(),
+            f"{control_user}@{control_host}",
             cmd,
         ]
         try:
@@ -96,16 +113,12 @@ def _copy_to_control(self, local_file, remote_file):
         :param remote_file:
             Remote filename
         """
-        control_host = self.config.get("control_host")
-        control_user = self.config.get("control_user", "ubuntu")
+        control_user, control_host = self.get_credentials()
         ssh_cmd = [
             "scp",
-            "-o",
-            "StrictHostKeyChecking=no",
-            "-o",
-            "UserKnownHostsFile=/dev/null",
+            *self.get_ssh_options(),
             local_file,
-            "{}@{}:{}".format(control_user, control_host, remote_file),
+            f"{control_user}@{control_host}:{remote_file}",
         ]
         try:
             output = subprocess.check_output(ssh_cmd, stderr=subprocess.STDOUT)
@@ -149,17 +162,27 @@ def reboot_control_host(self):
         self._run_control("true")
 
     def provision(self):
-        # If this is not a zapper, reboot before provisioning
-        if "zapper" not in self.config.get("control_switch_local_cmd", ""):
-            self.reboot_control_host()
-        try:
-            url = self.job_data["provision_data"]["url"]
-        except KeyError:
-            raise ProvisioningError(
-                'You must specify a "url" value in '
-                'the "provision_data" section of '
-                "your job_data"
-            )
+        # determine where to get the provisioning image from
+        source = self.job_data["provision_data"].get("url")
+        if source is None:
+            image_name = self.job_data["provision_data"].get("use_attachment")
+            if image_name is None:
+                raise ProvisioningError(
+                    'In the "provision_data" section of your job_data '
+                    'you must provide a value for "url" to specify '
+                    "where to download an image from or a value for "
+                    '"use_attachment" to specify which attachment to use '
+                    "as an image"
+                )
+            source = Path.cwd() / ATTACHMENTS_DIR / "provision" / image_name
+            if not source.exists():
+                raise ProvisioningError(
+                    'In the "provision_data" section of your job_data '
+                    'you have provided a value for "use_attachment" but '
+                    f"that attachment doesn't exist: {image_name}"
+                )
+
+        # determine where to write the provisioning image
         if "media" not in self.job_data["provision_data"]:
             media = None
             self.test_device = self.config["test_device"]
@@ -188,54 +211,103 @@ def provision(self):
                     "your job_data must be either "
                     "'sd' or 'usb'"
                 )
+
+        # If this is not a zapper, reboot before provisioning
+        if "zapper" not in self.config.get("control_switch_local_cmd", ""):
+            self.reboot_control_host()
         time.sleep(5)
-        logger.info(f"Flashing Test image on {self.test_device}")
+
+        self.flash_test_image(source)
+
+        if self.job_data["provision_data"].get("create_user", True):
+            with self.remote_mount():
+                image_type = self.get_image_type()
+                logger.info("Image type detected: {}".format(image_type))
+                logger.info("Creating Test User")
+                self.create_user(image_type)
+        else:
+            logger.info("Skipping test user creation (create_user=False)")
+        self.run_post_provision_script()
+        logger.info("Booting Test Image")
+        if media == "sd":
+            cmd = "zapper sdwire set DUT"
+        elif media == "usb":
+            cmd = "zapper typecmux set DUT"
+        else:
+            cmd = self.config.get("control_switch_device_cmd", "stm -dut")
+        self._run_control(cmd)
+        self.hardreset()
+        self.check_test_image_booted()
+
+    def download(self, url: str, local: Path, timeout: Optional[int]):
+        with requests.Session() as session:
+            response = session.get(url, stream=True, timeout=timeout)
+            response.raise_for_status()
+            with open(local, "wb") as file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:  # filter out keep-alive new chunks
+                        file.write(chunk)
+
+    def transfer_test_image(self, local: Path, timeout: Optional[int] = None):
+        ssh_options = " ".join(self.get_ssh_options())
+        control_user, control_host = self.get_credentials()
+        cmd = (
+            "set -o pipefail; "
+            f"cat {local} | "
+            f"ssh {ssh_options} {control_user}@{control_host} "
+            f'"zstdcat | sudo dd of={self.test_device} bs=16M"'
+        )
         try:
-            self.flash_test_image(url)
-            if self.job_data["provision_data"].get("create_user", True):
-                with self.remote_mount():
-                    image_type = self.get_image_type()
-                    logger.info("Image type detected: {}".format(image_type))
-                    logger.info("Creating Test User")
-                    self.create_user(image_type)
-            else:
-                logger.info("Skipping test user creation (create_user=False)")
-            self.run_post_provision_script()
-            logger.info("Booting Test Image")
-            if media == "sd":
-                cmd = "zapper sdwire set DUT"
-            elif media == "usb":
-                cmd = "zapper typecmux set DUT"
-            else:
-                cmd = self.config.get("control_switch_device_cmd", "stm -dut")
-            self._run_control(cmd)
-            self.hardreset()
-            self.check_test_image_booted()
-        except Exception:
-            raise
+            subprocess.run(
+                cmd,
+                capture_output=True,
+                check=True,
+                text=True,
+                shell=True,
+                executable="/bin/bash",
+                timeout=timeout,
+            )
+        except subprocess.CalledProcessError as error:
+            raise ProvisioningError(
+                f"Error while piping the test image to {self.test_device} "
+                f"through {control_user}@{control_host}: {error}"
+            ) from error
+        except subprocess.TimeoutExpired as error:
+            raise ProvisioningError(
+                f"Timeout while piping the test image to {self.test_device} "
+                f"through {control_user}@{control_host} "
+                f"using a timeout of {timeout}: {error}"
+            ) from error
 
-    def flash_test_image(self, url):
+    def flash_test_image(self, source: Union[str, Path]):
         """
-        Flash the image at :image_url to the sd card.
+        Flash the image at :source to the sd card.
 
-        :param url:
-            URL to download the image from
+        :param source:
+            URL or Path to retrieve the image from
         :raises ProvisioningError:
             If the command times out or anything else fails.
         """
         # First unmount, just in case
         self.unmount_writable_partition()
 
-        cmd = (
-            f"(set -o pipefail; curl -sf {shlex.quote(url)} | zstdcat| "
-            f"sudo dd of={self.test_device} bs=16M)"
-        )
-        logger.info("Running: %s", cmd)
-        try:
-            # XXX: I hope 30 min is enough? but maybe not!
-            self._run_control(cmd, timeout=1800)
-        except Exception:
-            raise ProvisioningError("timeout reached while flashing image!")
+        if isinstance(source, Path):
+            # the source is an existing attachment
+            logger.info(
+                f"Flashing Test image {source.name} on {self.test_device}"
+            )
+            self.transfer_test_image(local=source, timeout=1200)
+        else:
+            # the source is a URL
+            with tempfile.NamedTemporaryFile(delete=True) as source_file:
+                logger.info(f"Downloading test image from {source}")
+                self.download(source, local=source_file.name, timeout=1200)
+                url_name = Path(urllib.parse.urlparse(source).path).name
+                logger.info(
+                    f"Flashing Test image {url_name} on {self.test_device}"
+                )
+                self.transfer_test_image(local=source_file.name, timeout=1200)
+
         try:
             self._run_control("sync")
         except Exception:
@@ -247,10 +319,10 @@ def flash_test_image(self, url):
                 "sudo hdparm -z {}".format(self.test_device),
                 timeout=30,
             )
-        except Exception:
+        except Exception as error:
             raise ProvisioningError(
                 "Unable to run hdparm to rescan " "partitions"
-            )
+            ) from error
 
     def _get_part_labels(self):
         lsblk_data = self._run_control(
@@ -528,17 +600,14 @@ def check_test_image_booted(self):
                             return True
 
                     continue
-
+                device_ip = self.config["device_ip"]
                 cmd = [
                     "sshpass",
                     "-p",
                     test_password,
                     "ssh-copy-id",
-                    "-o",
-                    "StrictHostKeyChecking=no",
-                    "-o",
-                    "UserKnownHostsFile=/dev/null",
-                    "{}@{}".format(test_username, self.config["device_ip"]),
+                    *self.get_ssh_options(),
+                    f"{test_username}@{device_ip}",
                 ]
                 subprocess.check_output(
                     cmd, stderr=subprocess.STDOUT, timeout=60

docs/.wordlist.txt

@@ -95,6 +95,7 @@ TPM
 txt
 Ubuntu
 ubuntu
+url
 UI
 URI
 USB

docs/reference/test-phases.rst

@@ -222,7 +222,8 @@ In the `provisioning`, `firmware_update` and `test` phases, it is also possible
     job_queue: example-queue
     provision_data:
       attachments:
-        - local: "ubuntu-22.04.4-preinstalled-desktop-arm64+raspi.img.xz"
+        - local: ubuntu-22.04.4-preinstalled-desktop-arm64+raspi.img.xz
+      use_attachment: ubuntu-22.04.4-preinstalled-desktop-arm64+raspi.img.xz
     test_data:
       attachments:
         - local: "config.json"
@@ -264,6 +265,12 @@ In the `provisioning`, `firmware_update` and `test` phases, it is also possible
             │   └── ubuntu-logo.png
             └── script.sh
 
+In this example, there is no `url` field under the `provision_data` to specify where to download the provisioning image from.
+Instead, there is a `use_attachment` field that indicates which attachment should be used as a provisioning image.
+The presence of *either* `url` or `use_attachment` is required.
+
+At the moment, only the `muxpi` device connector supports provisioning using an attached image.
+
 Output 
 ------------