Allow registering replica-first-boot scripts

coriolis/constants.py

@@ -390,12 +390,11 @@
 PHASE_OSMORPHING_PRE_OS_MOUNT = "osmorphing_pre_os_mount"
 # Scripts that are executed after the OS partition is mounted (the default).
 PHASE_OSMORPHING_POST_OS_MOUNT = "osmorphing_post_os_mount"
-# We may eventually add "PHASE_REPLICA_FIRST_BOOT" for convenience, although
-# the users can already achieve this by using os-morphing scripts to schedule
-# scripts that will be executed at the next boot. This may require import
-# provider support.
+# Scripts that are executed when the replica VM starts for the first time.
+PHASE_REPLICA_FIRST_BOOT = "replica_first_boot"
 
 USER_SCRIPT_PHASES = [
     PHASE_OSMORPHING_PRE_OS_MOUNT,
     PHASE_OSMORPHING_POST_OS_MOUNT,
+    PHASE_REPLICA_FIRST_BOOT,
 ]

coriolis/osmorphing/base.py

@@ -28,6 +28,48 @@
 CLOUD_INIT_SERVICE_UNIT_NAME = "cloud-init"
 CLOUD_INIT_SERVICE_UNIT_NAME_FALLBACK = "cloud-init-main"
 
+FIRST_BOOT_SCRIPT_RUNNER = """#!/bin/bash
+function run_scripts {
+    script_dir=$1
+
+    for f in $script_dir/*.sh; do
+        if [ -x "$f" ]; then
+            echo "Invoking script: $f"
+            "$f"
+            echo "Exit code: $?"
+        fi
+    done
+}
+
+# Run Coriolis provided scripts.
+run_scripts /usr/lib/coriolis/firstboot/service
+
+# Run user provided scripts.
+run_scripts /usr/lib/coriolis/firstboot/user
+
+mkdir -p /var/lib/coriolis
+touch /var/lib/coriolis/firstboot-complete
+"""
+FIRST_BOOT_SCRIPT_RUNNER_PATH = "/usr/lib/coriolis/firstboot/run-firstboot.sh"
+FIRST_BOOT_SYSTEMD_UNIT = """
+[Unit]
+Description=Coriolis replica first-boot scripts.
+After=network-online.target
+Wants=network-online.target
+ConditionPathExists=!/var/lib/coriolis/firstboot-complete
+
+[Service]
+Type=oneshot
+ExecStart=/usr/lib/coriolis/firstboot/run-firstboot.sh
+RemainAfterExit=yes
+
+[Install]
+WantedBy=multi-user.target
+"""
+FIRST_BOOT_SYSTEMD_UNIT_NAME = "coriolis-firstboot.service"
+FIRST_BOOT_SYSTEMD_UNIT_PATH = (
+    f"/etc/systemd/system/{FIRST_BOOT_SYSTEMD_UNIT_NAME}")
+
 
 class BaseOSMorphingTools(object, with_metaclass(abc.ABCMeta)):
 
@@ -100,6 +142,15 @@ def get_packages(self):
     def run_user_script(self, user_script):
         pass
 
+    @abc.abstractmethod
+    def register_firstboot_script(
+        self,
+        script: str,
+        index: int = 0,
+        user_provided=True,
+    ):
+        pass
+
     @abc.abstractmethod
     def pre_packages_install(self, package_names):
         pass
@@ -719,3 +770,53 @@ def _setup_network_preservation(self, nics_info) -> None:
         self._add_net_udev_rules(net_ifaces_info)
 
         return
+
+    def register_firstboot_script(
+        self,
+        script: str,
+        index: int = 0,
+        user_provided=True,
+    ):
+        if len(script) == 0:
+            LOG.debug("Empty first-boot script, skipping...")
+            return
+
+        if user_provided:
+            script_dir = "/usr/lib/coriolis/firstboot/user"
+        else:
+            script_dir = "/usr/lib/coriolis/firstboot/service"
+        unique_id = str(uuid.uuid4()).split("-")[0]
+        script_path = os.path.join(script_dir, f"{index:02d}-{unique_id}.sh")
+
+        self._exec_cmd_chroot(f"mkdir -p {script_dir}")
+        self._write_file_sudo(script_path, script)
+        self._exec_cmd_chroot(f"chown root:root {script_path}")
+        self._exec_cmd_chroot(f"chmod 755 {script_path}")
+
+        # systemd unit used to launch first-boot scripts.
+        if not self._test_path(FIRST_BOOT_SYSTEMD_UNIT_PATH):
+            self._write_file_sudo(
+                FIRST_BOOT_SYSTEMD_UNIT_PATH, FIRST_BOOT_SYSTEMD_UNIT)
+            self._exec_cmd_chroot(
+                "chown root:root %s" % FIRST_BOOT_SYSTEMD_UNIT_PATH)
+            self._exec_cmd_chroot(
+                "chmod 644 %s" % FIRST_BOOT_SYSTEMD_UNIT_PATH)
+            wants_dir = "/etc/systemd/system/multi-user.target.wants"
+            self._exec_cmd_chroot("mkdir -p %s" % wants_dir)
+            self._exec_cmd_chroot(
+                "ln -sf %s %s/%s" % (
+                    FIRST_BOOT_SYSTEMD_UNIT_PATH,
+                    wants_dir,
+                    FIRST_BOOT_SYSTEMD_UNIT_NAME))
+
+        # A script that iterates over "/usr/lib/coriolis/firstboot/*.sh"
+        # scripts and runs them.
+        if not self._test_path(FIRST_BOOT_SCRIPT_RUNNER_PATH):
+            self._write_file_sudo(
+                FIRST_BOOT_SCRIPT_RUNNER_PATH, FIRST_BOOT_SCRIPT_RUNNER)
+            self._exec_cmd_chroot(
+                "chown root:root %s" % FIRST_BOOT_SCRIPT_RUNNER_PATH)
+            self._exec_cmd_chroot(
+                "chmod 755 %s" % FIRST_BOOT_SCRIPT_RUNNER_PATH)
+
+        LOG.info(f"Registered first-boot script: {script_path}")

coriolis/osmorphing/manager.py

@@ -307,3 +307,13 @@ def _morph_image(origin_provider, destination_provider, connection_info,
 
     LOG.info("Post packages install")
     import_os_morphing_tools.post_packages_install(packages_add)
+
+    first_boot_user_scripts = [
+        script["payload"] for script in user_scripts
+        if script["phase"] == constants.PHASE_REPLICA_FIRST_BOOT]
+    for script_idx, user_script in enumerate(first_boot_user_scripts):
+        event_manager.progress_update('Registering first-boot user script')
+        import_os_morphing_tools.register_firstboot_script(
+            user_script, script_idx, user_provided=True)
+    if not first_boot_user_scripts:
+        event_manager.progress_update('No first-boot user script specified')

coriolis/osmorphing/windows.py

@@ -456,14 +456,19 @@ def _get_cbslinit_scripts_dir(self, base_dir):
 
     def _write_local_script(self, base_dir, script_path, priority=50):
         scripts_dir = self._get_cbslinit_scripts_dir(base_dir)
-        script = "%s\\%d-%s" % (
+        remote_script_path = "%s\\%02d-%s" % (
             scripts_dir, priority,
             os.path.basename(script_path))
 
         with open(script_path, 'r') as fd:
             contents = fd.read()
             utils.write_winrm_file(
-                self._conn, script, contents)
+                self._conn, remote_script_path, contents)
+
+        LOG.info(
+            "Registered first-boot Coriolis script: %s -> %s",
+            script_path,
+            remote_script_path)
 
     def _write_cloudbase_init_conf(self, cloudbaseinit_base_dir,
                                    local_base_dir, com_port="COM1",
@@ -525,7 +530,7 @@ def _write_cloudbase_init_conf(self, cloudbaseinit_base_dir,
 
         self._write_local_script(
             cloudbaseinit_base_dir, disks_script,
-            priority=99)
+            priority=10)
 
     def _install_cloudbase_init(self, download_url,
                                 metadata_services=None, enabled_plugins=None,
@@ -723,3 +728,38 @@ def uninstall_packages(self, package_names):
 
     def post_packages_uninstall(self, package_names):
         pass
+
+    def register_firstboot_script(
+        self,
+        script: str,
+        index: int = 0,
+        user_provided=True,
+    ):
+        if len(script) == 0:
+            LOG.debug("Empty first-boot script, skipping...")
+            return
+
+        if user_provided:
+            # The default priority for Coriolis scripts is "50",
+            # some using below 50.
+            #
+            # The scripts are executed in alphabetical order, so the
+            # ones with a lower "priority" will be executed first.
+            #
+            # We'll bump the priority here so that user scripts will
+            # run after the Coriolis internal scripts.
+            index += 51
+
+        cbslinit_base_dir = self._get_cbslinit_base_dir()
+        script_dir = self._get_cbslinit_scripts_dir(cbslinit_base_dir)
+        unique_id = str(uuid.uuid4()).split("-")[0]
+        script_path = os.path.join(
+            script_dir, f"{index:02d}-{unique_id}.ps1")
+
+        self._conn.exec_ps_command(f"mkdir -Force {script_dir}")
+        utils.write_winrm_file(
+            self._conn,
+            script_path,
+            script)
+
+        LOG.info(f"Registered first-boot script: {script_path}")

coriolis/tests/integration/deployments/test_osmorphing.py

@@ -7,6 +7,8 @@
 installation in the target OS.
 """
 
+import os
+import re
 import uuid
 
 from coriolis.tests.integration import base as integration_base
@@ -135,3 +137,54 @@ def test_os_morphing_global_script_extended_format(self):
         # the replica OS disk was mounted.
         self.assertNotIn(self._dst_device, pre_mounts)
         self.assertIn(self._dst_device, post_mounts)
+
+    def test_os_morphing_global_script_first_boot(self):
+        payload = "mount > /boot_mounts"
+        user_scripts = {
+            'global': {
+                'linux': [
+                    {
+                        "phase": "replica_first_boot",
+                        "payload": "mount > /boot_mounts",
+                    },
+                ],
+                'windows': [
+                    {
+                        "phase": "replica_first_boot",
+                        "payload": "should-not-get-executed",
+                    },
+                ]
+            }
+        }
+        deployment_kwargs = {
+            "user_scripts": user_scripts,
+        }
+        self._execute_transfer_and_deployment(deployment_kwargs)
+
+        # TODO(lpetrut): the test import provider doesn't actually create
+        # replica instances (containers). If it did, we'd have no way to clean
+        # them up using Coriolis APIs.
+        #
+        # For this reason, we can't ensure that the first boot scripts
+        # actually get executed. We'll merely verify that those files
+        # have been injected at the expected location.
+        first_boot_script_dir = "usr/lib/coriolis/firstboot/user"
+        first_boot_scripts = test_utils.list_files_from_device(
+            self._dst_device, first_boot_script_dir)
+        if not first_boot_scripts:
+            raise AssertionError("Couldn't find first boot script dir.")
+
+        found = False
+        for file_name in first_boot_scripts:
+            if re.match(r"\d+-\w+\.sh", file_name):
+                first_boot_script_path = os.path.join(
+                    first_boot_script_dir, file_name)
+                first_boot_script = test_utils.read_file_from_device(
+                    self._dst_device,
+                    first_boot_script_path)
+                if payload == first_boot_script:
+                    found = True
+
+        if not found:
+            raise AssertionError(
+                "Couldn't find the expected first boot script.")

coriolis/tests/integration/utils.py

@@ -400,3 +400,18 @@ def read_file_from_device(device_path, rel_path):
                 return f.read()
         finally:
             _run(["umount", mount_point])
+
+
+def list_files_from_device(device_path, rel_path):
+    """Enumerates files from the filesystem of *device_path*.
+
+    Mounts the device read-only into a temporary directory, enumerates files,
+    then unmounts.
+    """
+    with tempfile.TemporaryDirectory() as mount_point:
+        _run(["mount", "-o", "ro", device_path, mount_point])
+
+        try:
+            return os.listdir(os.path.join(mount_point, rel_path))
+        finally:
+            _run(["umount", mount_point])

coriolis/tests/osmorphing/test_manager.py

@@ -29,6 +29,10 @@ def setUp(self):
                 "phase": constants.PHASE_OSMORPHING_POST_OS_MOUNT,
                 "payload": "post-os-mount-script",
             },
+            {
+                "phase": constants.PHASE_REPLICA_FIRST_BOOT,
+                "payload": "fist-boot-script",
+            },
         ]
 
         manager.CONF.proxy.url = "http://127.0.0.1:8080"
@@ -203,6 +207,14 @@ def install_packages(self, packages_add):
         def uninstall_packages(self, packages_remove):
             pass
 
+        def register_firstboot_script(
+            self,
+            script: str,
+            index: int = 0,
+            user_provided=True,
+        ):
+            pass
+
     @mock.patch.object(manager.osmount_factory, 'get_os_mount_tools')
     @mock.patch.object(manager.events, 'EventManager')
     @mock.patch.object(manager, 'run_os_detect')

coriolis/tests/osmorphing/test_windows.py

@@ -521,7 +521,7 @@ def test__write_cloudbase_init_conf(
             self.conn, conf_file_path, conf_content)
 
         mock_write_local_script.assert_called_once_with(
-            'C:\\Cloudbase-Init', mocked_full_path, priority=99)
+            'C:\\Cloudbase-Init', mocked_full_path, priority=10)
 
     @mock.patch.object(windows.utils, 'write_winrm_file')
     @mock.patch.object(windows.BaseWindowsMorphingTools, '_write_local_script')
@@ -946,3 +946,21 @@ def test_set_net_config_with_dhcp(
         mock_unload_registry_hive.assert_not_called()
 
         self.assertIsNone(result)
+
+    @mock.patch.object(windows.utils, 'write_winrm_file')
+    @mock.patch("uuid.uuid4")
+    def test_register_firstboot_script(self, mock_uuid, mock_write_winrm_file):
+        mock_uuid.return_value = "37c27abd-85ff-4cb8-8d31-4e7067e145ab"
+        mock_script = "mock-script"
+
+        self.morphing_tools.register_firstboot_script(
+            mock_script,
+            index=10,
+            user_provided=True)
+
+        self.morphing_tools._conn.exec_ps_command.assert_called_once_with(
+            "mkdir -Force C:\\Cloudbase-Init\\LocalScripts")
+        mock_write_winrm_file.assert_called_once_with(
+            self.morphing_tools._conn,
+            "C:\\Cloudbase-Init\\LocalScripts/61-37c27abd.ps1",
+            mock_script)