get-unique-id-func-refactored

Author: cx-hitesh-madgulkar

internal/wrappers/client.go

@@ -12,7 +12,6 @@ import (
 	"net/http/httptrace"
 	"net/url"
 	"os"
-	"os/user"
 	"runtime"
 	"strings"
 	"sync"
@@ -21,13 +20,11 @@ import (
 	applicationErrors "github.com/checkmarx/ast-cli/internal/constants/errors"
 	"github.com/checkmarx/ast-cli/internal/logger"
 	"github.com/golang-jwt/jwt/v5"
-	"github.com/google/uuid"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/viper"
 
 	commonParams "github.com/checkmarx/ast-cli/internal/params"
-	"github.com/checkmarx/ast-cli/internal/wrappers/configuration"
 	"github.com/checkmarx/ast-cli/internal/wrappers/kerberos"
 	"github.com/checkmarx/ast-cli/internal/wrappers/ntlm"
 )
@@ -984,85 +981,3 @@ func extractAZPFromToken(astToken string) (string, error) {
 	}
 	return azp, nil
 }
-
-func GetUniqueID() string {
-	var uniqueID string
-	isAllowed := false
-	accessToken, err := GetAccessToken()
-	if err != nil {
-		logger.PrintIfVerbose("Failed to get access token")
-		return ""
-	}
-	parser := jwt.NewParser(jwt.WithoutClaimsValidation())
-	token, _, err := parser.ParseUnverified(accessToken, jwt.MapClaims{})
-	if err != nil {
-		logger.PrintIfVerbose("Failed to parse JWT token " + err.Error())
-		return ""
-	}
-	claims, ok := token.Claims.(jwt.MapClaims)
-	if !ok {
-		logger.PrintIfVerbose("Failed to get JWT claims")
-		return ""
-	}
-
-	astLicense, ok := claims["ast-license"].(map[string]interface{})
-	if !ok {
-		logger.PrintIfVerbose("Failed to get ast-license from claims")
-		return ""
-	}
-
-	licenseData, ok := astLicense["LicenseData"].(map[string]interface{})
-	if !ok {
-		logger.PrintIfVerbose("Failed to get LicenseData from ast-license")
-		return ""
-	}
-
-	allowedEngines, ok := licenseData["allowedEngines"].([]interface{})
-	if !ok {
-		logger.PrintIfVerbose("Failed to get allowedEngines from LicenseData")
-		return ""
-	}
-
-	for _, engine := range allowedEngines {
-		engineStr, ok := engine.(string)
-		if !ok {
-			continue
-		}
-		if strings.EqualFold(engineStr, "Checkmarx Developer Assist") {
-			isAllowed = true
-			break
-		}
-	}
-
-	if !isAllowed {
-		logger.PrintIfVerbose("User does not not have permission to standalone dev asists feature")
-		return ""
-	}
-	uniqueID = viper.GetString(commonParams.UniqueIDConfigKey)
-	if uniqueID != "" {
-		return uniqueID
-	}
-	logger.PrintIfVerbose("Generating new unique id")
-	currentUser, err := user.Current()
-	if err != nil {
-		logger.PrintIfVerbose("Failed to get user: " + err.Error())
-		return ""
-	}
-	username := currentUser.Username
-	username = strings.TrimSpace(username)
-	logger.PrintIfVerbose("Username to be used for unique id: " + username)
-	if strings.Contains(username, "\\") {
-		username = strings.Split(username, "\\")[1]
-	}
-	uniqueID = uuid.New().String() + "_" + username
-
-	logger.PrintIfVerbose("Unique id: " + uniqueID)
-	viper.Set(commonParams.UniqueIDConfigKey, uniqueID)
-	configFilePath, _ := configuration.GetConfigFilePath()
-	err = configuration.SafeWriteSingleConfigKeyString(configFilePath, commonParams.UniqueIDConfigKey, uniqueID)
-	if err != nil {
-		logger.PrintIfVerbose("Failed to write config: " + err.Error())
-		return ""
-	}
-	return uniqueID
-}

internal/wrappers/jwt-helper.go

@@ -1,12 +1,17 @@
 package wrappers
 
 import (
+	"os/user"
 	"strings"
 
+	"github.com/checkmarx/ast-cli/internal/logger"
 	commonParams "github.com/checkmarx/ast-cli/internal/params"
+	"github.com/checkmarx/ast-cli/internal/wrappers/configuration"
 	"github.com/checkmarx/ast-cli/internal/wrappers/utils"
 	"github.com/golang-jwt/jwt/v5"
+	"github.com/google/uuid"
 	"github.com/pkg/errors"
+	"github.com/spf13/viper"
 )
 
 // JWTStruct model used to get all jwt fields
@@ -158,3 +163,49 @@ func (*JWTStruct) CheckPermissionByAccessToken(requiredPermission string) (hasPe
 	}
 	return permission, nil
 }
+
+func GetUniqueID() string {
+	var uniqueID string
+	// Check License first
+	jwtWrapper := NewJwtWrapper()
+	isAllowed, err := jwtWrapper.IsAllowedEngine("Checkmarx Developer Assist")
+	if err != nil {
+		logger.PrintIfVerbose("Failed to check engine allowance: " + err.Error())
+		return ""
+	}
+	if !isAllowed {
+		logger.PrintIfVerbose("User does not have permission to standalone dev assists feature")
+		return ""
+	}
+
+	// Check if unique id is already set
+	uniqueID = viper.GetString(commonParams.UniqueIDConfigKey)
+	if uniqueID != "" {
+		return uniqueID
+	}
+
+	// Generate new unique id
+	logger.PrintIfVerbose("Generating new unique id")
+	currentUser, err := user.Current()
+	if err != nil {
+		logger.PrintIfVerbose("Failed to get user: " + err.Error())
+		return ""
+	}
+	username := currentUser.Username
+	username = strings.TrimSpace(username)
+	logger.PrintIfVerbose("Username to be used for unique id: " + username)
+	if strings.Contains(username, "\\") {
+		username = strings.Split(username, "\\")[1]
+	}
+	uniqueID = uuid.New().String() + "_" + username
+
+	logger.PrintIfVerbose("Unique id: " + uniqueID)
+	viper.Set(commonParams.UniqueIDConfigKey, uniqueID)
+	configFilePath, _ := configuration.GetConfigFilePath()
+	err = configuration.SafeWriteSingleConfigKeyString(configFilePath, commonParams.UniqueIDConfigKey, uniqueID)
+	if err != nil {
+		logger.PrintIfVerbose("Failed to write config: " + err.Error())
+		return ""
+	}
+	return uniqueID
+}

internal/wrappers/jwt-helper_test.go

@@ -1,8 +1,11 @@
 package wrappers
 
 import (
+	"strings"
 	"testing"
 
+	commonParams "github.com/checkmarx/ast-cli/internal/params"
+	"github.com/spf13/viper"
 	"gotest.tools/assert"
 )
 
@@ -71,3 +74,44 @@ func TestGetEnabledEngines(t *testing.T) {
 		})
 	}
 }
+
+func TestGetUniqueID(t *testing.T) {
+	// Save original value and restore after test
+	originalID := viper.GetString(commonParams.UniqueIDConfigKey)
+	defer viper.Set(commonParams.UniqueIDConfigKey, originalID)
+
+	t.Run("returns existing unique ID from config", func(t *testing.T) {
+		// Setup: set existing ID
+		existingID := "test-uuid-456_testuser"
+		viper.Set(commonParams.UniqueIDConfigKey, existingID)
+
+		result := GetUniqueID()
+
+		if result != "" {
+			assert.Equal(t, existingID, result)
+		} else {
+			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' license")
+		}
+	})
+
+	t.Run("generates new unique ID when none exists", func(t *testing.T) {
+		// Setup: clear existing ID
+		viper.Set(commonParams.UniqueIDConfigKey, "")
+
+		result := GetUniqueID()
+
+		if result == "" {
+			t.Skip("Requires valid auth and 'Checkmarx Developer Assist' license")
+			return
+		}
+
+		// Verify format: UUID_username
+		assert.Assert(t, strings.Contains(result, "_"), "Should have UUID_username format")
+		assert.Assert(t, len(result) > 36, "Should contain UUID and username")
+
+		// Verify no backslash (Windows domain stripped)
+		parts := strings.Split(result, "_")
+		assert.Assert(t, len(parts) >= 2, "Should have at least 2 parts")
+		assert.Assert(t, !strings.Contains(parts[1], "\\"), "Username should not contain backslash")
+	})
+}