Move WMI disk functions to cim package

Author: laozc

pkg/cim/disk.go

@@ -23,6 +23,17 @@ const (
 	// GPTPartitionTypeMicrosoftReserved is the GUID for Microsoft Reserved Partition (MSR)
 	// Reserved by Windows for system use
 	GPTPartitionTypeMicrosoftReserved = "{e3c9e316-0b5c-4db8-817d-f92df00215ae}"
+
+	// ErrorCodeCreatePartitionAccessPathAlreadyInUse is the error code (42002) returned when the driver letter failed to assign after partition created
+	ErrorCodeCreatePartitionAccessPathAlreadyInUse = 42002
+)
+
+var (
+	DiskSelectorListForNumberAndLocation   = []string{"Number", "Location"}
+	DiskSelectorListForPartitionStyle      = []string{"PartitionStyle"}
+	DiskSelectorListForPathAndSerialNumber = []string{"Path", "SerialNumber"}
+	DiskSelectorListForIsOffline           = []string{"IsOffline"}
+	DiskSelectorListForSize                = []string{"Size"}
 )
 
 // QueryDiskByNumber retrieves disk information for a specific disk identified by its number.
@@ -76,3 +87,106 @@ func ListDisks(selectorList []string) ([]*storage.MSFT_Disk, error) {
 
 	return disks, nil
 }
+
+// InitializeDisk initializes a RAW disk with a particular partition style.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/initialize-msft-disk
+// for the WMI method definition.
+func InitializeDisk(disk *storage.MSFT_Disk, partitionStyle int) (int, error) {
+	result, err := disk.InvokeMethodWithReturn("Initialize", int32(partitionStyle))
+	return int(result), err
+}
+
+// RefreshDisk Refreshes the cached disk layout information.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-refresh
+// for the WMI method definition.
+func RefreshDisk(disk *storage.MSFT_Disk) (int, string, error) {
+	var status string
+	result, err := disk.InvokeMethodWithReturn("Refresh", &status)
+	return int(result), status, err
+}
+
+// CreatePartitionOnDisk creates a partition on a disk.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/createpartition-msft-disk
+// for the WMI method definition.
+func CreatePartitionOnDisk(disk *storage.MSFT_Disk, params ...interface{}) (int, error) {
+	result, err := disk.InvokeMethodWithReturn(
+		"CreatePartition",
+		params...)
+	return int(result), err
+}
+
+// SetDiskState takes a disk online or offline.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-online and
+// https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-offline
+// for the WMI method definition.
+func SetDiskState(disk *storage.MSFT_Disk, online bool) (int, string, error) {
+	method := "Offline"
+	if online {
+		method = "Online"
+	}
+
+	var status string
+	result, err := disk.InvokeMethodWithReturn(method, &status)
+	return int(result), status, err
+}
+
+// RescanDisks rescans all changes by updating the internal cache of software objects (that is, Disks, Partitions, Volumes)
+// for the storage setting.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-storagesetting-updatehoststoragecache
+// for the WMI method definition.
+func RescanDisks() (int, error) {
+	result, _, err := InvokeCimMethod(WMINamespaceStorage, "MSFT_StorageSetting", "UpdateHostStorageCache", nil)
+	return result, err
+}
+
+// GetDiskNumber returns the number of a disk.
+func GetDiskNumber(disk *storage.MSFT_Disk) (uint32, error) {
+	number, err := disk.GetProperty("Number")
+	if err != nil {
+		return 0, err
+	}
+	return uint32(number.(int32)), err
+}
+
+// GetDiskLocation returns the location of a disk.
+func GetDiskLocation(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertyLocation()
+}
+
+// GetDiskPartitionStyle returns the partition style of a disk.
+func GetDiskPartitionStyle(disk *storage.MSFT_Disk) (int32, error) {
+	retValue, err := disk.GetProperty("PartitionStyle")
+	if err != nil {
+		return 0, err
+	}
+	return retValue.(int32), err
+}
+
+// IsDiskOffline returns whether a disk is offline.
+func IsDiskOffline(disk *storage.MSFT_Disk) (bool, error) {
+	return disk.GetPropertyIsOffline()
+}
+
+// GetDiskSize returns the size of a disk.
+func GetDiskSize(disk *storage.MSFT_Disk) (int64, error) {
+	sz, err := disk.GetProperty("Size")
+	if err != nil {
+		return -1, err
+	}
+	return strconv.ParseInt(sz.(string), 10, 64)
+}
+
+// GetDiskPath returns the path of a disk.
+func GetDiskPath(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertyPath()
+}
+
+// GetDiskSerialNumber returns the serial number of a disk.
+func GetDiskSerialNumber(disk *storage.MSFT_Disk) (string, error) {
+	return disk.GetPropertySerialNumber()
+}

pkg/cim/volume.go

@@ -245,3 +245,27 @@ func GetPartitionDiskNumber(part *storage.MSFT_Partition) (uint32, error) {
 
 	return uint32(diskNumber.(int32)), nil
 }
+
+// SetPartitionState takes a partition online or offline.
+//
+// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-online and
+// https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-offline
+// for the WMI method definition.
+func SetPartitionState(part *storage.MSFT_Partition, online bool) (int, string, error) {
+	method := "Offline"
+	if online {
+		method = "Online"
+	}
+
+	var status string
+	result, err := part.InvokeMethodWithReturn(method, &status)
+	return int(result), status, err
+}
+
+func FilterForPartitionOnDisk(diskNumber uint32) *query.WmiQueryFilter {
+	return query.NewWmiQueryFilter("DiskNumber", strconv.Itoa(int(diskNumber)), query.Equals)
+}
+
+func FilterForPartitionsOfTypeNormal() *query.WmiQueryFilter {
+	return query.NewWmiQueryFilter("GptType", GPTPartitionTypeMicrosoftReserved, query.NotEquals)
+}