Stage virtio balloon service for instalation

coriolis/osmorphing/windows.py

@@ -85,6 +85,9 @@ QEMU_GUEST_AGENT_INSTALL_FROM_DIR_SCRIPT_FORMAT = (
 
 VIRTIO_WIN_ISO_PATH = "c:\\virtio-win.iso"
 QEMU_GA_MSI_NAME = "qemu-ga.msi"
+BALLOON_SERVICE_GUEST_DIR = "Program Files\\Balloon"
+BALLOON_SERVICE_INSTALL_SCRIPT_FORMAT = (
+    '& "%(balloon_dir)s\\blnsvr.exe" -i')
 
 INTERFACES_PATH_FORMAT = (
     "HKLM:\\%s\\ControlSet001\\Services\\Tcpip\\Parameters\\Interfaces")
@@ -829,9 +832,44 @@ class BaseWindowsMorphingTools(base.BaseOSMorphingTools):
             ]
 
             self._install_dism_drivers(driver_paths)
+            self._stage_balloon_service(virtio_drive, virtio_dir, arch)
         finally:
             self._dismount_disk_image(VIRTIO_WIN_ISO_PATH)
 
+    def _stage_balloon_service(self, virtio_drive, virtio_dir, arch):
+        """Copy blnsvr.exe into the guest and register a first-boot script.
+
+        Offline DISM driver injection only installs the balloon kernel
+        driver. The separate BalloonService (blnsvr.exe) is required for
+        the driver to report guest memory statistics to KVM/Proxmox hosts.
+        """
+        balloon_source = "%s:\\Balloon\\%s\\%s" % (
+            virtio_drive, virtio_dir, arch)
+        blnsvr_source = "%s\\blnsvr.exe" % balloon_source
+        if not self._conn.test_path(blnsvr_source):
+            LOG.warning(
+                "VirtIO Balloon user-mode service (blnsvr.exe) was not "
+                "found at '%s'. Skipping BalloonService setup.",
+                blnsvr_source)
+            return
+
+        balloon_dest = "%s%s" % (self._os_root_dir, BALLOON_SERVICE_GUEST_DIR)
+        self._event_manager.progress_update(
+            "Staging VirtIO BalloonService (blnsvr.exe)")
+        self._conn.exec_ps_command(
+            "New-Item -ItemType Directory -Force -Path '%s' | Out-Null"
+            % balloon_dest)
+        self._conn.exec_ps_command(
+            "Copy-Item -Path '%s\\*' -Destination '%s' -Recurse -Force"
+            % (balloon_source, balloon_dest))
+
+        guest_balloon_dir = "C:\\%s" % BALLOON_SERVICE_GUEST_DIR
+        local_script = BALLOON_SERVICE_INSTALL_SCRIPT_FORMAT % {
+            "balloon_dir": guest_balloon_dir}
+        self.register_firstboot_script(
+            local_script, user_provided=False,
+            script_filename="coriolis_balloon_service_install.ps1")
+
     def _install_dism_drivers(self, driver_paths):
         """Installs the drivers located at the given paths into the offline
         Windows image via DISM.

coriolis/tests/osmorphing/test_windows.py

@@ -1026,6 +1026,8 @@ class BaseWindowsMorphingToolsTestCase(test_base.CoriolisBaseTestCase):
         }
     )
     @ddt.unpack
+    @mock.patch.object(
+        windows.BaseWindowsMorphingTools, "_stage_balloon_service")
     @mock.patch.object(windows.BaseWindowsMorphingTools, "_mount_disk_image")
     @mock.patch.object(windows.BaseWindowsMorphingTools, "_get_sid")
     @mock.patch.object(windows.BaseWindowsMorphingTools, "_grant_permissions")
@@ -1041,6 +1043,7 @@ class BaseWindowsMorphingToolsTestCase(test_base.CoriolisBaseTestCase):
         mock_grant_permissions,
         mock_get_sid,
         mock_mount_disk_image,
+        mock_stage_balloon_service,
         version_number=None,
         edition_id=None,
         exp_virtio_dir=None,
@@ -1085,6 +1088,43 @@ class BaseWindowsMorphingToolsTestCase(test_base.CoriolisBaseTestCase):
             for driver in drivers]
         mock_add_dism_driver.assert_has_calls(
             [mock.call(path) for path in exp_driver_paths])
+        mock_stage_balloon_service.assert_called_once_with(
+            fake_image_mountpoint, exp_virtio_dir, "amd64")
+
+    @mock.patch.object(
+        windows.BaseWindowsMorphingTools, "register_firstboot_script")
+    def test_stage_balloon_service(self, mock_register_firstboot_script):
+        self.conn.test_path.return_value = True
+
+        self.morphing_tools._stage_balloon_service("e", "w10", "amd64")
+
+        self.conn.test_path.assert_called_once_with(
+            "e:\\Balloon\\w10\\amd64\\blnsvr.exe")
+        expected_copy_calls = [
+            mock.call(
+                "New-Item -ItemType Directory -Force -Path "
+                "'C:\\Program Files\\Balloon' | Out-Null"),
+            mock.call(
+                "Copy-Item -Path 'e:\\Balloon\\w10\\amd64\\*' "
+                "-Destination 'C:\\Program Files\\Balloon' "
+                "-Recurse -Force"),
+        ]
+        self.conn.exec_ps_command.assert_has_calls(expected_copy_calls)
+        mock_register_firstboot_script.assert_called_once_with(
+            '& "C:\\Program Files\\Balloon\\blnsvr.exe" -i',
+            user_provided=False,
+            script_filename="coriolis_balloon_service_install.ps1")
+
+    @mock.patch.object(
+        windows.BaseWindowsMorphingTools, "register_firstboot_script")
+    def test_stage_balloon_service_missing_blnsvr(
+            self, mock_register_firstboot_script):
+        self.conn.test_path.return_value = False
+
+        self.morphing_tools._stage_balloon_service("e", "w10", "amd64")
+
+        self.conn.exec_ps_command.assert_not_called()
+        mock_register_firstboot_script.assert_not_called()
 
     @ddt.data(
         {