diff --git a/os_win/tests/unit/utils/storage/virtdisk/test_vhdutils.py b/os_win/tests/unit/utils/storage/virtdisk/test_vhdutils.py index 6bb8ede7..c10f5fb0 100644 --- a/os_win/tests/unit/utils/storage/virtdisk/test_vhdutils.py +++ b/os_win/tests/unit/utils/storage/virtdisk/test_vhdutils.py @@ -15,6 +15,7 @@ import ctypes import os from unittest import mock +import uuid import ddt import six @@ -134,6 +135,14 @@ class VHDUtilsTestCase(test_base.BaseTestCase): self._mock_close.assert_called_once_with( mock.sentinel.handle) + def test_guid_from_str(self): + buff = list(range(16)) + py_uuid = uuid.UUID(bytes=bytes(buff)) + guid = wintypes.GUID.from_str(str(py_uuid)) + guid_bytes = ctypes.cast(ctypes.byref(guid), + ctypes.POINTER(wintypes.BYTE * 16)).contents + self.assertEqual(buff, guid_bytes[:]) + @mock.patch.object(vhdutils.VHDUtils, '_get_vhd_device_id') def _test_create_vhd(self, mock_get_dev_id, new_vhd_type): create_params_struct = ( @@ -151,7 +160,8 @@ class VHDUtilsTestCase(test_base.BaseTestCase): new_vhd_type=new_vhd_type, src_path=mock.sentinel.src_path, max_internal_size=mock.sentinel.max_internal_size, - parent_path=mock.sentinel.parent_path) + parent_path=mock.sentinel.parent_path, + guid=mock.sentinel.guid) self._fake_vst_struct.assert_called_once_with( DeviceId=mock_get_dev_id.return_value, @@ -174,6 +184,11 @@ class VHDUtilsTestCase(test_base.BaseTestCase): self.assertEqual( vhdutils.VIRTUAL_DISK_DEFAULT_SECTOR_SIZE, fake_create_params.Version2.SectorSizeInBytes) + self.assertEqual( + vhdutils.wintypes.GUID.from_str.return_value, + fake_create_params.Version2.UniqueId) + vhdutils.wintypes.GUID.from_str.assert_called_once_with( + mock.sentinel.guid) self._mock_run.assert_called_once_with( vhdutils.virtdisk.CreateVirtualDisk, @@ -503,6 +518,44 @@ class VHDUtilsTestCase(test_base.BaseTestCase): **self._run_args) self._mock_close.assert_called_once_with(mock.sentinel.handle) + @mock.patch.object(vhdutils.VHDUtils, '_open') + def test_set_vhd_guid(self, mock_open): + set_vdisk_info_struct = ( + self._vdisk_struct.SET_VIRTUAL_DISK_INFO) + open_params_struct = ( + self._vdisk_struct.OPEN_VIRTUAL_DISK_PARAMETERS) + + fake_set_params = set_vdisk_info_struct.return_value + fake_open_params = open_params_struct.return_value + mock_open.return_value = mock.sentinel.handle + + self._vhdutils.set_vhd_guid(mock.sentinel.vhd_path, + mock.sentinel.guid) + + self.assertEqual(w_const.OPEN_VIRTUAL_DISK_VERSION_2, + fake_open_params.Version) + self.assertFalse(fake_open_params.Version2.GetInfoOnly) + + self._vhdutils._open.assert_called_once_with( + mock.sentinel.vhd_path, + open_flag=w_const.OPEN_VIRTUAL_DISK_FLAG_NO_PARENTS, + open_access_mask=0, + open_params=vhdutils.ctypes.byref(fake_open_params)) + vhdutils.wintypes.GUID.from_str.assert_called_once_with( + mock.sentinel.guid) + + self.assertEqual(w_const.SET_VIRTUAL_DISK_INFO_VIRTUAL_DISK_ID, + fake_set_params.Version) + self.assertEqual(vhdutils.wintypes.GUID.from_str.return_value, + fake_set_params.VirtualDiskId) + + self._mock_run.assert_called_once_with( + vhdutils.virtdisk.SetVirtualDiskInformation, + mock.sentinel.handle, + vhdutils.ctypes.byref(fake_set_params), + **self._run_args) + self._mock_close.assert_called_once_with(mock.sentinel.handle) + @mock.patch.object(vhdutils.VHDUtils, 'get_internal_vhd_size_by_file_size') @mock.patch.object(vhdutils.VHDUtils, '_resize_vhd') @mock.patch.object(vhdutils.VHDUtils, '_check_resize_needed') diff --git a/os_win/utils/storage/virtdisk/vhdutils.py b/os_win/utils/storage/virtdisk/vhdutils.py index f8b69a05..292638b3 100644 --- a/os_win/utils/storage/virtdisk/vhdutils.py +++ b/os_win/utils/storage/virtdisk/vhdutils.py @@ -134,7 +134,7 @@ class VHDUtils(object): self._win32_utils.close_handle(handle) def create_vhd(self, new_vhd_path, new_vhd_type, src_path=None, - max_internal_size=0, parent_path=None): + max_internal_size=0, parent_path=None, guid=None): new_device_id = self._get_vhd_device_id(new_vhd_path) vst = vdisk_struct.VIRTUAL_STORAGE_TYPE( @@ -152,6 +152,8 @@ class VHDUtils(object): w_const.CREATE_VHD_PARAMS_DEFAULT_BLOCK_SIZE) params.Version2.SectorSizeInBytes = ( VIRTUAL_DISK_DEFAULT_SECTOR_SIZE) + if guid: + params.Version2.UniqueId = wintypes.GUID.from_str(guid) handle = wintypes.HANDLE() create_virtual_disk_flag = CREATE_VIRTUAL_DISK_FLAGS.get( @@ -362,6 +364,27 @@ class VHDUtils(object): ctypes.byref(params), cleanup_handle=handle) + def set_vhd_guid(self, vhd_path, guid): + # VHDX parents will not be updated, regardless of the open flag. + open_params = vdisk_struct.OPEN_VIRTUAL_DISK_PARAMETERS() + open_params.Version = w_const.OPEN_VIRTUAL_DISK_VERSION_2 + open_params.Version2.GetInfoOnly = False + + handle = self._open( + vhd_path, + open_flag=w_const.OPEN_VIRTUAL_DISK_FLAG_NO_PARENTS, + open_access_mask=0, + open_params=ctypes.byref(open_params)) + + params = vdisk_struct.SET_VIRTUAL_DISK_INFO() + params.Version = w_const.SET_VIRTUAL_DISK_INFO_VIRTUAL_DISK_ID + params.VirtualDiskId = wintypes.GUID.from_str(guid) + + self._run_and_check_output(virtdisk.SetVirtualDiskInformation, + handle, + ctypes.byref(params), + cleanup_handle=handle) + def resize_vhd(self, vhd_path, new_max_size, is_file_max_size=True, validate_new_size=True): if is_file_max_size: diff --git a/os_win/utils/winapi/constants.py b/os_win/utils/winapi/constants.py index 08d8b9d0..62f87847 100644 --- a/os_win/utils/winapi/constants.py +++ b/os_win/utils/winapi/constants.py @@ -330,3 +330,5 @@ GET_VIRTUAL_DISK_INFO_PROVIDER_SUBTYPE = 7 GET_VIRTUAL_DISK_INFO_IS_LOADED = 13 SET_VIRTUAL_DISK_INFO_PARENT_PATH = 1 +SET_VIRTUAL_DISK_INFO_IDENTIFIER = 2 +SET_VIRTUAL_DISK_INFO_VIRTUAL_DISK_ID = 5 diff --git a/os_win/utils/winapi/libs/virtdisk.py b/os_win/utils/winapi/libs/virtdisk.py index c124940c..726ab704 100644 --- a/os_win/utils/winapi/libs/virtdisk.py +++ b/os_win/utils/winapi/libs/virtdisk.py @@ -180,11 +180,19 @@ class GET_VIRTUAL_DISK_INFO(ctypes.Structure): PGET_VIRTUAL_DISK_INFO = ctypes.POINTER(GET_VIRTUAL_DISK_INFO) -# Only this version is used, we avoid defining a union. +class _SET_VIRTUAL_DISK_INFO_U(ctypes.Union): + _fields_ = [ + ('ParentFilePath', wintypes.LPCWSTR), + ('UniqueIdentifier', wintypes.GUID), + ('VirtualDiskId', wintypes.GUID), + ] + + class SET_VIRTUAL_DISK_INFO(ctypes.Structure): + _anonymous_ = ['_setinfo'] _fields_ = [ ('Version', wintypes.DWORD), - ('ParentFilePath', wintypes.LPCWSTR) + ('_setinfo', _SET_VIRTUAL_DISK_INFO_U) ] diff --git a/os_win/utils/winapi/wintypes.py b/os_win/utils/winapi/wintypes.py index e1546b6e..0bd9c5b5 100644 --- a/os_win/utils/winapi/wintypes.py +++ b/os_win/utils/winapi/wintypes.py @@ -21,6 +21,7 @@ import ctypes import sys +import uuid BYTE = ctypes.c_byte WORD = ctypes.c_ushort @@ -83,6 +84,12 @@ class GUID(ctypes.Structure): ("Data4", BYTE * 8) ] + @classmethod + def from_str(cls, guid_str): + py_uuid = uuid.UUID(guid_str) + uuid_buff = (BYTE * 16)(*py_uuid.bytes) + return ctypes.cast(uuid_buff, ctypes.POINTER(GUID)).contents + class OVERLAPPED(ctypes.Structure): _fields_ = [