Mitigating PowerShell Terminal Spoofing

cns/service/main.go

@@ -787,8 +787,10 @@ func main() {
 	}
 
 	// Setting the remote ARP MAC address to 12-34-56-78-9a-bc on windows for external traffic if HNS is enabled
-	execClient := platform.NewExecClient(nil)
-	err = platform.SetSdnRemoteArpMacAddress(execClient)
+
+	registry := platform.NewRegistryClient(nil)
+
+	err = platform.SetSdnRemoteArpMacAddress(registry)
 	if err != nil {
 		logger.Errorf("Failed to set remote ARP MAC address: %v", err)
 		return

platform/mockregistryclient.go

@@ -0,0 +1,64 @@
+package platform
+
+import (
+	"errors"
+
+	"golang.org/x/sys/windows/registry"
+)
+
+type MockRegistryClient struct {
+	returnError       bool
+	keys              map[string]*mockRegistryKey
+	setOpenKey        setOpenKey
+	setRestartService setRestartService
+}
+
+type (
+	setOpenKey        func(k registry.Key, path string, access uint32) (RegistryKey, error)
+	setRestartService func(serviceName string) error
+)
+
+type mockRegistryKey struct {
+	Values map[string]string
+}
+
+func NewMockRegistryClient(returnErr bool) *MockRegistryClient {
+	return &MockRegistryClient{
+		returnError: returnErr,
+		keys:        make(map[string]*mockRegistryKey),
+	}
+}
+
+func (r *MockRegistryClient) SetOpenKey(fn setOpenKey) {
+	r.setOpenKey = fn
+}
+
+func (r *MockRegistryClient) OpenKey(k registry.Key, path string, access uint32) (RegistryKey, error) {
+	return r.setOpenKey(k, path, access)
+
+}
+
+func (k *mockRegistryKey) GetStringValue(name string) (string, uint32, error) {
+	if value, exists := k.Values[name]; exists {
+		return value, registry.SZ, nil
+	}
+	return "", 0, errors.New("value does not exist")
+}
+
+func (k *mockRegistryKey) SetStringValue(name string, value string) error {
+	k.Values[name] = name
+	return nil
+}
+
+func (k *mockRegistryKey) Close() error {
+	return nil
+}
+
+func (r *MockRegistryClient) SetRestartService(fn setRestartService) {
+	r.setRestartService = fn
+}
+
+func (r *MockRegistryClient) restartService(serviceName string) error {
+	return r.setRestartService(serviceName)
+
+}

platform/os_windows.go

@@ -20,6 +20,7 @@
 	"github.com/pkg/errors"
 	"go.uber.org/zap"
 	"golang.org/x/sys/windows"
+	"golang.org/x/sys/windows/registry"
 )
 
 const (
@@ -61,24 +62,12 @@
 	// for vlan tagged arp requests
 	SDNRemoteArpMacAddress = "12-34-56-78-9a-bc"
 
-	// Command to get SDNRemoteArpMacAddress registry key
-	GetSdnRemoteArpMacAddressCommand = "(Get-ItemProperty " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress).SDNRemoteArpMacAddress"
-
-	// Command to set SDNRemoteArpMacAddress registry key
-	SetSdnRemoteArpMacAddressCommand = "Set-ItemProperty " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
-
-	// Command to check if system has hns state path or not
-	CheckIfHNSStatePathExistsCommand = "Test-Path " +
-		"-Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State"
-
 	// Command to fetch netadapter and pnp id
+	//TODO can we replace this (and things in endpoint_windows) with "golang.org/x/sys/windows"
+	//var adapterInfo windows.IpAdapterInfo
+	//var bufferSize uint32 = uint32(unsafe.Sizeof(adapterInfo))
 	GetMacAddressVFPPnpIDMapping = "Get-NetAdapter | Select-Object MacAddress, PnpDeviceID| Format-Table -HideTableHeaders"
 
-	// Command to restart HNS service
-	RestartHnsServiceCommand = "Restart-Service -Name hns"
-
 	// Interval between successive checks for mellanox adapter's PriorityVLANTag value
 	defaultMellanoxMonitorInterval = 30 * time.Second
 
@@ -257,40 +246,51 @@
 }
 
 // SetSdnRemoteArpMacAddress sets the regkey for SDNRemoteArpMacAddress needed for multitenancy if hns is enabled
-func SetSdnRemoteArpMacAddress(execClient ExecClient) error {
-	exists, err := execClient.ExecutePowershellCommand(CheckIfHNSStatePathExistsCommand)
+func SetSdnRemoteArpMacAddress(registerClient RegistryClient) error {
+	key, err := registerClient.OpenKey(registry.LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\hns\\State", registry.READ|registry.SET_VALUE)
 	if err != nil {
+		if err == registry.ErrNotExist {
+			log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
+			return nil
+		}
 		errMsg := fmt.Sprintf("Failed to check the existent of hns state path due to error %s", err.Error())
 		log.Printf(errMsg)
 		return errors.Errorf(errMsg)
 	}
-	if strings.EqualFold(exists, "false") {
+
+	if key == nil {
 		log.Printf("hns state path does not exist, skip setting SdnRemoteArpMacAddress")
 		return nil
 	}
+
 	if sdnRemoteArpMacAddressSet == false {
-		result, err := execClient.ExecutePowershellCommand(GetSdnRemoteArpMacAddressCommand)
+
+		//Was (Get-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress).SDNRemoteArpMacAddress"
+		result, _, err := key.GetStringValue("SDNRemoteArpMacAddress")
 		if err != nil {
 			return err
 		}
 
 		// Set the reg key if not already set or has incorrect value
 		if result != SDNRemoteArpMacAddress {
-			if _, err = execClient.ExecutePowershellCommand(SetSdnRemoteArpMacAddressCommand); err != nil {
+
+			//was "Set-ItemProperty -Path HKLM:\\SYSTEM\\CurrentControlSet\\Services\\hns\\State -Name SDNRemoteArpMacAddress -Value \"12-34-56-78-9a-bc\""
+
+			if err := key.SetStringValue("SDNRemoteArpMacAddress", SDNRemoteArpMacAddress); err != nil {
 				log.Printf("Failed to set SDNRemoteArpMacAddress due to error %s", err.Error())
 				return err
 			}
-
 			log.Printf("[Azure CNS] SDNRemoteArpMacAddress regKey set successfully. Restarting hns service.")
-			if _, err := execClient.ExecutePowershellCommand(RestartHnsServiceCommand); err != nil {
+
+			//	was "Restart-Service -Name hns"
+			if err := registerClient.restartService("hns"); err != nil {
 				log.Printf("Failed to Restart HNS Service due to error %s", err.Error())
 				return err
 			}
 		}
 
 		sdnRemoteArpMacAddressSet = true
 	}
-
 	return nil
 }
 
@@ -364,6 +364,7 @@
 	pidstr = strings.Trim(pidstr, "\r\n")
 	cmd := fmt.Sprintf("Get-Process -Id %s|Format-List", pidstr)
 	p := NewExecClient(nil)
+	//TODO not removing this because it seems to only be called in test?
 	out, err := p.ExecutePowershellCommand(cmd)
 	if err != nil {
 		log.Printf("Process is not running. Output:%v, Error %v", out, err)

platform/os_windows_test.go

@@ -4,14 +4,14 @@
 	"context"
 	"errors"
 	"os/exec"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/Azure/azure-container-networking/platform/windows/adapter/mocks"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/windows/registry"
 )
 
 var errTestFailure = errors.New("test failure")
@@ -116,34 +116,54 @@
 }
 
 func TestSetSdnRemoteArpMacAddress_hnsNotEnabled(t *testing.T) {
-	mockExecClient := NewMockExecClient(false)
-	// testing skip setting SdnRemoteArpMacAddress when hns not enabled
-	mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
-		return "False", nil
+
+	mockRegistryClient := NewMockRegistryClient(false)
+	// mockRegistryClient.SetOpenKey(func(k registry.Key, path string, access uint32) (RegistryKey, error) {
+	// 	return nil,nil
+	// })
+	mockRegistryClient.SetOpenKey(func(k registry.Key, path string, access uint32) (RegistryKey, error) {
+		return nil, nil
+	})
+
+	mockRegistryClient.SetRestartService(func(s string) error {
+		return nil
 	})
-	err := SetSdnRemoteArpMacAddress(mockExecClient)
+
+	err := SetSdnRemoteArpMacAddress(mockRegistryClient)
 	assert.NoError(t, err)
 	assert.Equal(t, false, sdnRemoteArpMacAddressSet)
 
 	// testing the scenario when there is an error in checking if hns is enabled or not
-	mockExecClient.SetPowershellCommandResponder(func(_ string) (string, error) {
-		return "", errTestFailure
+
+	mockRegistryClient.SetOpenKey(func(k registry.Key, path string, access uint32) (RegistryKey, error) {
+		return &mockRegistryKey{
+			Values: map[string]string{
+				"": "",
+			},
+		}, errTestFailure
 	})
-	err = SetSdnRemoteArpMacAddress(mockExecClient)
+
+	err = SetSdnRemoteArpMacAddress(mockRegistryClient)
 	assert.ErrorAs(t, err, &errTestFailure)
 	assert.Equal(t, false, sdnRemoteArpMacAddressSet)
 }
 
 func TestSetSdnRemoteArpMacAddress_hnsEnabled(t *testing.T) {
-	mockExecClient := NewMockExecClient(false)
-	// happy path
-	mockExecClient.SetPowershellCommandResponder(func(cmd string) (string, error) {
-		if strings.Contains(cmd, "Test-Path") {
-			return "True", nil
-		}
-		return "", nil
+
+	mockRegistryClient := NewMockRegistryClient(false) // happy path
+
+	mockRegistryClient.SetOpenKey(func(k registry.Key, path string, access uint32) (RegistryKey, error) {
+		return &mockRegistryKey{
+			Values: map[string]string{
+				"SDNRemoteArpMacAddress": "MockData1",
+			},
+		}, nil
+	})
+	mockRegistryClient.SetRestartService(func(s string) error {
+		return nil
 	})
-	err := SetSdnRemoteArpMacAddress(mockExecClient)
+
+	err := SetSdnRemoteArpMacAddress(mockRegistryClient)
 	assert.NoError(t, err)
 	assert.Equal(t, true, sdnRemoteArpMacAddressSet)
 	// reset sdnRemoteArpMacAddressSet

platform/registryInterface_windows.go

@@ -0,0 +1,102 @@
+package platform
+
+import (
+	"fmt"
+	"time"
+
+	"go.uber.org/zap"
+	"golang.org/x/sys/windows/registry"
+	"golang.org/x/sys/windows/svc"
+	"golang.org/x/sys/windows/svc/mgr"
+)
+
+type RegistryClient interface {
+	OpenKey(k registry.Key, path string, access uint32) (RegistryKey, error)
+	restartService(serviceName string) error
+}
+
+type RegistryKey interface {
+	GetStringValue(name string) (val string, valtype uint32, err error)
+	SetStringValue(name, value string) error
+	Close() error
+}
+
+type registryClient struct {
+	Timeout time.Duration
+	logger  *zap.Logger
+}
+
+type registryKey struct {
+	key registry.Key
+}
+
+func NewRegistryClient(logger *zap.Logger) RegistryClient {
+	return &registryClient{
+		Timeout: defaultExecTimeout * time.Second,
+		logger:  logger,
+	}
+}
+
+func (r *registryClient) OpenKey(k registry.Key, path string, access uint32) (RegistryKey, error) {
+	key, err := registry.OpenKey(k, path, access)
+	if err != nil {
+		return nil, err
+	}
+	return &registryKey{key}, nil
+}
+
+func (k *registryKey) GetStringValue(name string) (string, uint32, error) {
+	return k.key.GetStringValue(name)
+}
+
+func (k *registryKey) SetStringValue(name string, value string) error {
+	return k.key.SetStringValue(name, value)
+}
+
+func (k *registryKey) Close() error {
+	return k.key.Close()
+}
+
+// straight out of chat gpt
+func (r *registryClient) restartService(serviceName string) error {
+	// Connect to the service manager
+	m, err := mgr.Connect()
+	if err != nil {
+		return fmt.Errorf("could not connect to service manager: %v", err)
+	}
+	defer m.Disconnect()
+
+	// Open the service by name
+	service, err := m.OpenService(serviceName)
+	if err != nil {
+		return fmt.Errorf("could not access service: %v", err)
+	}
+	defer service.Close()
+
+	// Stop the service
+	_, err = service.Control(svc.Stop)
+	if err != nil {
+		return fmt.Errorf("could not stop service: %v", err)
+	}
+
+	// Wait for the service to stop
+	status, err := service.Query()
+	if err != nil {
+		return fmt.Errorf("could not query service status: %v", err)
+	}
+	for status.State != svc.Stopped {
+		time.Sleep(500 * time.Millisecond)
+		status, err = service.Query()
+		if err != nil {
+			return fmt.Errorf("could not query service status: %v", err)
+		}
+	}
+
+	// Start the service again
+	err = service.Start()
+	if err != nil {
+		return fmt.Errorf("could not start service: %v", err)
+	}
+
+	return nil
+}