In many customer facing applications, whether e-commerce sites or SaaS platforms, the concept of user impersonation can become really important when these applications grow, especially for customer support. It's typically a straight forward mechanism to implement within an application, however, I ran into a brick wall with AWS Cognito. 

AWS Cognito is AWS's answer to identity management for B2C and B2B applications, with support for all standard authentication flows and protocols, including SAML and OpenID. Although Cognito is great and takes all the pain away from managing authentication ourselves, it's quite a generic tool - this means that if we wanted to implement something that isn't supported out-of-the-box, we have to get creative.

So, I've written this brief guide to explain the implementation details and niches of implementing a user impersonation flow with AWS Cognito. 

## What is user impersonation?

User impersonation is a method which can be used, typically by an admin or a customer support user, to impersonate a user and gain access to their account. This allows an admin user to view an application as a specified user, without requiring the user to share their password or other credentials.

## Why might we need it?

There are quite a few use cases where user impersonation could be helpful. However, a prime example could be: a user is having issues with their account, or experiencing a specific bug in the application. A customer support user could use user impersonation to gain access to the user's account, allowing them to view the application exactly how the user would - giving them better visibility of an issue.

## How do we do it?

As expected, Amazon have designed Cognito with security best pratises in mind. And with that, they have made a seemingly conscious decision to not implement a user impersonation feature, which means if your application has a requirement for that, it becomes quite awkward.

However, there is still a way this can be done. Cognito's user pools support integration with AWS Lambda, allowing custom logic to be implemented throughout the authentication flow. So where do we start?

### Creating the Cognito User Pool

To get started, the first thing we will need is a User Pool in Cognito. This will be where our users records, which includes any custom user attributes etc. A user pool can be created in Cognito with the Terraform code below:

```hcl
# ./main.tf

resource "aws_cognito_user_pool" "default" {
  name = "app-users"
}
```

User pools can be configured in a number of ways to best suit an application, but to get started we can just give it a name. As we now have a user pool, we will now need a user pool client which we can authenticate with. Typically, I would have a user pool client which will be used by users when logging into the application, then another for the custom user impersonation flow. But for now, lets just configure one for the custom flow.

```hcl
# ./main.tf

resource "aws_cognito_user_pool_client" "admin" {
  name                = "user-impersonation"
  user_pool_id        = aws_cognito_user_pool.default.id
  generate_secret     = true // Ensure a client secret is generated.

  // Allow the "CUSTOM_AUTH" flow. Also allow refresh token auth, as it's
  // required, but we won't use it.
  explicit_auth_flows = ["ALLOW_CUSTOM_AUTH", "ALLOW_REFRESH_TOKEN_AUTH"]
}
```

### Creating the define auth challenge Lambda function

Now that we have a user pool configured, we can start to implement the code to support the custom authentication flow. The first step of this, is to create an AWS Lambda function which will be used to define the authentication flow for the user impersonation flow.

AWS have put together a useful set of documentation on the custom auth flow support of Cognito, [here](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-challenge.html), which shows how it can be utilised for different use cases.

In this example, we will be using Golang for our application code, although Lambda supports a wide range of languages. This function will be quite simple and just require one Go file, but before we get to that we will need to initialise a new Go module, and install the AWS SDKs.

```bash
# ./

$ go mod init custom-auth
```

Once the module has been initialised, we can install the AWS modules required for our Lambda function.

```bash
# ./

$ go get -u github.com/aws/aws-lambda-go/events
$ go get -u github.com/aws/aws-lambda-go/lambda
```

We now have everything we need, and can start with the function code.

```go
# ./cmd/define-challenge-function/main.go

package main

import (
	"context"
	"errors"

	"github.com/aws/aws-lambda-go/events"
	"github.com/aws/aws-lambda-go/lambda"
)

// The lambda function handler code.
func handler(ctx context.Context, event events.CognitoEventUserPoolsDefineAuthChallenge) (events.CognitoEventUserPoolsDefineAuthChallenge, error) {
	// Ensure our user exists in the user pool.
	if event.Request.UserNotFound {
		event.Response.IssueTokens = false
		event.Response.FailAuthentication = true
		return event, errors.New("user not found")
	}

	count := len(event.Request.Session)
	challenge := (*events.CognitoEventUserPoolsChallengeResult)(nil)
	if count > 0 {
		// If there is a challenge in the session, get the most recent one.
		challenge = event.Request.Session[count-1]
	}

	// Firstly, check if the custom auth challenge has been verified, then allow the
	// issuing of authentication tokens.
	if challenge != nil && challenge.ChallengeName == "CUSTOM_CHALLENGE" && challenge.ChallengeResult {
		event.Response.IssueTokens = true
		event.Response.FailAuthentication = false
	} else { // Otherwise, skip and forward to the custom challenge.
		event.Response.IssueTokens = false
		event.Response.FailAuthentication = false
		event.Response.ChallengeName = "CUSTOM_CHALLENGE"
	}

	return event, nil
}

func main() {
    // Lambda start here!
	lambda.Start(handler)
}
```

The code above has two main aspects: the handler function, and the "main" entry-point to the code, which starts the Lambda function. The `handler` function is used to handle the "define auth challenge" event in Cognito, it defines how the auth flow should behave. If the session has the custom auth flow challenge, and it has been verified, it informs Cognito that the flow is valid and it can issue the auth tokens. However, if the challenge hasn't been verified, or a challenge doesn't yet exist, the function tells Cognito to use the custom auth challenge.

Next, let's build the function code, and create the Lambda resources. 

```bash
# ./

$ GOOS=linux GOARCH=amd64 go build -o define-challenge-function cmd/define-challenge-function/main.go
```

The command above is used to build the Go code, and output an executable file called "define-challenge-function". Once we have the compiled executable, we can zip it in preparation for creating the Lambda function in Terraform.

```bash
# ./

$ zip define-challenge-function.zip define-challenge-function
```

We should now have a zip file of the Lambda function code executable. We can now use this to create the Lambda function in Terraform.

```hcl
# ./main.tf

// Create the Lambda function's execution role.
data "aws_iam_policy_document" "assume_role" {
  statement {
    effect = "Allow"

    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }

    actions = ["sts:AssumeRole"]
  }
}

resource "aws_iam_role" "lambda" {
  name               = "user-impersonation-lambda"
  assume_role_policy = data.aws_iam_policy_document.assume_role.json
}

// Create the Lambda function.
resource "aws_lambda_function" "define_challenge" {
  filename         = "define-challenge-function.zip"
  source_code_hash = filebase64sha256("define-challenge-function.zip")
  function_name    = "app-users-define-auth-challenge"
  role             = aws_iam_role.lambda.arn
  handler          = "main"
  runtime          = "go1.x"
  timeout          = 30
}

// Next, create a Lambda permission resource, which allows Cognito
// to execute our Lambda function.
resource "aws_lambda_permission" "define_challenge" {
  statement_id  = "AllowExecutionFromCognito"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.define_challenge.function_name
  principal     = "cognito-idp.amazonaws.com"
  source_arn    = aws_cognito_user_pool.default.arn
}
```

And that's the first Lambda function done!

### Creating the create auth challenge Lambda function

The next step in the custom auth flow, is the creation of the auth challenge. This is used by Cognito, when the challenge has been defined by the previous Lambda function, to create the session specific arguments. To streamline this article, we will repeat the previous steps, but with the function code below, instead.

```go
# ./cmd/create-challenge/main.go

// The lambda function handler code.
func handler(ctx context.Context, event events.CognitoEventUserPoolsCreateAuthChallenge) (events.CognitoEventUserPoolsCreateAuthChallenge, error) {
	// Make sure we only an auth challenge for our custom flow.
	if event.Request.ChallengeName != "CUSTOM_CHALLENGE" {
		return event, nil
	}

	// Define our secret key, which we will use to verify the authentication flow.
	//
	// It would be a good idea to randomly generate this key, and store it in
	// something like AWS Secrets Manager - which also allows rotation.
	const secret = "my-secret-auth-key"

	// Set the private challenge paramters, used to verify the
	// auth challenge later in the flow.
	event.Response.PrivateChallengeParameters = map[string]string{
		"secret": secret,
	}

	return event, nil
}
```

The code can be used similar the the code in the last step, just with this `handler` instead. This handler is used to define any specific parameters for the custom challenge, and in this example we will define a private parameter called `secret`, with a secret string which we will use to verify the flow with.

### Create the verify auth challenge Lambda function

And, for the final Lambda function in the flow! The next function is used to verify the custom auth challenge, which will allow Cognito to issue an authentication token for a user.

As before, here is what the handler function code looks like for the verification step.

```go
# ./cmd/verify-challenge/main.go

// The lambda function handler code.
func handler(ctx context.Context, event events.CognitoEventUserPoolsVerifyAuthChallenge) (events.CognitoEventUserPoolsVerifyAuthChallenge, error) {
	// Get the "secret" as defined in the "create auth challenge" function.
	secret := event.Request.PrivateChallengeParameters["secret"]

	// Verify that the secret matches the auth challenge answer.
	valid := event.Request.ChallengeAnswer == secret

	// Set the answer in the response.
	event.Response.AnswerCorrect = valid

	return event, nil
}
```

This function is very simple, it verifies the auth challenge response with the defined "secret" we configured in the create auth challenge function. If they match, we tell Cognito that the answer is correct, which will complete the custom auth flow.

And that's all the Lambda function code we will need! However, we need to tell Cognito about the function. To do this, we need to update the user pool resource in the Terraform code.

```hcl
# ./main.tf

resource "aws_cognito_user_pool" "default" {
  name = "app-users"

  # The Lambda config is used to tell the user pool which functions
  # to use for various steps in the flows. 
  lambda_config {
    define_auth_challenge          = aws_lambda_function.define_challenge.arn
    verify_auth_challenge_response = aws_lambda_function.verify_challenge.arn
    create_auth_challenge          = aws_lambda_function.create_challenge.arn
  }
}
```

That should be all the Terraform we need, now to apply it to create the Lambda functions and the User Pools.

```bash
# ./

$ terraform apply
```

While the Terraform is creating our resources, we can write the final code to get an authentication token for a user. To get an authentication token, we need to perform two steps: firstly, we need to initialise the authentication flow, then we will need to respond to the custom auth challenge, with our secret key.

Here is some more Go code, which can be used to get the token. This code can be adapted to fit in to an application and used to make an authenticated API request as an impersonated user.

```go
package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"log"

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

func main() {
	userPoolId := "<USER POOL ID>"
	clientId := "<USER POOL CLIENT ID>"
	clientSecret := "<USER POOL CLIENT SECRET>"

	// The Cognito username of the user to login as.
	username := "example@domain.com"

	// To authenticate the an admin auth flow with Cognito, we need to generate
	// a secret hash. The secret hash is a HMAC SHA256 hash made up of the user's
	// username and the client id, where the HMAC key is the user pool client secret.
	//
	// Formula: BASE64(HMACSHA256(clientSecret, username + clientId))

	// Generate the secret hash, as described above.
	hash := hmac.New(sha256.New, []byte(clientSecret))
	hash.Write([]byte(username + clientId))
	secretHash := hash.Sum(nil)
	encoded := base64.StdEncoding.EncodeToString(secretHash)

	// Initialize a session which will use the default shared credentials,
	// and initalize a new Cognito SDK client.
	sess, err := session.NewSession()
	if err != nil {
		log.Fatalf("Error creating session: %v", err)
	}
	cognitoClient := cognitoidentityprovider.New(sess)

	// Use the Cognito client to initiate an admin auth flow.
	initiateResponse, err := cognitoClient.AdminInitiateAuth(&cognitoidentityprovider.AdminInitiateAuthInput{
		AuthFlow:   aws.String("CUSTOM_AUTH"),
		ClientId:   &clientId,
		UserPoolId: &userPoolId,
		AuthParameters: map[string]*string{
			"USERNAME":    &username,
			"SECRET_HASH": &encoded,
		},
	})
	if err != nil {
		log.Fatalf("Error initiating admin auth: %v", err)
	}

	// The secret key as defined by the create auth challenge function.
	secret := "my-secret-auth-key"

	// Once we've initialised the auth flow, we now need to respond to
	// the custom auth challenge. We can do this below, with our username,
	// secret hash and secret key as the answer.
	challengeResponse, err := cognitoClient.AdminRespondToAuthChallenge(&cognitoidentityprovider.AdminRespondToAuthChallengeInput{
		ChallengeName: initiateResponse.ChallengeName,
		Session:       initiateResponse.Session,
		ClientId:      &clientId,
		UserPoolId:    &userPoolId,
		ChallengeResponses: map[string]*string{
			"USERNAME":    &username,
			"SECRET_HASH": &encoded,
			"ANSWER":      &secret,
		},
	})
	if err != nil {
		log.Fatalf("Error responding to admin auth challenge: %v", err)
	}

	// Print the authentication result.
	fmt.Println("Authentication successful: ", *challengeResponse.AuthenticationResult)
}
```

By now the Terraform should have applied. The code can now be tested, or fit into an application. The `challengeResponse.AuthenticationResult` object should provide all authentication information needed to provide a good session, for example: token expiration, an access token and an ID token.

## Improvements

As with any example code, it can be improved and adapted for a specific applications needs, from tech stack to security practises. Whilst this example is using Terraform and Golang, a range of other tools could be used, for example you could use CloudFormation and NodeJS.

From an implementation perspective, a notable improvement that could be made it the "secret key" we use to create and verify the authentication challenge. This secret could be randomly generated and stored in AWS Secrets Manager, and automatically rotated on a schedule. This improves the security of the implementation as if the user pool client secret and the challenge secret key got leaked, the challenge's secret could be rotated very easily, without the need to redeploy the Lambda function code.

## Source code

If following along with a blog post isn't your thing, or you'd like to jump into a repository, here is a link to [Github](https://github.com/reecerussell/aws-cogntio-user-impersonation) with the full code.

User Impersonation With AWS Cognito

A guide to user impersonation with AWS Cognito