query-orchestration/cmd/cognito_test/user.creation.tool/cognito.go

// cognito.go handles AWS Cognito User Pool operations for user identity management.
// This file provides functions to create, delete, enable, disable, and retrieve users from
// AWS Cognito User Pools, including credential validation, user attribute management,
// and integration with AWS SDK v2 for identity provider operations.
package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

// Cognito types
type CognitoConfig struct {
	UserPoolID string
	Region     string
}

type UserInput struct {
	Email     string
	FirstName string
	LastName  string
}

type CognitoUserResponse struct {
	SubjectID string
	Username  string
	Email     string
	FirstName string
	LastName  string
	Status    string
	Enabled   bool
}

func sanityCheckCredentials() {
	awsProfile := os.Getenv("AWS_PROFILE")

	var cfgOpts []func(*config.LoadOptions) error
	if awsProfile != "" {
		cfgOpts = append(cfgOpts, config.WithSharedConfigProfile(awsProfile))
	}

	// Pass options with ... to expand the slice
	cfg, err := config.LoadDefaultConfig(context.Background(), cfgOpts...)
	if err != nil {
		fmt.Fprintf(os.Stderr, "AWS: Failed to load config: %v\n", err)
		os.Exit(1)
	}

	stsSvc := sts.NewFromConfig(cfg)
	_, err = stsSvc.GetCallerIdentity(context.Background(), &sts.GetCallerIdentityInput{})
	if err != nil {
		fmt.Fprintf(os.Stderr, "AWS: Credential check failed: %v\n", err)
		os.Exit(1)
	}

	permitKey := os.Getenv("PERMIT_KEY")
	if permitKey == "" {
		fmt.Fprintf(os.Stderr, "Permit.io: PERMIT_KEY environment variable is not set\n")
		os.Exit(1)
	}

	req, _ := http.NewRequest("GET", "https://api.permit.io/v2/projects", nil)
	req.Header.Add("Authorization", "Bearer "+permitKey)
	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil || resp.StatusCode < 200 || resp.StatusCode >= 300 {
		fmt.Fprintf(os.Stderr, "Permit.io: Credential check failed: %v, status=%v\n", err, resp.StatusCode)
		os.Exit(1)
	}
	resp.Body.Close()
}

// Initialize a Cognito client with the correct profile
func initializeCognitoClient() (*cognitoidentityprovider.Client, error) {
	awsProfile := os.Getenv("AWS_PROFILE")

	var cfgOpts []func(*config.LoadOptions) error
	if awsProfile != "" {
		cfgOpts = append(cfgOpts, config.WithSharedConfigProfile(awsProfile))
	}

	cfg, err := config.LoadDefaultConfig(context.Background(), cfgOpts...)
	if err != nil {
		return nil, err
	}

	client := cognitoidentityprovider.NewFromConfig(cfg)
	return client, nil
}

func createCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, config *CognitoConfig, user map[string]string) (*CognitoUserResponse, error) {
	// Use email as username
	username := user["email"]

	// Prepare user attributes
	userAttributes := []types.AttributeType{
		{
			Name:  aws.String("email"),
			Value: aws.String(user["email"]),
		},
		{
			Name:  aws.String("given_name"),
			Value: aws.String(user["first_name"]),
		},
		{
			Name:  aws.String("family_name"),
			Value: aws.String(user["last_name"]),
		},
		{
			Name:  aws.String("email_verified"),
			Value: aws.String("true"),
		},
	}

	// Prepare the AdminCreateUser input
	createUserInput := &cognitoidentityprovider.AdminCreateUserInput{
		UserPoolId:             aws.String(config.UserPoolID),
		Username:               aws.String(username),
		UserAttributes:         userAttributes,
		DesiredDeliveryMediums: []types.DeliveryMediumType{types.DeliveryMediumTypeEmail},
	}

	// Create the user
	result, err := client.AdminCreateUser(ctx, createUserInput)
	if err != nil {
		return nil, err
	}

	// Extract the subject ID from attributes
	var subjectID string
	for _, attr := range result.User.Attributes {
		if aws.ToString(attr.Name) == "sub" {
			subjectID = aws.ToString(attr.Value)
			break
		}
	}

	// Build response
	response := &CognitoUserResponse{
		SubjectID: subjectID,
		Username:  aws.ToString(result.User.Username),
		Email:     user["email"],
		FirstName: user["first_name"],
		LastName:  user["last_name"],
		Status:    string(result.User.UserStatus),
		Enabled:   result.User.Enabled,
	}

	return response, nil
}

func deleteCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) error {
	deleteInput := &cognitoidentityprovider.AdminDeleteUserInput{
		UserPoolId: aws.String(userPoolID),
		Username:   aws.String(username),
	}

	_, err := client.AdminDeleteUser(ctx, deleteInput)
	if err != nil {
		return err
	}

	return nil
}

// getCognitoUser retrieves an existing user from Cognito
func getCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) (*CognitoUserResponse, error) {
	input := &cognitoidentityprovider.AdminGetUserInput{
		UserPoolId: aws.String(userPoolID),
		Username:   aws.String(username),
	}

	result, err := client.AdminGetUser(ctx, input)
	if err != nil {
		return nil, err
	}

	// Extract attributes
	var subjectID, email, firstName, lastName string
	for _, attr := range result.UserAttributes {
		switch aws.ToString(attr.Name) {
		case "sub":
			subjectID = aws.ToString(attr.Value)
		case "email":
			email = aws.ToString(attr.Value)
		case "given_name":
			firstName = aws.ToString(attr.Value)
		case "family_name":
			lastName = aws.ToString(attr.Value)
		}
	}

	return &CognitoUserResponse{
		SubjectID: subjectID,
		Username:  aws.ToString(result.Username),
		Email:     email,
		FirstName: firstName,
		LastName:  lastName,
		Status:    string(result.UserStatus),
		Enabled:   result.Enabled,
	}, nil
}

// userExistsInCognito checks if a user exists in Cognito
func userExistsInCognito(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) bool {
	_, err := getCognitoUser(ctx, client, userPoolID, username)
	return err == nil
}

// disableCognitoUser disables a user in Cognito
func disableCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) error {
	input := &cognitoidentityprovider.AdminDisableUserInput{
		UserPoolId: aws.String(userPoolID),
		Username:   aws.String(username),
	}

	_, err := client.AdminDisableUser(ctx, input)
	return err
}

// enableCognitoUser enables a user in Cognito
func enableCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) error {
	input := &cognitoidentityprovider.AdminEnableUserInput{
		UserPoolId: aws.String(userPoolID),
		Username:   aws.String(username),
	}

	_, err := client.AdminEnableUser(ctx, input)
	return err
}

// isCognitoUserEnabled checks if a user is enabled in Cognito
func isCognitoUserEnabled(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) (bool, error) {
	user, err := getCognitoUser(ctx, client, userPoolID, username)
	if err != nil {
		return false, err
	}
	return user.Enabled, nil
}