refactor secret masking

.vscode/launch.json

@@ -0,0 +1,13 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Launch Package",
+            "type": "go",
+            "request": "launch",
+            "mode": "debug",
+            "program": "internal/secrets/maskSecrets.go",
+            "args": ["../../../kics/assets/queries/common/passwords_and_secrets/test/positive16.yaml"]
+        }
+    ]
+}

go.mod

@@ -2,4 +2,13 @@ module github.com/checkmarxDev/gpt-wrapper
 
 go 1.17
 
-require github.com/google/uuid v1.3.0
+require (
+	github.com/google/uuid v1.3.0
+	github.com/rs/zerolog v1.31.0
+)
+
+require (
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
+	golang.org/x/sys v0.12.0 // indirect
+)

go.sum

@@ -1,2 +1,17 @@
+github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.31.0 h1:FcTR3NnLWW+NnTwwhFWiJSZr4ECLpqCm6QsEnyvbV4A=
+github.com/rs/zerolog v1.31.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
+golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/secrets/maskSecrets.go

@@ -3,7 +3,6 @@ package secrets
 import (
 	_ "embed"
 	"encoding/json"
-	"fmt"
 	"math"
 	"regexp"
 	"strings"
@@ -37,11 +36,11 @@ type AllowRule struct {
 type SecretRule struct {
 	ID          string      `json:"id"`
 	Name        string      `json:"name"`
+	Multiline   bool        `json:"multiline"`
 	Regex       string      `json:"regex"`
 	Entropies   []Entropy   `json:"entropies"`
-	Multiline   Multiline   `json:"multiline"`
 	AllowRules  []AllowRule `json:"allowRules"`
-	SpecialMask string      `json:"specialMask"`
+	GroupToMask int         `json:"groupMask"`
 }
 
 type SecretRules struct {
@@ -52,10 +51,11 @@ type SecretRules struct {
 type SecretRegex struct {
 	QueryName   string
 	Regex       *regexp.Regexp
-	Multiline   Multiline
+	RegexStr    string
+	Multiline   bool
+	GroupToMask int
 	Entropies   []Entropy
 	AllowRules  []*regexp.Regexp
-	SpecialMask *regexp.Regexp
 }
 
 type Result struct {
@@ -78,7 +78,6 @@ func LoadRegexps() ([]SecretRegex, []*regexp.Regexp, error) {
 	var regexes []SecretRegex
 	for _, regexStruct := range secretRules.Rules {
 		regex := regexStruct.Regex
-		specialMask := regexStruct.SpecialMask
 		var allowRules []*regexp.Regexp
 
 		for _, rule := range regexStruct.AllowRules {
@@ -88,12 +87,6 @@ func LoadRegexps() ([]SecretRegex, []*regexp.Regexp, error) {
 			}
 		}
 
-		var specialMaskCompiled *regexp.Regexp
-		if specialMask == "" {
-			specialMaskCompiled = nil
-		} else {
-			specialMaskCompiled, _ = regexp.Compile(specialMask)
-		}
 		regexCompiled, _ := regexp.Compile(regex)
 
 		secretRegex := &SecretRegex{
@@ -102,7 +95,8 @@ func LoadRegexps() ([]SecretRegex, []*regexp.Regexp, error) {
 			AllowRules:  allowRules,
 			Multiline:   regexStruct.Multiline,
 			Entropies:   regexStruct.Entropies,
-			SpecialMask: specialMaskCompiled,
+			GroupToMask: regexStruct.GroupToMask,
+			RegexStr:    regexStruct.Regex,
 		}
 		regexes = append(regexes, *secretRegex)
 	}
@@ -118,31 +112,17 @@ func LoadRegexps() ([]SecretRegex, []*regexp.Regexp, error) {
 }
 
 // getLineNumber calculates the line number based on the match index
-func getLineNumber(str string, index int) int {
+func getLineNumber(text, portion string) int {
+	index := strings.Index(text, portion) + 1
 	lineNumber := 1
 	for i := 0; i < index; i++ {
-		if str[i] == '\n' && i != index-1 {
+		if text[i] == '\n' && i != index-1 {
 			lineNumber++
 		}
 	}
 	return lineNumber
 }
 
-func getLines(str string, firstLine int, lastLine int) string {
-	lineNumber := 1
-	var returnLines []byte
-	for i := 0; i < len(str) && lineNumber <= lastLine; i++ {
-		if lineNumber >= firstLine && (str[i] != '\n' || lineNumber < lastLine) {
-			returnLines = append(returnLines, str[i])
-		}
-		if str[i] == '\n' {
-			lineNumber++
-		}
-
-	}
-	return string(returnLines)
-}
-
 // CheckEntropyInterval - verifies if a given token's entropy is within expected bounds
 func CheckEntropyInterval(entropy Entropy, token string) (isEntropyInInterval bool, entropyLevel float64) {
 	base64Entropy := calculateEntropy(token, Base64Chars)
@@ -179,6 +159,54 @@ func calculateEntropy(token, charSet string) float64 {
 	return math.Log2(length) - freq/length
 }
 
+// maskRegexByMatchGroup masks a content using a regex and the group to be masked
+func maskRegexByMatchGroup(groupToMask int, matchContent string, query *SecretRegex) string {
+	query.Regex = regexp.MustCompile(".*" + query.RegexStr) // add .* to match the last appearance
+	groups := query.Regex.FindAllStringSubmatch(matchContent, -1)
+	lastMatch := groups[len(groups)-1]
+	if len(lastMatch) < groupToMask {
+		return matchContent
+	}
+	return strings.Replace(matchContent, lastMatch[groupToMask], "<masked>", 1)
+}
+
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// IsAllowRule check if string matches any of the allow rules for the secret queries
+func IsAllowRule(s string, query *SecretRegex, allowRules []*regexp.Regexp) bool {
+	query.Regex = regexp.MustCompile(query.RegexStr)
+	regexMatch := query.Regex.FindStringIndex(s)
+	if regexMatch != nil {
+		allowRuleMatches := AllowRuleMatches(s, append(query.AllowRules, allowRules...))
+
+		for _, allowMatch := range allowRuleMatches {
+			allowStart, allowEnd := allowMatch[0], allowMatch[1]
+			regexStart, regexEnd := regexMatch[0], regexMatch[1]
+
+			if (allowStart <= regexEnd && allowStart >= regexStart) || (regexStart <= allowEnd && regexStart >= allowStart) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// AllowRuleMatches return all the allow rules matches for the secret queries
+func AllowRuleMatches(s string, allowRules []*regexp.Regexp) [][]int {
+	allowRuleMatches := [][]int{}
+	for i := range allowRules {
+		res := allowRules[i].FindAllStringIndex(s, -1)
+		allowRuleMatches = append(allowRuleMatches, res...)
+	}
+	return allowRuleMatches
+}
+
 // ReplaceMatches If matches between the regex and the file content, then replace the match with the string "<masked>"
 func ReplaceMatches(fileName string, result string, regexs []SecretRegex, allowRegexes []*regexp.Regexp) (string, []Result, []maskedSecret.MaskedSecret) {
 	var results []Result
@@ -188,20 +216,17 @@ func ReplaceMatches(fileName string, result string, regexs []SecretRegex, allowR
 	lines := strings.Split(strings.ReplaceAll(result, "\r\n", "\n"), "\n")
 	// Replace matches
 	for _, re := range regexs {
-		if re.Multiline.DetectLineGroup != 0 {
+		if re.Multiline {
 			multilineRegexes = append(multilineRegexes, re)
 			continue
 		}
 		for index, line := range lines {
-
 			originalLine := lines[index]
 			lines[index] = re.Regex.ReplaceAllStringFunc(line, func(match string) string {
-				for _, allowRule := range append(re.AllowRules, allowRegexes...) {
-					if allowRule.FindString(line) != "" {
-						return match
-					}
+				if IsAllowRule(line, &re, append(re.AllowRules, allowRegexes...)) {
+					return match
 				}
-
+				re.Regex = regexp.MustCompile(re.RegexStr)
 				groups := re.Regex.FindAllStringSubmatch(result, -1)
 				for _, entropy := range re.Entropies {
 					if len(groups) < entropy.Group {
@@ -211,11 +236,7 @@ func ReplaceMatches(fileName string, result string, regexs []SecretRegex, allowR
 					}
 				}
 
-				startOfMatch := ""
-				if re.SpecialMask != nil {
-					startOfMatch = re.SpecialMask.FindString(line)
-				}
-				maskedSecret := fmt.Sprintf("%s<masked>", startOfMatch)
+				maskedSecret := maskRegexByMatchGroup(re.GroupToMask, match, &re)
 				results = append(results, Result{QueryName: "Passwords And Secrets - " + re.QueryName, Line: index + 1, FileName: fileName, Severity: "HIGH"})
 				return maskedSecret
 			})
@@ -232,14 +253,13 @@ func ReplaceMatches(fileName string, result string, regexs []SecretRegex, allowR
 	result = strings.Join(lines[:], "\n")
 	for _, re := range multilineRegexes {
 		// Find all matches of the regular expression in the string
-		groups := re.Regex.FindStringSubmatchIndex(result)
+		re.Regex = regexp.MustCompile(re.RegexStr)
+		groups := re.Regex.FindStringSubmatch(result)
 
 		// Iterate over each match
 		for groups != nil {
 			maskedSecretElement := maskedSecret.MaskedSecret{}
-			firstLine := getLineNumber(result, groups[0])
-			lastLine := getLineNumber(result, groups[1])
-			fullContext := getLines(result, firstLine, lastLine)
+			fullContext := groups[0]
 			allowed := false
 
 			for _, allowRule := range append(re.AllowRules, allowRegexes...) {
@@ -254,40 +274,23 @@ func ReplaceMatches(fileName string, result string, regexs []SecretRegex, allowR
 			}
 
 			// Extract the matched substring
-			matchString := result[groups[0]:groups[1]]
-
-			if len(groups) <= re.Multiline.DetectLineGroup*2 {
-				groups = nil
-				continue
-			}
+			matchString := groups[0]
 
-			stringToMask := result[groups[re.Multiline.DetectLineGroup*2]:groups[re.Multiline.DetectLineGroup*2+1]]
-			lineOfSecret := getLineNumber(result, groups[re.Multiline.DetectLineGroup*2])
-
-			startOfMatch := ""
-			if re.SpecialMask != nil {
-				partOfMatches := re.SpecialMask.FindAllStringIndex(stringToMask, -1)
-				if len(partOfMatches) != 0 {
-					partOfMatch := partOfMatches[len(partOfMatches)-1]
-					startOfMatch = stringToMask[0:partOfMatch[1]]
-				}
-			}
-			maskedSecret := fmt.Sprintf("%s<masked>", startOfMatch)
+			lineOfSecret := getLineNumber(result, groups[re.GroupToMask])
 
+			maskedMatchString := maskRegexByMatchGroup(re.GroupToMask, matchString, &re)
 			results = append(results, Result{QueryName: "Passwords And Secrets - " + re.QueryName, Line: lineOfSecret, FileName: fileName, Severity: "HIGH"})
 
-			maskedMatchString := strings.Replace(matchString, stringToMask, maskedSecret, 1)
-
 			// Add the masked string to return
 			maskedSecretElement.Masked = maskedMatchString
 			maskedSecretElement.Secret = matchString
-			maskedSecretElement.Line = firstLine
+			maskedSecretElement.Line = lineOfSecret
 
 			maskedSecrets = append(maskedSecrets, maskedSecretElement)
 
 			result = strings.Replace(result, matchString, maskedMatchString, 1)
 
-			groups = re.Regex.FindStringSubmatchIndex(result)
+			groups = re.Regex.FindStringSubmatch(result)
 		}
 	}
 	return result, results, maskedSecrets

internal/secrets/maskSecrets_test.go

@@ -5,12 +5,14 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"
 	"testing"
+
+	"github.com/rs/zerolog/log"
 )
 
 func TestSecretsDetection(t *testing.T) {
-
 	expectedResultsPath := "test/positive_expected_result.json"
 	expectedResults, err := os.ReadFile(expectedResultsPath)
 	if err != nil {
@@ -47,7 +49,7 @@ func TestSecretsDetection(t *testing.T) {
 	}
 
 	for _, ree := range rs {
-		if ree.Multiline.DetectLineGroup != 0 {
+		if ree.Multiline {
 			multiLineQueries = append(multiLineQueries, ree.QueryName)
 		}
 	}
@@ -58,7 +60,6 @@ func TestSecretsDetection(t *testing.T) {
 }
 
 func processFile(t *testing.T, path string, rs []SecretRegex, allowrs []*regexp.Regexp) []Result {
-
 	fileContent, err := os.ReadFile(path)
 	if err != nil {
 		t.Fatal(err)
@@ -70,8 +71,29 @@ func processFile(t *testing.T, path string, rs []SecretRegex, allowrs []*regexp.
 
 }
 
+func diffActualExpectedVulnerabilities(actual, expected []Result) []string {
+	m := make(map[string]bool)
+	diff := make([]string, 0)
+	for i := range expected {
+		m[expected[i].QueryName+":"+filepath.Base(expected[i].FileName)+":"+strconv.Itoa(expected[i].Line)] = true
+	}
+	for i := range actual {
+		if _, ok := m[actual[i].QueryName+":"+filepath.Base(actual[i].FileName)+":"+strconv.Itoa(actual[i].Line)]; !ok {
+			diff = append(diff, actual[i].FileName+":"+strconv.Itoa(actual[i].Line))
+		}
+	}
+
+	return diff
+}
+
 func compareExpectedWithActual(expected, actual []Result, multiLineQueries []string) bool {
 	if len(expected) != len(actual) {
+		log.Error().Msgf(
+			"Count of actual issues and expected vulnerabilities doesn't match\n -- \n"+
+				"not present in expected and present in actual: %v\n"+
+				"not present in actual and present in expected: %v\n",
+			diffActualExpectedVulnerabilities(actual, expected),
+			diffActualExpectedVulnerabilities(expected, actual))
 		return false
 	}