query-orchestration/internal/usermanagement/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 user attribute management and integration with AWS SDK v2.
package usermanagement

import (
	"context"
	"os"

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

// CreateCognitoUser creates a new user in AWS Cognito User Pool with the specified attributes.
// It uses the email as the username and sets the email as verified by default.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - config: Cognito configuration containing User Pool ID and region
//   - attrs: User attributes including email, first name, and last name
//
// Returns:
//   - *CognitoUserResponse: The created user's details including Subject ID
//   - error: Error if the user creation fails
func CreateCognitoUser(ctx context.Context, client *cognitoidentityprovider.Client, config *CognitoConfig, attrs UserAttributes) (*CognitoUserResponse, error) {
	// Use email as username
	username := attrs.Email

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

	// Tag the user with the environment ID if specified
	if attrs.EnvironmentID != "" {
		userAttributes = append(userAttributes, types.AttributeType{
			Name:  aws.String(CognitoEnvIDAttrName),
			Value: aws.String(attrs.EnvironmentID),
		})
	}

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

	// Suppress welcome email if COGNITO_SUPPRESS_EMAILS environment variable is set to "true"
	// This is useful for integration tests to avoid hitting AWS Cognito's 50 emails/day limit
	if os.Getenv("COGNITO_SUPPRESS_EMAILS") == "true" {
		createUserInput.MessageAction = types.MessageActionTypeSuppress
	}

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

	// Extract attributes from the response
	var subjectID, environmentID string
	for _, attr := range result.User.Attributes {
		switch aws.ToString(attr.Name) {
		case "sub":
			subjectID = aws.ToString(attr.Value)
		case CognitoEnvIDAttrName:
			environmentID = aws.ToString(attr.Value)
		}
	}

	// Build response
	response := &CognitoUserResponse{
		SubjectID:     subjectID,
		Username:      aws.ToString(result.User.Username),
		Email:         attrs.Email,
		FirstName:     attrs.FirstName,
		LastName:      attrs.LastName,
		Status:        string(result.User.UserStatus),
		Enabled:       result.User.Enabled,
		EnvironmentID: environmentID,
	}

	return response, nil
}

// GetCognitoUser retrieves an existing user from AWS Cognito User Pool by username.
// The username is typically the user's email address.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to retrieve (typically email)
//
// Returns:
//   - *CognitoUserResponse: The user's details including Subject ID and attributes
//   - error: Error if the user retrieval fails or user doesn't exist
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, environmentID 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)
		case CognitoEnvIDAttrName:
			environmentID = 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,
		EnvironmentID: environmentID,
	}, nil
}

// UserExistsInCognito checks if a user exists in AWS Cognito User Pool by username.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to check (typically email)
//
// Returns:
//   - bool: true if the user exists, false otherwise
func UserExistsInCognito(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string) bool {
	_, err := GetCognitoUser(ctx, client, userPoolID, username)
	return err == nil
}

// DeleteCognitoUser permanently deletes a user from AWS Cognito User Pool.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to delete (typically email)
//
// Returns:
//   - error: Error if the deletion fails
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)
	return err
}

// DisableCognitoUser disables a user in AWS Cognito User Pool, preventing them from signing in.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to disable (typically email)
//
// Returns:
//   - error: Error if the disable operation fails
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 previously disabled user in AWS Cognito User Pool.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to enable (typically email)
//
// Returns:
//   - error: Error if the enable operation fails
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 currently enabled in AWS Cognito User Pool.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to check (typically email)
//
// Returns:
//   - bool: true if the user is enabled, false if disabled
//   - error: Error if the check fails or user doesn't exist
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
}

// UpdateCognitoUserAttributes updates user attributes in AWS Cognito User Pool.
// Only updates attributes that are provided (non-empty).
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - username: The username to update (typically email)
//   - attrs: User attributes to update (only non-empty fields are updated)
//
// Returns:
//   - error: Error if the update fails
func UpdateCognitoUserAttributes(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID, username string, attrs UserAttributes) error {
	var userAttributes []types.AttributeType

	if attrs.FirstName != "" {
		userAttributes = append(userAttributes, types.AttributeType{
			Name:  aws.String("given_name"),
			Value: aws.String(attrs.FirstName),
		})
	}

	if attrs.LastName != "" {
		userAttributes = append(userAttributes, types.AttributeType{
			Name:  aws.String("family_name"),
			Value: aws.String(attrs.LastName),
		})
	}

	if attrs.EnvironmentID != "" {
		userAttributes = append(userAttributes, types.AttributeType{
			Name:  aws.String(CognitoEnvIDAttrName),
			Value: aws.String(attrs.EnvironmentID),
		})
	}

	if len(userAttributes) == 0 {
		return nil
	}

	input := &cognitoidentityprovider.AdminUpdateUserAttributesInput{
		UserPoolId:     aws.String(userPoolID),
		Username:       aws.String(username),
		UserAttributes: userAttributes,
	}

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

// ListCognitoUsersResult represents the result of listing Cognito users with pagination.
type ListCognitoUsersResult struct {
	Users         []CognitoUserResponse
	PaginationKey *string
}

// ListCognitoUsers lists users from AWS Cognito User Pool with pagination support.
//
// Parameters:
//   - ctx: Context for the operation
//   - client: AWS Cognito Identity Provider client
//   - userPoolID: The Cognito User Pool ID
//   - limit: Maximum number of users to return (max 60)
//   - paginationToken: Token for pagination (nil for first page)
//
// Returns:
//   - *ListCognitoUsersResult: Result containing users and next pagination token
//   - error: Error if the list operation fails
func ListCognitoUsers(ctx context.Context, client *cognitoidentityprovider.Client, userPoolID string, limit int, paginationToken *string) (*ListCognitoUsersResult, error) {
	// Cognito max limit is 60
	if limit > 60 {
		limit = 60
	}
	if limit < 1 {
		limit = 60
	}

	// Ensure limit is within safe bounds for int32 conversion
	var limitInt32 int32
	if limit > 60 || limit < 0 {
		limitInt32 = 60
	} else {
		limitInt32 = int32(limit) // #nosec G115 -- limit is validated to be within 0-60 range
	}

	input := &cognitoidentityprovider.ListUsersInput{
		UserPoolId: aws.String(userPoolID),
		Limit:      aws.Int32(limitInt32),
	}

	if paginationToken != nil && *paginationToken != "" {
		input.PaginationToken = paginationToken
	}

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

	// Convert Cognito users to our response format
	users := make([]CognitoUserResponse, 0, len(result.Users))
	for _, cognitoUser := range result.Users {
		var subjectID, email, firstName, lastName, environmentID string
		for _, attr := range cognitoUser.Attributes {
			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)
			case CognitoEnvIDAttrName:
				environmentID = aws.ToString(attr.Value)
			}
		}

		users = append(users, CognitoUserResponse{
			SubjectID:     subjectID,
			Username:      aws.ToString(cognitoUser.Username),
			Email:         email,
			FirstName:     firstName,
			LastName:      lastName,
			Status:        string(cognitoUser.UserStatus),
			Enabled:       cognitoUser.Enabled,
			EnvironmentID: environmentID,
		})
	}

	return &ListCognitoUsersResult{
		Users:         users,
		PaginationKey: result.PaginationToken,
	}, nil
}