Repository: GreenmaskIO/greenmask
Pull request #126 (2024-05-17 13:04:20)

Fix/units and allowed types
* Added text and varchar support for IP and MAC addresses
* Fixed unit tests
* Fixed unit names in `RandomUnixTimestamp` transformer
* Removed old `RandomUnixTime` transformer


Patch:
diff --git a/internal/db/postgres/transformers/email.go b/internal/db/postgres/transformers/email.go
index 83ec4c1..8558e78 100644
--- a/internal/db/postgres/transformers/email.go
+++ b/internal/db/postgres/transformers/email.go
@@ -24,7 +24,7 @@ var emailTransformerAllowedChars = []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijk
 
 var EmailTransformerDefinition = utils.NewTransformerDefinition(
 	utils.NewTransformerProperties(
-		"Email",
+		"RandomEmail",
 		"Generate random email",
 	),
 
diff --git a/internal/db/postgres/transformers/random_faker.go b/internal/db/postgres/transformers/random_faker.go
index dfeba42..88868bb 100644
--- a/internal/db/postgres/transformers/random_faker.go
+++ b/internal/db/postgres/transformers/random_faker.go
@@ -50,14 +50,6 @@ var FakerTransformersDes = map[string]*FakerTransformerDef{
 		Description:    "Generates a random longitude value.",
 	},
 
-	// Faker Datetime
-	"RandomUnixTime": {
-		Generator: func(opts ...options.OptionFunc) string {
-			return fmt.Sprintf("%d", faker.UnixTime())
-		},
-		SupportedTypes: []string{"int4", "int8", "numeric"},
-		Description:    "Generates a random Unix timestamp.",
-	},
 	"RandomMonthName": {
 		Generator:      faker.MonthName,
 		SupportedTypes: []string{"text", "varchar"},
@@ -90,16 +82,6 @@ var FakerTransformersDes = map[string]*FakerTransformerDef{
 	},
 
 	// Faker Internet
-	"RandomEmail": {
-		Generator:      faker.Email,
-		SupportedTypes: []string{"text", "varchar"},
-		Description:    "Generates a random email address.",
-	},
-	"RandomMacAddress": {
-		Generator:      faker.MacAddress,
-		SupportedTypes: []string{"text", "varchar", "macaddr", "macaddr8"},
-		Description:    "Generates a random MAC address.",
-	},
 	"RandomDomainName": {
 		Generator:      faker.DomainName,
 		SupportedTypes: []string{"text", "varchar"},
@@ -115,16 +97,6 @@ var FakerTransformersDes = map[string]*FakerTransformerDef{
 		SupportedTypes: []string{"text", "varchar"},
 		Description:    "Generates a random username.",
 	},
-	"RandomIPv4": {
-		Generator:      faker.IPv4,
-		SupportedTypes: []string{"text", "varchar", "inet"},
-		Description:    "Generates a random IPv4 address.",
-	},
-	"RandomIPv6": {
-		Generator:      faker.IPv6,
-		SupportedTypes: []string{"text", "varchar", "inet"},
-		Description:    "Generates a random IPv6 address.",
-	},
 	"RandomPassword": {
 		Generator:      faker.Password,
 		SupportedTypes: []string{"text", "varchar"},
diff --git a/internal/db/postgres/transformers/random_ip.go b/internal/db/postgres/transformers/random_ip.go
index 96f454b..3e1d9fc 100644
--- a/internal/db/postgres/transformers/random_ip.go
+++ b/internal/db/postgres/transformers/random_ip.go
@@ -44,7 +44,7 @@ var RandomIpDefinition = utils.NewTransformerDefinition(
 		SetCastDbType("cidr").
 		SetDynamicMode(
 			toolkit.NewDynamicModeProperties().
-				SetCompatibleTypes("cidr", "text", "varchar"),
+				SetCompatibleTypes("text", "varchar", "cidr"),
 		),
 
 	engineParameterDefinition,
@@ -146,8 +146,8 @@ func (rbt *RandomIp) Transform(ctx context.Context, r *toolkit.Record) (*toolkit
 		return nil, fmt.Errorf("unable to transform value: %w", err)
 	}
 
-	newRawValue := toolkit.NewRawValue([]byte(ipVal.String()), false)
-	if err = r.SetRawColumnValueByIdx(rbt.columnIdx, newRawValue); err != nil {
+	newVal := toolkit.NewRawValue([]byte(ipVal.String()), false)
+	if err = r.SetRawColumnValueByIdx(rbt.columnIdx, newVal); err != nil {
 		return nil, fmt.Errorf("unable to set new value: %w", err)
 	}
 	return r, nil
diff --git a/internal/db/postgres/transformers/random_mac.go b/internal/db/postgres/transformers/random_mac.go
index 2a9d78d..fa7f210 100644
--- a/internal/db/postgres/transformers/random_mac.go
+++ b/internal/db/postgres/transformers/random_mac.go
@@ -22,6 +22,7 @@ import (
 	"github.com/greenmaskio/greenmask/internal/db/postgres/transformers/utils"
 	"github.com/greenmaskio/greenmask/internal/generators/transformers"
 	"github.com/greenmaskio/greenmask/pkg/toolkit"
+	"github.com/jackc/pgx/v5/pgtype"
 )
 
 const (
@@ -46,7 +47,7 @@ var RandomMacAddressDefinition = utils.NewTransformerDefinition(
 		"Column name",
 	).SetIsColumn(toolkit.NewColumnProperties().
 		SetAffected(true).
-		SetAllowedColumnTypes("macaddr"),
+		SetAllowedColumnTypes("macaddr", "varchar", "text"),
 	).SetRequired(true),
 
 	toolkit.MustNewParameterDefinition(
@@ -172,7 +173,11 @@ func (rbt *RandomMac) Done(ctx context.Context) error {
 
 func (rbt *RandomMac) Transform(ctx context.Context, r *toolkit.Record) (*toolkit.Record, error) {
 
-	_, err := r.ScanColumnValueByIdx(rbt.columnIdx, &rbt.originalMac)
+	rawVal, err := r.GetRawColumnValueByIdx(rbt.columnIdx)
+	if err != nil {
+		return nil, fmt.Errorf("unable to scan value: %w", err)
+	}
+	err = r.Driver.ScanValueByTypeOid(pgtype.MacaddrOID, rawVal.Data, &rbt.originalMac)
 	if err != nil {
 		return nil, fmt.Errorf("unable to scan value: %w", err)
 	}
@@ -182,7 +187,8 @@ func (rbt *RandomMac) Transform(ctx context.Context, r *toolkit.Record) (*toolki
 		return nil, fmt.Errorf("unable to transform value: %w", err)
 	}
 
-	if err = r.SetColumnValueByIdx(rbt.columnIdx, macAddr); err != nil {
+	newVal := toolkit.NewRawValue([]byte(macAddr.String()), false)
+	if err = r.SetRawColumnValueByIdx(rbt.columnIdx, newVal); err != nil {
 		return nil, fmt.Errorf("unable to set new value: %w", err)
 	}
 
diff --git a/internal/db/postgres/transformers/random_unix_timestamp.go b/internal/db/postgres/transformers/random_unix_timestamp.go
index 2e00d94..daedf92 100644
--- a/internal/db/postgres/transformers/random_unix_timestamp.go
+++ b/internal/db/postgres/transformers/random_unix_timestamp.go
@@ -13,12 +13,16 @@ import (
 )
 
 const (
-	secondsUnit = "sec"
-	milliUnit   = "milli"
-	microUnit   = "micro"
-	nanoUnit    = "nano"
+	secondsUnit = "second"
+	milliUnit   = "millisecond"
+	microUnit   = "microsecond"
+	nanoUnit    = "nanosecond"
 )
 
+var timestampUnitValues = []string{
+	secondsUnit, milliUnit, microUnit, nanoUnit,
+}
+
 var unixTimestampTransformerDefinition = utils.NewTransformerDefinition(
 	utils.NewTransformerProperties(
 		"RandomUnixTimestamp",
@@ -219,12 +223,13 @@ func getUnixByUnit(v time.Time, unit string) int64 {
 
 func validateDateUnitParameterValue(p *toolkit.ParameterDefinition, v toolkit.ParamsValue) (toolkit.ValidationWarnings, error) {
 
-	if !slices.Contains([]string{secondsUnit, milliUnit, microUnit, nanoUnit}, string(v)) {
+	if !slices.Contains(timestampUnitValues, string(v)) {
 		return toolkit.ValidationWarnings{
 			toolkit.NewValidationWarning().
 				SetSeverity(toolkit.ErrorValidationSeverity).
 				AddMeta("ParameterValue", string(v)).
-				SetMsg("wrong truncation part value: must be one of nano, second, minute, hour, day, month, year"),
+				AddMeta("AllowedValues", truncateParts).
+				SetMsg("wrong timestamp unit value"),
 		}, nil
 	}
 	return nil, nil
diff --git a/pkg/toolkit/tesing_helpers.go b/pkg/toolkit/tesing_helpers.go
index ea13c52..b4225f9 100644
--- a/pkg/toolkit/tesing_helpers.go
+++ b/pkg/toolkit/tesing_helpers.go
@@ -198,6 +198,22 @@ var columnList = []*Column{
 		NotNull:  false,
 		Length:   -1,
 	},
+	{
+		Name:     "ip_address",
+		TypeName: "inet",
+		TypeOid:  pgtype.InetOID,
+		Num:      22,
+		NotNull:  false,
+		Length:   -1,
+	},
+	{
+		Name:     "net_mask",
+		TypeName: "cidr",
+		TypeOid:  pgtype.CIDROID,
+		Num:      23,
+		NotNull:  false,
+		Length:   -1,
+	},
 }
 
 // GetDriverAndRecord - return adhoc table for testing


Test patch:
diff --git a/internal/db/postgres/transformers/random_ip_test.go b/internal/db/postgres/transformers/random_ip_test.go
index 69f617f..4ef4c89 100644
--- a/internal/db/postgres/transformers/random_ip_test.go
+++ b/internal/db/postgres/transformers/random_ip_test.go
@@ -2,10 +2,11 @@ package transformers
 
 import (
 	"context"
+	"testing"
+
 	"github.com/greenmaskio/greenmask/internal/db/postgres/transformers/utils"
 	"github.com/greenmaskio/greenmask/pkg/toolkit"
 	"github.com/stretchr/testify/require"
-	"testing"
 )
 
 func TestRandomIpTransformer_Transform_random_dynamic(t *testing.T) {
@@ -19,7 +20,24 @@ func TestRandomIpTransformer_Transform_random_dynamic(t *testing.T) {
 		expected      string
 	}{
 		{
-			name:       "IPv4 dynamic test",
+			name:       "IPv4 dynamic test with strict types",
+			columnName: "ip_address",
+			record: map[string]*toolkit.RawValue{
+				"ip_address": toolkit.NewRawValue([]byte("192.168.1.10"), false),
+				"net_mask":   toolkit.NewRawValue([]byte("192.168.1.0/30"), false),
+			},
+			params: map[string]toolkit.ParamsValue{
+				"engine": toolkit.ParamsValue("random"),
+			},
+			dynamicParams: map[string]*toolkit.DynamicParamValue{
+				"subnet": {
+					Column: "net_mask",
+				},
+			},
+			expected: "192.168.1.[1,2]",
+		},
+		{
+			name:       "IPv4 dynamic test with strings",
 			columnName: "data",
 			record: map[string]*toolkit.RawValue{
 				"data":  toolkit.NewRawValue([]byte("192.168.1.10"), false),
diff --git a/internal/db/postgres/transformers/random_mac_test.go b/internal/db/postgres/transformers/random_mac_test.go
index b51be9f..30dbae0 100644
--- a/internal/db/postgres/transformers/random_mac_test.go
+++ b/internal/db/postgres/transformers/random_mac_test.go
@@ -99,6 +99,17 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 			managementType: managementTypeNameLocal,
 			castType:       castTypeNameGroup,
 		},
+		{
+			name:       "Random mac addr in text without keepOriginalVendor with Universal and Individual",
+			columnName: "data",
+			original:   "03:1a:2b:3c:4d:5e",
+			params: map[string]toolkit.ParamsValue{
+				"engine":               toolkit.ParamsValue("hash"),
+				"keep_original_vendor": toolkit.ParamsValue("false"),
+			},
+			managementType: managementTypeNameLocal,
+			castType:       castTypeNameGroup,
+		},
 	}
 
 	for _, tt := range tests {
@@ -130,18 +141,19 @@ func TestRandomMacTransformer_Transform_random(t *testing.T) {
 			)
 			require.NoError(t, err)
 			var res net.HardwareAddr
-			isNull, err := r.ScanColumnValueByName(tt.columnName, &res)
+			rawVal, err := r.GetRawColumnValueByName(tt.columnName)
+			require.NoError(t, err)
+			require.False(t, rawVal.IsNull)
+			require.NotEmpty(t, rawVal.Data)
+			err = r.Driver.ScanValueByTypeName("macaddr", rawVal.Data, &res)
 			require.NoError(t, err)
-			require.False(t, isNull)
 
 			newMacAddrInfo, err := transformers.ExploreMacAddress(res)
 			require.NoError(t, err)
 
 			if string(tt.params["keep_original_vendor"]) == "true" {
-				require.Equal(t, tt.original[:8], res.String()[0:8])
-			}
-
-			if tt.castType != castTypeNameAny && tt.managementType != managementTypeNameAny {
+				require.Equal(t, tt.original[:8], res.String()[:8])
+			} else if tt.castType != castTypeNameAny && tt.managementType != managementTypeNameAny {
 				require.True(
 					t,
 					castTypeNameToIndex(tt.castType) == newMacAddrInfo.CastType && managementTypeNameToIndex(tt.managementType) == newMacAddrInfo.ManagementType,