Skip to content

Commit a63a4ec

Browse files
committed
Move WMI functions to cim package
1 parent 73e19af commit a63a4ec

File tree

8 files changed

+453
-223
lines changed

8 files changed

+453
-223
lines changed

pkg/cim/disk.go

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,17 @@ const (
2323
// GPTPartitionTypeMicrosoftReserved is the GUID for Microsoft Reserved Partition (MSR)
2424
// Reserved by Windows for system use
2525
GPTPartitionTypeMicrosoftReserved = "{e3c9e316-0b5c-4db8-817d-f92df00215ae}"
26+
27+
// ErrorCodeCreatePartitionAccessPathAlreadyInUse is the error code (42002) returned when the driver letter failed to assign after partition created
28+
ErrorCodeCreatePartitionAccessPathAlreadyInUse = 42002
29+
)
30+
31+
var (
32+
DiskSelectorListForDiskNumberAndLocation = []string{"Number", "Location"}
33+
DiskSelectorListForPartitionStyle = []string{"PartitionStyle"}
34+
DiskSelectorListForPathAndSerialNumber = []string{"Path", "SerialNumber"}
35+
DiskSelectorListForIsOffline = []string{"IsOffline"}
36+
DiskSelectorListForSize = []string{"Size"}
2637
)
2738

2839
// QueryDiskByNumber retrieves disk information for a specific disk identified by its number.
@@ -77,6 +88,60 @@ func ListDisks(selectorList []string) ([]*storage.MSFT_Disk, error) {
7788
return disks, nil
7889
}
7990

91+
// InitializeDisk initializes a RAW disk with a particular partition style.
92+
//
93+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/initialize-msft-disk
94+
// for the WMI method definition.
95+
func InitializeDisk(disk *storage.MSFT_Disk, partitionStyle int) (int, error) {
96+
result, err := disk.InvokeMethodWithReturn("Initialize", int32(partitionStyle))
97+
return int(result), err
98+
}
99+
100+
// RefreshDisk Refreshes the cached disk layout information.
101+
//
102+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-refresh
103+
// for the WMI method definition.
104+
func RefreshDisk(disk *storage.MSFT_Disk) (int, string, error) {
105+
var status string
106+
result, err := disk.InvokeMethodWithReturn("Refresh", &status)
107+
return int(result), status, err
108+
}
109+
110+
// CreatePartition creates a partition on a disk.
111+
//
112+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/createpartition-msft-disk
113+
// for the WMI method definition.
114+
func CreatePartition(disk *storage.MSFT_Disk, params ...interface{}) (int, error) {
115+
result, err := disk.InvokeMethodWithReturn("CreatePartition", params...)
116+
return int(result), err
117+
}
118+
119+
// SetDiskState takes a disk online or offline.
120+
//
121+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-online and
122+
// https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-disk-offline
123+
// for the WMI method definition.
124+
func SetDiskState(disk *storage.MSFT_Disk, online bool) (int, string, error) {
125+
method := "Offline"
126+
if online {
127+
method = "Online"
128+
}
129+
130+
var status string
131+
result, err := disk.InvokeMethodWithReturn(method, &status)
132+
return int(result), status, err
133+
}
134+
135+
// RescanDisks rescans all changes by updating the internal cache of software objects (that is, Disks, Partitions, Volumes)
136+
// for the storage setting.
137+
//
138+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-storagesetting-updatehoststoragecache
139+
// for the WMI method definition.
140+
func RescanDisks() (int, error) {
141+
result, _, err := InvokeCimMethod(WMINamespaceStorage, "MSFT_StorageSetting", "UpdateHostStorageCache", nil)
142+
return result, err
143+
}
144+
80145
// GetDiskNumber returns the number of a disk.
81146
func GetDiskNumber(disk *storage.MSFT_Disk) (uint32, error) {
82147
number, err := disk.GetProperty("Number")
@@ -85,3 +150,41 @@ func GetDiskNumber(disk *storage.MSFT_Disk) (uint32, error) {
85150
}
86151
return uint32(number.(int32)), err
87152
}
153+
154+
// GetDiskLocation returns the location of a disk.
155+
func GetDiskLocation(disk *storage.MSFT_Disk) (string, error) {
156+
return disk.GetPropertyLocation()
157+
}
158+
159+
// GetDiskPartitionStyle returns the partition style of a disk.
160+
func GetDiskPartitionStyle(disk *storage.MSFT_Disk) (int32, error) {
161+
retValue, err := disk.GetProperty("PartitionStyle")
162+
if err != nil {
163+
return 0, err
164+
}
165+
return retValue.(int32), err
166+
}
167+
168+
// IsDiskOffline returns whether a disk is offline.
169+
func IsDiskOffline(disk *storage.MSFT_Disk) (bool, error) {
170+
return disk.GetPropertyIsOffline()
171+
}
172+
173+
// GetDiskSize returns the size of a disk.
174+
func GetDiskSize(disk *storage.MSFT_Disk) (int64, error) {
175+
sz, err := disk.GetProperty("Size")
176+
if err != nil {
177+
return -1, err
178+
}
179+
return strconv.ParseInt(sz.(string), 10, 64)
180+
}
181+
182+
// GetDiskPath returns the path of a disk.
183+
func GetDiskPath(disk *storage.MSFT_Disk) (string, error) {
184+
return disk.GetPropertyPath()
185+
}
186+
187+
// GetDiskSerialNumber returns the serial number of a disk.
188+
func GetDiskSerialNumber(disk *storage.MSFT_Disk) (string, error) {
189+
return disk.GetPropertySerialNumber()
190+
}

pkg/cim/smb.go

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
package cim
55

66
import (
7+
"strings"
8+
79
"github.com/microsoft/wmi/pkg/base/query"
810
cim "github.com/microsoft/wmi/pkg/wmiinstance"
911
)
@@ -17,8 +19,24 @@ const (
1719
SmbMappingStatusConnecting
1820
SmbMappingStatusReconnecting
1921
SmbMappingStatusUnavailable
22+
23+
credentialDelimiter = ":"
2024
)
2125

26+
// escapeQueryParameter escapes a parameter for WMI Queries
27+
func escapeQueryParameter(s string) string {
28+
s = strings.ReplaceAll(s, "'", "''")
29+
s = strings.ReplaceAll(s, "\\", "\\\\")
30+
return s
31+
}
32+
33+
func escapeUserName(userName string) string {
34+
// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L169-L170
35+
userName = strings.ReplaceAll(userName, "\\", "\\\\")
36+
userName = strings.ReplaceAll(userName, credentialDelimiter, "\\"+credentialDelimiter)
37+
return userName
38+
}
39+
2240
// QuerySmbGlobalMappingByRemotePath retrieves the SMB global mapping from its remote path.
2341
//
2442
// The equivalent WMI query is:
@@ -28,7 +46,7 @@ const (
2846
// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
2947
// for the WMI class definition.
3048
func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, error) {
31-
smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", remotePath)
49+
smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
3250
instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
3351
if err != nil {
3452
return nil, err
@@ -37,12 +55,22 @@ func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, err
3755
return instances[0], err
3856
}
3957

40-
// RemoveSmbGlobalMappingByRemotePath removes a SMB global mapping matching to the remote path.
58+
// GetSmbGlobalMappingStatus returns the status of an SMB global mapping.
59+
func GetSmbGlobalMappingStatus(inst *cim.WmiInstance) (int32, error) {
60+
statusProp, err := inst.GetProperty("Status")
61+
if err != nil {
62+
return SmbMappingStatusUnavailable, err
63+
}
64+
65+
return statusProp.(int32), nil
66+
}
67+
68+
// RemoveSmbGlobalMappingByRemotePath removes an SMB global mapping matching to the remote path.
4169
//
4270
// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
4371
// for the WMI class definition.
4472
func RemoveSmbGlobalMappingByRemotePath(remotePath string) error {
45-
smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", remotePath)
73+
smbQuery := query.NewWmiQuery("MSFT_SmbGlobalMapping", "RemotePath", escapeQueryParameter(remotePath))
4674
instances, err := QueryInstances(WMINamespaceSmb, smbQuery)
4775
if err != nil {
4876
return err
@@ -51,3 +79,22 @@ func RemoveSmbGlobalMappingByRemotePath(remotePath string) error {
5179
_, err = instances[0].InvokeMethod("Remove", true)
5280
return err
5381
}
82+
83+
// NewSmbGlobalMapping creates a new SMB global mapping to the remote path.
84+
//
85+
// Refer to https://pkg.go.dev/github.com/microsoft/wmi/server2019/root/microsoft/windows/smb#MSFT_SmbGlobalMapping
86+
// for the WMI class definition.
87+
func NewSmbGlobalMapping(remotePath, username, password string, requirePrivacy bool) (int, error) {
88+
params := map[string]interface{}{
89+
"RemotePath": remotePath,
90+
"RequirePrivacy": requirePrivacy,
91+
}
92+
if username != "" {
93+
// refer to https://github.com/PowerShell/PowerShell/blob/9303de597da55963a6e26a8fe164d0b256ca3d4d/src/Microsoft.PowerShell.Commands.Management/cimSupport/cmdletization/cim/cimConverter.cs#L166-L178
94+
// on how SMB credential is handled in PowerShell
95+
params["Credential"] = escapeUserName(username) + credentialDelimiter + password
96+
}
97+
98+
result, _, err := InvokeCimMethod(WMINamespaceSmb, "MSFT_SmbGlobalMapping", "Create", params)
99+
return result, err
100+
}

pkg/cim/volume.go

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,23 @@ import (
77
"fmt"
88
"strconv"
99

10+
"github.com/go-ole/go-ole"
1011
"github.com/microsoft/wmi/pkg/base/query"
1112
"github.com/microsoft/wmi/pkg/errors"
1213
"github.com/microsoft/wmi/server2019/root/microsoft/windows/storage"
14+
"k8s.io/klog/v2"
15+
)
16+
17+
const (
18+
FileSystemUnknown = 0
19+
)
20+
21+
var (
22+
VolumeSelectorListForFileSystemType = []string{"FileSystemType"}
23+
VolumeSelectorListForStats = []string{"UniqueId", "SizeRemaining", "Size"}
24+
VolumeSelectorListUniqueID = []string{"UniqueId"}
25+
26+
PartitionSelectorListObjectID = []string{"ObjectId"}
1327
)
1428

1529
// QueryVolumeByUniqueID retrieves a specific volume by its unique identifier,
@@ -78,6 +92,68 @@ func ListVolumes(selectorList []string) ([]*storage.MSFT_Volume, error) {
7892
return volumes, nil
7993
}
8094

95+
// FormatVolume formats the specified volume.
96+
//
97+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/format-msft-volume
98+
// for the WMI method definition.
99+
func FormatVolume(volume *storage.MSFT_Volume, params ...interface{}) (int, error) {
100+
result, err := volume.InvokeMethodWithReturn("Format", params...)
101+
return int(result), err
102+
}
103+
104+
// FlushVolume flushes the cached data in the volume's file system to disk.
105+
//
106+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-volume-flush
107+
// for the WMI method definition.
108+
func FlushVolume(volume *storage.MSFT_Volume) (int, error) {
109+
result, err := volume.Flush()
110+
return int(result), err
111+
}
112+
113+
// GetVolumeUniqueID returns the unique ID (object ID) of a volume.
114+
func GetVolumeUniqueID(volume *storage.MSFT_Volume) (string, error) {
115+
return volume.GetPropertyUniqueId()
116+
}
117+
118+
// GetVolumeFileSystemType returns the file system type of a volume.
119+
func GetVolumeFileSystemType(volume *storage.MSFT_Volume) (int32, error) {
120+
fsType, err := volume.GetProperty("FileSystemType")
121+
if err != nil {
122+
return 0, err
123+
}
124+
return fsType.(int32), nil
125+
}
126+
127+
// GetVolumeSize returns the size of a volume.
128+
func GetVolumeSize(volume *storage.MSFT_Volume) (int64, error) {
129+
volumeSizeVal, err := volume.GetProperty("Size")
130+
if err != nil {
131+
return -1, err
132+
}
133+
134+
volumeSize, err := strconv.ParseInt(volumeSizeVal.(string), 10, 64)
135+
if err != nil {
136+
return -1, err
137+
}
138+
139+
return volumeSize, err
140+
}
141+
142+
// GetVolumeSizeRemaining returns the remaining size of a volume.
143+
func GetVolumeSizeRemaining(volume *storage.MSFT_Volume) (int64, error) {
144+
volumeSizeRemainingVal, err := volume.GetProperty("SizeRemaining")
145+
if err != nil {
146+
return -1, err
147+
}
148+
149+
volumeSizeRemaining, err := strconv.ParseInt(volumeSizeRemainingVal.(string), 10, 64)
150+
if err != nil {
151+
return -1, err
152+
}
153+
154+
return volumeSizeRemaining, err
155+
}
156+
81157
// ListPartitionsOnDisk retrieves all partitions or a partition with the specified number on a disk.
82158
//
83159
// The equivalent WMI query is:
@@ -245,3 +321,78 @@ func GetPartitionDiskNumber(part *storage.MSFT_Partition) (uint32, error) {
245321

246322
return uint32(diskNumber.(int32)), nil
247323
}
324+
325+
// SetPartitionState takes a partition online or offline.
326+
//
327+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-online and
328+
// https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-offline
329+
// for the WMI method definition.
330+
func SetPartitionState(part *storage.MSFT_Partition, online bool) (int, string, error) {
331+
method := "Offline"
332+
if online {
333+
method = "Online"
334+
}
335+
336+
var status string
337+
result, err := part.InvokeMethodWithReturn(method, &status)
338+
return int(result), status, err
339+
}
340+
341+
// GetPartitionSupportedSize retrieves the minimum and maximum sizes that the partition can be resized to using the ResizePartition method.
342+
//
343+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-getsupportedsizes
344+
// for the WMI method definition.
345+
func GetPartitionSupportedSize(part *storage.MSFT_Partition) (result int, sizeMin, sizeMax int64, status string, err error) {
346+
sizeMin = -1
347+
sizeMax = -1
348+
349+
var sizeMinVar, sizeMaxVar ole.VARIANT
350+
invokeResult, err := part.InvokeMethodWithReturn("GetSupportedSize", &sizeMinVar, &sizeMaxVar, &status)
351+
if invokeResult != 0 || err != nil {
352+
result = int(invokeResult)
353+
}
354+
klog.V(5).Infof("got sizeMin (%v) sizeMax (%v) from partition (%v), status: %s", sizeMinVar, sizeMaxVar, part, status)
355+
356+
sizeMin, err = strconv.ParseInt(sizeMinVar.ToString(), 10, 64)
357+
if err != nil {
358+
return
359+
}
360+
361+
sizeMax, err = strconv.ParseInt(sizeMaxVar.ToString(), 10, 64)
362+
return
363+
}
364+
365+
// ResizePartition resizes a partition.
366+
//
367+
// Refer to https://learn.microsoft.com/en-us/windows-hardware/drivers/storage/msft-partition-resize
368+
// for the WMI method definition.
369+
func ResizePartition(part *storage.MSFT_Partition, size int64) (int, string, error) {
370+
var status string
371+
result, err := part.InvokeMethodWithReturn("Resize", strconv.Itoa(int(size)), &status)
372+
return int(result), status, err
373+
}
374+
375+
// GetPartitionSize returns the size of a partition.
376+
func GetPartitionSize(part *storage.MSFT_Partition) (int64, error) {
377+
sizeProp, err := part.GetProperty("Size")
378+
if err != nil {
379+
return -1, err
380+
}
381+
382+
size, err := strconv.ParseInt(sizeProp.(string), 10, 64)
383+
if err != nil {
384+
return -1, err
385+
}
386+
387+
return size, err
388+
}
389+
390+
// FilterForPartitionOnDisk creates a WMI query filter to query a disk by its number.
391+
func FilterForPartitionOnDisk(diskNumber uint32) *query.WmiQueryFilter {
392+
return query.NewWmiQueryFilter("DiskNumber", strconv.Itoa(int(diskNumber)), query.Equals)
393+
}
394+
395+
// FilterForPartitionsOfTypeNormal creates a WMI query filter for all non-reserved partitions.
396+
func FilterForPartitionsOfTypeNormal() *query.WmiQueryFilter {
397+
return query.NewWmiQueryFilter("GptType", GPTPartitionTypeMicrosoftReserved, query.NotEquals)
398+
}

0 commit comments

Comments
 (0)