Repository: GreenmaskIO/greenmask
Pull request #125 (2024-05-16 18:22:36)

fix: Revised RandomMac transformer
* Replaced logic from manual parsing to net.HardwareAddr object
* Revised encode-decoding logic via driver
* Revised tests and renamed some functions


Patch:
diff --git a/internal/db/postgres/transformers/random_mac.go b/internal/db/postgres/transformers/random_mac.go
index 8c4116e..2a9d78d 100644
--- a/internal/db/postgres/transformers/random_mac.go
+++ b/internal/db/postgres/transformers/random_mac.go
@@ -17,6 +17,8 @@ package transformers
 import (
 	"context"
 	"fmt"
+	"net"
+
 	"github.com/greenmaskio/greenmask/internal/db/postgres/transformers/utils"
 	"github.com/greenmaskio/greenmask/internal/generators/transformers"
 	"github.com/greenmaskio/greenmask/pkg/toolkit"
@@ -44,7 +46,7 @@ var RandomMacAddressDefinition = utils.NewTransformerDefinition(
 		"Column name",
 	).SetIsColumn(toolkit.NewColumnProperties().
 		SetAffected(true).
-		SetAllowedColumnTypes("text", "varchar", "macaddr"),
+		SetAllowedColumnTypes("macaddr"),
 	).SetRequired(true),
 
 	toolkit.MustNewParameterDefinition(
@@ -89,6 +91,7 @@ type RandomMac struct {
 	managementType     int
 	t                  *transformers.MacAddress
 	settings           *RandomMacSettings
+	originalMac        net.HardwareAddr
 }
 
 func NewMacAddressTransformer(ctx context.Context, driver *toolkit.Driver, parameters map[string]toolkit.Parameterizer) (utils.Transformer, toolkit.ValidationWarnings, error) {
@@ -169,18 +172,17 @@ func (rbt *RandomMac) Done(ctx context.Context) error {
 
 func (rbt *RandomMac) Transform(ctx context.Context, r *toolkit.Record) (*toolkit.Record, error) {
 
-	val, err := r.GetRawColumnValueByIdx(rbt.columnIdx)
+	_, err := r.ScanColumnValueByIdx(rbt.columnIdx, &rbt.originalMac)
 	if err != nil {
 		return nil, fmt.Errorf("unable to scan value: %w", err)
 	}
 
-	macAddr, err := rbt.t.Generate(val.Data, rbt.keepOriginalVendor, rbt.castType, rbt.managementType)
+	macAddr, err := rbt.t.Generate(rbt.originalMac, rbt.keepOriginalVendor, rbt.castType, rbt.managementType)
 	if err != nil {
 		return nil, fmt.Errorf("unable to transform value: %w", err)
 	}
 
-	newRawValue := toolkit.NewRawValue(macAddr, false)
-	if err = r.SetRawColumnValueByIdx(rbt.columnIdx, newRawValue); err != nil {
+	if err = r.SetColumnValueByIdx(rbt.columnIdx, macAddr); err != nil {
 		return nil, fmt.Errorf("unable to set new value: %w", err)
 	}
 
diff --git a/internal/generators/transformers/random_mac.go b/internal/generators/transformers/random_mac.go
index a04076b..14acf04 100644
--- a/internal/generators/transformers/random_mac.go
+++ b/internal/generators/transformers/random_mac.go
@@ -16,6 +16,8 @@ package transformers
 
 import (
 	"fmt"
+	"net"
+
 	"github.com/greenmaskio/greenmask/internal/generators"
 )
 
@@ -37,7 +39,7 @@ type MacAddress struct {
 }
 
 type MacAddressInfo struct {
-	MacAddressStr  string
+	MacAddress     net.HardwareAddr
 	ManagementType int
 	CastType       int
 }
@@ -52,13 +54,17 @@ func (macAddr *MacAddress) GetRequiredGeneratorByteLength() int {
 	return macAddr.byteLength
 }
 
-func (macAddr *MacAddress) Generate(original []byte, keepOriginalVendor bool, castType int, managementType int) ([]byte, error) {
-	hostBytes, err := macAddr.generator.Generate(original)
+func (macAddr *MacAddress) Generate(original net.HardwareAddr, keepOriginalVendor bool, castType int, managementType int) (net.HardwareAddr, error) {
+	randoBytes, err := macAddr.generator.Generate(original)
 	if err != nil {
 		return nil, fmt.Errorf("error generating random bytes: %w", err)
 	}
+	randomMac, err := RandomBytesToHardwareAddr(randoBytes)
+	if err != nil {
+		return nil, fmt.Errorf("error converting random bytes to hardware address: %w", err)
+	}
 
-	result, err := ModifyMacAddress(hostBytes, original, keepOriginalVendor, castType, managementType)
+	result, err := ModifyMacAddress(randomMac, original, keepOriginalVendor, castType, managementType)
 	if err != nil {
 		return nil, fmt.Errorf("can't modify mac address: %w", err)
 	}
@@ -66,27 +72,11 @@ func (macAddr *MacAddress) Generate(original []byte, keepOriginalVendor bool, ca
 	return result, nil
 }
 
-func ModifyMacAddress(newMac, original []byte, keepOriginalVendor bool, castType, managementType int) ([]byte, error) {
+func ModifyMacAddress(newMac, original net.HardwareAddr, keepOriginalVendor bool, castType, managementType int) ([]byte, error) {
 	if keepOriginalVendor {
-		originalMacAddrInfo, err := ParseMacAddr(original)
-		if err != nil {
-			return nil, fmt.Errorf("can't get original mac address info: %v", err)
-		}
-
-		// Logic with control U/L bits
-		if originalMacAddrInfo.ManagementType == ManagementTypeLocal {
-			newMac[0] |= 0x02
-		} else {
-			newMac[0] &^= 0x02
-		}
-
-		// Logic with control I/G bits
-		if originalMacAddrInfo.CastType == CastTypeGroup {
-			newMac[0] |= 0x01
-		} else {
-			newMac[0] &^= 0x01
-		}
-
+		newMac[0] = original[0]
+		newMac[1] = original[1]
+		newMac[2] = original[2]
 	} else {
 		if managementType == ManagementTypeUniversal || managementType == ManagementTypeLocal {
 			if managementType == ManagementTypeLocal {
@@ -108,7 +98,8 @@ func ModifyMacAddress(newMac, original []byte, keepOriginalVendor bool, castType
 	return newMac, nil
 }
 
-func ParseMacAddr(macAddress []byte) (*MacAddressInfo, error) {
+// ExploreMacAddress - explore mac address and return info about it
+func ExploreMacAddress(macAddress net.HardwareAddr) (*MacAddressInfo, error) {
 	firstByte := macAddress[0]
 	managementType := ManagementTypeUniversal
 	if firstByte&0x02 == 0x02 {
@@ -120,12 +111,7 @@ func ParseMacAddr(macAddress []byte) (*MacAddressInfo, error) {
 		castType = CastTypeGroup
 	}
 
-	macSrt, err := MacBytesToString(macAddress)
-	if err != nil {
-		return nil, fmt.Errorf("can't create mac address string: %v", err)
-	}
-
-	return &MacAddressInfo{ManagementType: managementType, CastType: castType, MacAddressStr: macSrt}, nil
+	return &MacAddressInfo{ManagementType: managementType, CastType: castType, MacAddress: macAddress}, nil
 }
 
 func (macAddr *MacAddress) SetGenerator(g generators.Generator) error {
@@ -136,10 +122,11 @@ func (macAddr *MacAddress) SetGenerator(g generators.Generator) error {
 	return nil
 }
 
-func MacBytesToString(originalBytes []byte) (macString string, err error) {
-	if len(originalBytes) < 6 {
-		return "", fmt.Errorf("incorrect size of MAC-address")
-	}
-
-	return fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x", originalBytes[0], originalBytes[1], originalBytes[2], originalBytes[3], originalBytes[4], originalBytes[5]), nil
+func RandomBytesToHardwareAddr(originalBytes []byte) (net.HardwareAddr, error) {
+	return net.ParseMAC(
+		fmt.Sprintf(
+			"%02x:%02x:%02x:%02x:%02x:%02x",
+			originalBytes[0], originalBytes[1], originalBytes[2], originalBytes[3], originalBytes[4], originalBytes[5],
+		),
+	)
 }


Test patch:
diff --git a/internal/db/postgres/transformers/random_mac_test.go b/internal/db/postgres/transformers/random_mac_test.go
index cb13680..b51be9f 100644
--- a/internal/db/postgres/transformers/random_mac_test.go
+++ b/internal/db/postgres/transformers/random_mac_test.go
@@ -3,11 +3,13 @@ package transformers
 import (
 	"context"
 	"fmt"
+	"net"
+	"testing"
+
 	"github.com/greenmaskio/greenmask/internal/db/postgres/transformers/utils"
 	"github.com/greenmaskio/greenmask/internal/generators/transformers"
 	"github.com/greenmaskio/greenmask/pkg/toolkit"
 	"github.com/stretchr/testify/require"
-	"testing"
 )
 
 func TestRandomMacTransformer_Transform_random(t *testing.T) {
@@ -15,15 +17,15 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 	tests := []struct {
 		name           string
 		columnName     string
-		original       []byte
+		original       string
 		params         map[string]toolkit.ParamsValue
 		castType       string
 		managementType string
 	}{
 		{
 			name:       "Random mac addr with keepOriginalVendor with Universal and Individual",
-			columnName: "data",
-			original:   []byte("00:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "00:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("true"),
@@ -33,8 +35,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr with keepOriginalVendor with Universal and Group",
-			columnName: "data",
-			original:   []byte("01:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "01:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("true"),
@@ -44,8 +46,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr with keepOriginalVendor with Local and Group",
-			columnName: "data",
-			original:   []byte("03:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "03:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("true"),
@@ -55,8 +57,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr without keepOriginalVendor with Universal and Group",
-			columnName: "data",
-			original:   []byte("03:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "03:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("false"),
@@ -66,8 +68,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr without keepOriginalVendor with Universal and Individual",
-			columnName: "data",
-			original:   []byte("03:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "03:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("false"),
@@ -77,8 +79,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr without keepOriginalVendor with Local and Individual",
-			columnName: "data",
-			original:   []byte("03:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "03:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("false"),
@@ -88,8 +90,8 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 		},
 		{
 			name:       "Random mac addr without keepOriginalVendor with Universal and Individual",
-			columnName: "data",
-			original:   []byte("03:1a:2b:3c:4d:5e"),
+			columnName: "macaddress",
+			original:   "03:1a:2b:3c:4d:5e",
 			params: map[string]toolkit.ParamsValue{
 				"engine":               toolkit.ParamsValue("hash"),
 				"keep_original_vendor": toolkit.ParamsValue("false"),
@@ -101,15 +103,15 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			parsedMacOriginal, err := transformers.ParseMacAddr(tt.original)
+			originalMacAddrInfo, err := transformers.ExploreMacAddress([]byte(tt.original))
 			require.NoError(t, err)
-			require.NotEmpty(t, parsedMacOriginal)
+			require.NotEmpty(t, originalMacAddrInfo)
 
 			tt.params["column"] = toolkit.ParamsValue(tt.columnName)
 			tt.params["cast_type"] = toolkit.ParamsValue(tt.castType)
 			tt.params["management_type"] = toolkit.ParamsValue(tt.managementType)
 
-			driver, record := getDriverAndRecord(tt.columnName, string(tt.original))
+			driver, record := getDriverAndRecord(tt.columnName, tt.original)
 			def, ok := utils.DefaultTransformerRegistry.Get("RandomMac")
 			require.True(t, ok)
 
@@ -127,27 +129,23 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 				record,
 			)
 			require.NoError(t, err)
-			var res []byte
+			var res net.HardwareAddr
 			isNull, err := r.ScanColumnValueByName(tt.columnName, &res)
 			require.NoError(t, err)
 			require.False(t, isNull)
 
-			parsedMacGenerated, err := transformers.ParseMacAddr(res)
+			newMacAddrInfo, err := transformers.ExploreMacAddress(res)
 			require.NoError(t, err)
 
 			if string(tt.params["keep_original_vendor"]) == "true" {
-				require.True(
-					t,
-					parsedMacOriginal.CastType == parsedMacGenerated.CastType && parsedMacOriginal.ManagementType == parsedMacGenerated.ManagementType,
-					fmt.Sprintf("Mac address info is't equals %+v != %+v", parsedMacOriginal, parsedMacGenerated),
-				)
+				require.Equal(t, tt.original[:8], res.String()[0:8])
 			}
 
 			if tt.castType != castTypeNameAny && tt.managementType != managementTypeNameAny {
 				require.True(
 					t,
-					castTypeNameToIndex(tt.castType) == parsedMacGenerated.CastType && managementTypeNameToIndex(tt.managementType) == parsedMacGenerated.ManagementType,
-					fmt.Sprintf("Mac address info is't equals %+v != %+v", parsedMacOriginal, parsedMacGenerated),
+					castTypeNameToIndex(tt.castType) == newMacAddrInfo.CastType && managementTypeNameToIndex(tt.managementType) == newMacAddrInfo.ManagementType,
+					fmt.Sprintf("Mac address info is't equals %+v != %+v", originalMacAddrInfo, newMacAddrInfo),
 				)
 			}
 		})
diff --git a/internal/db/postgres/transformers/tesing_helpers_test.go b/internal/db/postgres/transformers/tesing_helpers_test.go
index 0408c56..510aeb5 100644
--- a/internal/db/postgres/transformers/tesing_helpers_test.go
+++ b/internal/db/postgres/transformers/tesing_helpers_test.go
@@ -164,6 +164,14 @@ var columnList = []*toolkit.Column{
 		NotNull:  false,
 		Length:   -1,
 	},
+	{
+		Name:     "macaddress",
+		TypeName: "numeric",
+		TypeOid:  pgtype.MacaddrOID,
+		Num:      17,
+		NotNull:  false,
+		Length:   -1,
+	},
 }
 
 // getDriverAndRecord - return adhoc table for testing
diff --git a/internal/generators/transformers/random_mac_test.go b/internal/generators/transformers/random_mac_test.go
index d28fb0c..cb1cbe0 100644
--- a/internal/generators/transformers/random_mac_test.go
+++ b/internal/generators/transformers/random_mac_test.go
@@ -2,12 +2,13 @@ package transformers
 
 import (
 	"fmt"
-	"github.com/greenmaskio/greenmask/internal/generators"
-	"github.com/rs/zerolog/log"
-	"github.com/stretchr/testify/require"
 	"net"
 	"testing"
 	"time"
+
+	"github.com/greenmaskio/greenmask/internal/generators"
+	"github.com/rs/zerolog/log"
+	"github.com/stretchr/testify/require"
 )
 
 func TestMacAddress_Generate(t *testing.T) {
@@ -90,7 +91,7 @@ func TestMacAddress_Generate(t *testing.T) {
 			_, err := net.ParseMAC(string(tt.original))
 			require.NoError(t, err)
 
-			parsedMacOriginal, err := ParseMacAddr(tt.original)
+			parsedMacOriginal, err := ExploreMacAddress(tt.original)
 			require.NoError(t, err)
 
 			tr, err := NewMacAddress()
@@ -105,14 +106,7 @@ func TestMacAddress_Generate(t *testing.T) {
 			res, err = tr.Generate(tt.original, tt.keepOriginalVendor, tt.castType, tt.managementType)
 			require.NoError(t, err)
 
-			resultStr, err := MacBytesToString(res)
-			require.NotEmpty(t, resultStr)
-			require.NoError(t, err)
-
-			_, err = net.ParseMAC(resultStr)
-			require.NoError(t, err)
-
-			parsedMacGenerated, err := ParseMacAddr(res)
+			parsedMacGenerated, err := ExploreMacAddress(res)
 			require.NoError(t, err)
 
 			if tt.keepOriginalVendor == true {
@@ -133,7 +127,7 @@ func TestMacAddress_Generate(t *testing.T) {
 
 			log.Debug().
 				Str("macAddrOriginal", string(tt.original)).
-				Str("macAddr", parsedMacGenerated.MacAddressStr).
+				Str("macAddr", parsedMacGenerated.MacAddress.String()).
 				Msg("result")
 		})
 	}