query-orchestration/internal/rbac/cognito.go

package rbac

import (
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// KeyProvider interface for different key sources
type KeyProvider interface {
	GetPublicKey(kid string) (*rsa.PublicKey, error)
}

// LocalKeyProvider implements KeyProvider using local PEM files
type LocalKeyProvider struct {
	PrivateKeyPath string
	PublicKeyPath  string
}

// CognitoKeyProvider implements KeyProvider using AWS Cognito JWKS endpoint
type CognitoKeyProvider struct {
	Region     string
	UserPoolID string
	JwksCache  map[string]*rsa.PublicKey // Cache for JWKS keys by Key ID
}

// JWKS structure for Cognito public keys
type JWKS struct {
	Keys []struct {
		Kid string `json:"kid"`
		Kty string `json:"kty"`
		N   string `json:"n"`
		E   string `json:"e"`
		Use string `json:"use"`
		Alg string `json:"alg"`
	} `json:"keys"`
}

// TokenDetails contains the extracted information from a JWT token
type TokenDetails struct {
	Token      *jwt.Token
	Claims     jwt.MapClaims
	UserID     string
	Email      string
	Groups     []string
	Expiration time.Time
}

// Load private key from file
func LoadPrivateKey(keyPath string) (*rsa.PrivateKey, error) {
	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("error reading private key file: %w", err)
	}
	return jwt.ParseRSAPrivateKeyFromPEM(keyData)
}

// Load public key from file
func LoadPublicKey(keyPath string) (*rsa.PublicKey, error) {
	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("error reading public key file: %w", err)
	}
	return jwt.ParseRSAPublicKeyFromPEM(keyData)
}

// GetPublicKey implements KeyProvider for LocalKeyProvider
func (lkp *LocalKeyProvider) GetPublicKey(kid string) (*rsa.PublicKey, error) {
	// For local provider, kid is ignored since we only have one key
	return LoadPublicKey(lkp.PublicKeyPath)
}

// GetPrivateKey returns the private key for signing test tokens
func (lkp *LocalKeyProvider) GetPrivateKey() (*rsa.PrivateKey, error) {
	return LoadPrivateKey(lkp.PrivateKeyPath)
}

// jwkToPublicKey converts JWK components to RSA public key
func jwkToPublicKey(n, e string) (*rsa.PublicKey, error) {
	// Decode the base64 URL-encoded modulus
	nBytes, err := base64.RawURLEncoding.DecodeString(n)
	if err != nil {
		return nil, fmt.Errorf("error decoding modulus: %w", err)
	}

	// Decode the base64 URL-encoded exponent
	eBytes, err := base64.RawURLEncoding.DecodeString(e)
	if err != nil {
		return nil, fmt.Errorf("error decoding exponent: %w", err)
	}

	// Convert modulus bytes to big.Int
	modulus := new(big.Int)
	modulus.SetBytes(nBytes)

	// Convert exponent bytes to int
	var exponent int
	for i := 0; i < len(eBytes); i++ {
		exponent = exponent<<8 + int(eBytes[i])
	}

	// Create RSA public key
	return &rsa.PublicKey{
		N: modulus,
		E: exponent,
	}, nil
}

// GetPublicKey implements KeyProvider for CognitoKeyProvider
func (ckp *CognitoKeyProvider) GetPublicKey(kid string) (*rsa.PublicKey, error) {
	// Check cache first
	if ckp.JwksCache == nil {
		ckp.JwksCache = make(map[string]*rsa.PublicKey)
	}

	if key, exists := ckp.JwksCache[kid]; exists {
		return key, nil
	}

	// Fetch JWKS from Cognito
	//baseURL := "https://cognito-idp.amazonaws.com"
	//jwksPath := fmt.Sprintf("/%s/%s/.well-known/jwks.json", ckp.Region, ckp.UserPoolID)
	jwksURL := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", ckp.Region, ckp.UserPoolID)

	if err := IsValidJWKSURL(jwksURL); err != nil {
		return nil, fmt.Errorf("invalid JWKS URL: %w", err)
	}

	client := &http.Client{
		Timeout: 10 * time.Second,
	}
	req, err := http.NewRequest("GET", jwksURL, nil)
	if err != nil {
		return nil, fmt.Errorf("error creating request for JWKS: %w", err)
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("error fetching JWKS from Cognito: %w", err)
	}
	defer resp.Body.Close()

	jwksData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("error reading JWKS response: %w", err)
	}

	var jwks JWKS
	if err := json.Unmarshal(jwksData, &jwks); err != nil {
		return nil, fmt.Errorf("error parsing JWKS JSON: %w", err)
	}

	// Find the key with matching kid
	for _, key := range jwks.Keys {
		if key.Kid == kid {
			// Convert JWK to RSA public key
			publicKey, err := jwkToPublicKey(key.N, key.E)
			if err != nil {
				return nil, fmt.Errorf("error converting JWK to public key: %w", err)
			}

			// Store in cache
			ckp.JwksCache[kid] = publicKey
			return publicKey, nil
		}
	}

	return nil, fmt.Errorf("no matching key found with kid: %s", kid)
}

// NewLocalKeyProvider creates a new LocalKeyProvider
func NewLocalKeyProvider(privateKeyPath, publicKeyPath string) *LocalKeyProvider {
	return &LocalKeyProvider{
		PrivateKeyPath: privateKeyPath,
		PublicKeyPath:  publicKeyPath,
	}
}

// NewCognitoKeyProvider creates a new CognitoKeyProvider
func NewCognitoKeyProvider(region, userPoolID string) *CognitoKeyProvider {
	return &CognitoKeyProvider{
		Region:     region,
		UserPoolID: userPoolID,
		JwksCache:  make(map[string]*rsa.PublicKey),
	}
}

// GenerateMockJWT creates a test JWT with given groups
func GenerateMockJWT(keyProvider *LocalKeyProvider, userID, email string, groups []string) (string, error) {
	privateKey, err := keyProvider.GetPrivateKey()
	if err != nil {
		return "", err
	}

	// Create a header with kid value (important for verification)
	headers := map[string]interface{}{
		"kid": "test-key-id", // In a real scenario this would match the key ID from Cognito
		"alg": "RS256",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
		"sub":            userID,
		"email":          email,
		"exp":            time.Now().Add(time.Hour).Unix(),
		"iat":            time.Now().Unix(),
		"iss":            "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_XXXXXXXXX",
		"cognito:groups": groups,
	})

	// Set the header
	for k, v := range headers {
		token.Header[k] = v
	}

	return token.SignedString(privateKey)
}

// PrintTokenDetails prints the JWT token details for debugging
func PrintTokenDetails(token *jwt.Token) {
	fmt.Println("\n=== JWT Token Details ===")

	// Print Header
	headerJSON, _ := json.MarshalIndent(token.Header, "", "  ")
	fmt.Printf("\nHeader:\n%s\n", string(headerJSON))

	// Print Claims
	claimsJSON, _ := json.MarshalIndent(token.Claims, "", "  ")
	fmt.Printf("\nClaims:\n%s\n", string(claimsJSON))

	// Print Signature
	fmt.Printf("\nSignature: %s\n", token.Signature)

	// Print Method
	fmt.Printf("\nSigning Method: %s\n", token.Method.Alg())

	// Print Valid status
	fmt.Printf("\nValid: %v\n", token.Valid)

	fmt.Println("\n=====================")
}

// ValidateJWT validates a JWT token and returns the claims
func ValidateJWT(tokenString string, keyProvider KeyProvider) (*jwt.Token, jwt.MapClaims, error) {
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, errors.New("unexpected signing method")
		}

		// Get key ID from token header
		kid, ok := token.Header["kid"].(string)
		if !ok {
			return nil, errors.New("no key ID (kid) in token header")
		}

		// Get public key from provider
		return keyProvider.GetPublicKey(kid)
	})

	if err != nil {
		return nil, nil, err
	}

	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		return token, claims, nil
	}
	return nil, nil, errors.New("invalid token")
}

// ExtractTokenDetails extracts and organizes token information
func ExtractTokenDetails(token *jwt.Token, claims jwt.MapClaims) *TokenDetails {
	details := &TokenDetails{
		Token:  token,
		Claims: claims,
	}

	if sub, ok := claims["sub"].(string); ok {
		details.UserID = sub
	}

	if email, ok := claims["email"].(string); ok {
		details.Email = email
	}

	if expFloat, ok := claims["exp"].(float64); ok {
		details.Expiration = time.Unix(int64(expFloat), 0)
	}

	if groupsInterface, ok := claims["cognito:groups"].([]interface{}); ok {
		details.Groups = make([]string, 0, len(groupsInterface))
		for _, group := range groupsInterface {
			if groupStr, ok := group.(string); ok {
				details.Groups = append(details.Groups, groupStr)
			}
		}
	}

	return details
}

// HasGroup checks if a user belongs to a specific group
func HasGroup(claims jwt.MapClaims, requiredGroup string) bool {
	if groups, ok := claims["cognito:groups"].([]interface{}); ok {
		for _, group := range groups {
			if groupStr, ok := group.(string); ok && strings.EqualFold(groupStr, requiredGroup) {
				return true
			}
		}
	}
	return false
}

// HasAnyGroup checks if a user belongs to any of the specified groups
func HasAnyGroup(claims jwt.MapClaims, requiredGroups []string) bool {
	for _, group := range requiredGroups {
		if HasGroup(claims, group) {
			return true
		}
	}
	return false
}

// HasAllGroups checks if a user belongs to all of the specified groups
func HasAllGroups(claims jwt.MapClaims, requiredGroups []string) bool {
	for _, group := range requiredGroups {
		if !HasGroup(claims, group) {
			return false
		}
	}
	return true
}

// IsValidJWKSURL validates that a JWKS URL is secure and from a trusted domain
func IsValidJWKSURL(urlStr string) error {
	parsed, err := url.Parse(urlStr)
	if err != nil {
		return fmt.Errorf("invalid URL format: %w", err)
	}

	// Add your trusted domains here
	trustedDomains := []string{
		"cognito-idp.us-east-1.amazonaws.com",
		"cognito-idp.us-west-2.amazonaws.com",
		// Add other AWS Cognito regions as needed
	}

	isAllowed := false
	for _, domain := range trustedDomains {
		if parsed.Host == domain {
			isAllowed = true
			break
		}
	}

	if !isAllowed {
		return fmt.Errorf("untrusted JWKS domain: %s", parsed.Host)
	}

	// Ensure HTTPS
	if parsed.Scheme != "https" {
		return fmt.Errorf("JWKS URL must use HTTPS")
	}

	return nil
}

// VerifyGroupMembership verifies a user belongs to specified groups
func VerifyGroupMembership(claims jwt.MapClaims, expectedGroups []string) error {
	groups, ok := claims["cognito:groups"].([]interface{})
	if !ok {
		return fmt.Errorf("expected 'cognito:groups' claim, but not found")
	}

	for _, eg := range expectedGroups {
		found := false
		for _, g := range groups {
			if gStr, ok := g.(string); ok && strings.EqualFold(gStr, eg) {
				found = true
				break
			}
		}
		if !found {
			return fmt.Errorf("expected group %s not found in token", eg)
		}
	}

	return nil
}