This commit is contained in:
jay brown
2025-03-18 16:05:17 -07:00
parent 3fef82354a
commit 78a0cea7b7
+73 -57
View File
@@ -3,91 +3,107 @@ package main
import (
"fmt"
"log"
"queryorchestration/internal/rbac" // Replace with your actual module name
"queryorchestration/internal/rbac"
"queryorchestration/internal/serviceconfig"
"github.com/golang-jwt/jwt/v5"
)
// for testing run like `LOG_LEVEL=INFO go run main.go`
func main() {
func setupConfig() (*serviceconfig.BaseConfig, error) {
cfg := &serviceconfig.BaseConfig{}
errGettingConfig := serviceconfig.InitializeConfig(cfg)
if errGettingConfig != nil {
log.Fatal(errGettingConfig)
if err := serviceconfig.InitializeConfig(cfg); err != nil {
return nil, fmt.Errorf("failed to initialize config: %w", err)
}
errorInitializingAuthProvider := cfg.InitializeAuthProvider()
if errorInitializingAuthProvider != nil {
log.Fatal(errorInitializingAuthProvider)
if err := cfg.InitializeAuthProvider(); err != nil {
return nil, fmt.Errorf("failed to initialize auth provider: %w", err)
}
cfg.PrintConfig("secret")
// use the config here for the key provider parameters
// Example with local key provider for testing
if cfg.AuthProvider == nil {
log.Fatal("AuthProvider is nil")
return nil, fmt.Errorf("auth provider is nil")
}
localProvider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
return cfg, nil
}
// Generate a mock JWT with multiple Cognito groups
func generateAndValidateToken(localProvider *rbac.LocalKeyProvider) (*rbac.TokenDetails, error) {
groups := []string{"admins", "developers", "querybuilders", "uploaders", "exporters"}
tokenString, err := rbac.GenerateTestJWT(localProvider, "test-user-id", "test@example.com", groups)
if err != nil {
fmt.Println("Error generating token:", err)
return
return nil, fmt.Errorf("error generating token: %w", err)
}
fmt.Printf("raw token: %s\n", tokenString)
token, claims, err := rbac.ValidateJWT(tokenString, localProvider)
if err != nil {
fmt.Println("Token validation failed:", err)
return nil, fmt.Errorf("token validation failed: %w", err)
}
if !token.Valid {
return nil, fmt.Errorf("invalid token")
}
return rbac.ExtractTokenDetails(token, claims), nil
}
func checkUserGroups(claims jwt.MapClaims) {
groupsInterface, ok := claims["cognito:groups"]
if !ok {
fmt.Println("No groups found in token.")
return
}
// Print token details for debugging
rbac.PrintTokenDetails(token)
groups, ok := groupsInterface.([]interface{})
if !ok {
fmt.Println("Invalid groups format in token.")
return
}
// Example with Cognito key provider for production
/*
cognitoProvider := rbac.NewCognitoKeyProvider("us-east-1", "us-east-1_XXXXXXXXX")
token, claims, err := rbac.ValidateJWT(tokenString, cognitoProvider)
*/
if token.Valid {
// Extract the token details
details := rbac.ExtractTokenDetails(token, claims)
fmt.Println("Token is valid.")
fmt.Println("User ID:", details.UserID)
fmt.Println("Email:", details.Email)
fmt.Println("Expires:", details.Expiration)
if len(details.Groups) > 0 {
fmt.Println("User Groups:")
for _, group := range details.Groups {
fmt.Println("-", group)
if len(groups) > 0 {
fmt.Println("User Groups:")
for _, group := range groups {
if groupStr, ok := group.(string); ok {
fmt.Println("-", groupStr)
}
} else {
fmt.Println("No groups found in token.")
}
// Check if user is in a specific group
if rbac.HasGroup(claims, "admins") {
fmt.Println("User is an admin.")
}
// Check if user is in any of these groups
if rbac.HasAnyGroup(claims, []string{"developers", "querybuilders"}) {
fmt.Println("User can modify or query data.")
}
// Check if user is in all of these groups
if rbac.HasAllGroups(claims, []string{"uploaders", "exporters"}) {
fmt.Println("User can both upload and export data.")
}
} else {
fmt.Println("Invalid token")
fmt.Println("No groups found in token.")
}
if rbac.HasGroup(claims, "admins") {
fmt.Println("User is an admin.")
}
if rbac.HasAnyGroup(claims, []string{"developers", "querybuilders"}) {
fmt.Println("User can modify or query data.")
}
if rbac.HasAllGroups(claims, []string{"uploaders", "exporters"}) {
fmt.Println("User can both upload and export data.")
}
}
// for testing run like `LOG_LEVEL=INFO go run main.go`
func main() {
_, err := setupConfig()
if err != nil {
log.Fatal(err)
}
localProvider := rbac.NewLocalKeyProvider("private_key.pem", "public_key.pem")
details, err := generateAndValidateToken(localProvider)
if err != nil {
log.Fatal(err)
}
fmt.Println("Token is valid.")
fmt.Println("User ID:", details.UserID)
fmt.Println("Email:", details.Email)
fmt.Println("Expires:", details.Expiration)
checkUserGroups(details.Claims)
}