Move WMI related functions to cim package

docs/IMPLEMENTATION.md

@@ -121,6 +121,20 @@ func CallMethod(disp *ole.IDispatch, name string, params ...interface{}) (result
 }
 ```
 
+### Association
+
+Association can be used to retrieve all instances that are associated with
+a particular source instance.
+
+There are a few Association classes in WMI.
+
+For example, association class [MSFT_PartitionToVolume](https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partitiontovolume)
+can be used to retrieve a volume (`MSFT_Volume`) from a partition (`MSFT_Partition`), and vice versa.
+
+```go
+collection, err := part.GetAssociated("MSFT_PartitionToVolume", "MSFT_Volume", "Volume", "Partition")
+```
+
 <a name="debug-powershell"></a>
 ## Debug with PowerShell
 
@@ -181,6 +195,13 @@ PS C:\Users\Administrator> $vol.FileSystem
 NTFS
 ```
 
+### Association
+
+```powershell
+PS C:\Users\Administrator> $partition = (Get-CimInstance -Namespace root\Microsoft\Windows\Storage -ClassName MSFT_Partition -Filter "DiskNumber = 0")[0]
+PS C:\Users\Administrator> Get-CimAssociatedInstance -InputObject $partition -Association MSFT_PartitionToVolume
+```
+
 ### Call Class Method
 
 You may get Class Methods for a single CIM class using `$class.CimClassMethods`.

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 (
+	DiskSelectorListForDiskNumberAndLocation = []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.
@@ -77,6 +88,60 @@ 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
+}
+
+// CreatePartition 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 CreatePartition(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")
@@ -85,3 +150,41 @@ func GetDiskNumber(disk *storage.MSFT_Disk) (uint32, error) {
 	}
 	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/smb.go

@@ -4,6 +4,8 @@
 package cim
 
 import (
+	"strings"
+
 	"github.com/microsoft/wmi/pkg/base/query"
 	cim "github.com/microsoft/wmi/pkg/wmiinstance"
 )
@@ -17,8 +19,24 @@ const (
 	SmbMappingStatusConnecting
 	SmbMappingStatusReconnecting
 	SmbMappingStatusUnavailable
+
+	credentialDelimiter = ":"
 )
 
+// escapeQueryParameter escapes a parameter for WMI Queries
+func escapeQueryParameter(s string) string {
+	s = strings.ReplaceAll(s, "'", "''")
+	s = strings.ReplaceAll(s, "\\", "\\\\")
+	return s
+}
+
+func escapeUserName(userName string) string {
+	// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L169-L170
+	userName = strings.ReplaceAll(userName, "\\", "\\\\")
+	userName = strings.ReplaceAll(userName, credentialDelimiter, "\\"+credentialDelimiter)
+	return userName
+}
+
 // QuerySmbGlobalMappingByRemotePath retrieves the SMB global mapping from its remote path.
 //
 // The equivalent WMI query is:
@@ -28,7 +46,7 @@ const (
 // Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
 // for the WMI class definition.
 func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, error) {
-	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", remotePath)
+	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
 	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
 	if err != nil {
 		return nil, err
@@ -37,12 +55,22 @@ func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, err
 	return instances[0], err
 }
 
-// RemoveSmbGlobalMappingByRemotePath removes a SMB global mapping matching to the remote path.
+// GetSmbGlobalMappingStatus returns the status of an SMB global mapping.
+func GetSmbGlobalMappingStatus(inst *cim.WmiInstance) (int32, error) {
+	statusProp, err := inst.GetProperty("Status")
+	if err != nil {
+		return SmbMappingStatusUnavailable, err
+	}
+
+	return statusProp.(int32), nil
+}
+
+// RemoveSmbGlobalMappingByRemotePath removes an SMB global mapping matching to the remote path.
 //
 // Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
 // for the WMI class definition.
 func RemoveSmbGlobalMappingByRemotePath(remotePath string) error {
-	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", remotePath)
+	smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
 	instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
 	if err != nil {
 		return err
@@ -51,3 +79,22 @@ func RemoveSmbGlobalMappingByRemotePath(remotePath string) error {
 	_, err = instances[0].InvokeMethod("Remove", true)
 	return err
 }
+
+// NewSmbGlobalMapping creates a new SMB global mapping to the remote path.
+//
+// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
+// for the WMI class definition.
+func NewSmbGlobalMapping(remotePath, username, password string, requirePrivacy bool) (int, error) {
+	params := map[string]interface{}{
+		"RemotePath":     remotePath,
+		"RequirePrivacy": requirePrivacy,
+	}
+	if username != "" {
+		// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L166-L178
+		// on how SMB credential is handled in PowerShell
+		params["Credential"] = escapeUserName(username) + credentialDelimiter + password
+	}
+
+	result, _, err := InvokeCimMethod(WMINamespaceSmb, "MSFT_SmbGlobalMapping", "Create", params)
+	return result, err
+}