Move WMI related functions to cim package

Author: laozc

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
@@ -42,7 +60,7 @@ func QuerySmbGlobalMappingByRemotePath(remotePath string) (*cim.WmiInstance, err
 // 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 +69,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
+}

pkg/os/smb/api.go

@@ -10,10 +10,6 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-const (
-	credentialDelimiter = ":"
-)
-
 type API interface {
 	IsSmbMapped(remotePath string) (bool, error)
 	NewSmbLink(remotePath, localPath string) error
@@ -33,17 +29,6 @@ func New(requirePrivacy bool) *SmbAPI {
 	}
 }
 
-func remotePathForQuery(remotePath string) string {
-	return strings.ReplaceAll(remotePath, "\\", "\\\\")
-}
-
-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
-	escaped := strings.ReplaceAll(userName, "\\", "\\\\")
-	escaped = strings.ReplaceAll(escaped, credentialDelimiter, "\\"+credentialDelimiter)
-	return escaped
-}
-
 func createSymlink(link, target string, isDir bool) error {
 	linkPtr, err := syscall.UTF16PtrFromString(link)
 	if err != nil {
@@ -68,7 +53,7 @@ func createSymlink(link, target string, isDir bool) error {
 }
 
 func (*SmbAPI) IsSmbMapped(remotePath string) (bool, error) {
-	inst, err := cim.QuerySmbGlobalMappingByRemotePath(remotePathForQuery(remotePath))
+	inst, err := cim.QuerySmbGlobalMappingByRemotePath(remotePath)
 	if err != nil {
 		return false, cim.IgnoreNotFound(err)
 	}
@@ -106,17 +91,7 @@ func (*SmbAPI) NewSmbLink(remotePath, localPath string) error {
 }
 
 func (api *SmbAPI) NewSmbGlobalMapping(remotePath, username, password string) error {
-	params := map[string]interface{}{
-		"RemotePath":     remotePath,
-		"RequirePrivacy": api.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 := cim.InvokeCimMethod(cim.WMINamespaceSmb, "MSFT_SmbGlobalMapping", "Create", params)
+	result, err := cim.NewSmbGlobalMapping(remotePath, username, password, api.RequirePrivacy)
 	if err != nil {
 		return fmt.Errorf("NewSmbGlobalMapping failed. result: %d, err: %v", result, err)
 	}
@@ -125,7 +100,7 @@ func (api *SmbAPI) NewSmbGlobalMapping(remotePath, username, password string) er
 }
 
 func (*SmbAPI) RemoveSmbGlobalMapping(remotePath string) error {
-	err := cim.RemoveSmbGlobalMappingByRemotePath(remotePathForQuery(remotePath))
+	err := cim.RemoveSmbGlobalMappingByRemotePath(remotePath)
 	if err != nil {
 		return fmt.Errorf("error remove smb mapping '%s'. err: %v", remotePath, err)
 	}