Various Python 3 utility scripts related to AWS Cognito authentication

aws_cognito_get_creds.py

#!/usr/bin/env python3
# Author: Sean Pesce
#
# Obtain AWS Cognito user identity ID and credentials

import argparse
import getpass
import json
import os
import requests
import sys

ID_TOKEN_ENV_VAR = 'AWS_COGNITO_ID_TOKEN'


def get_aws_cognito_identity_id(identity_pool, user_pool, id_token, region):
    headers = {
        'Content-type': 'application/x-amz-json-1.1',
        'x-amz-api-version': '2016-04-18',
        'X-Amz-Target': 'AWSCognitoIdentityService.GetId',
    }
    url = f'https://cognito-identity.{region}.amazonaws.com/'
    body = {
        'IdentityPoolId': identity_pool,
        'Logins': {
            f'cognito-idp.{region}.amazonaws.com/{user_pool}': id_token,
        },
    }
    body = json.dumps(body)
    response = requests.post(url, headers=headers, data=body)
    assert response.status_code == 200, f'[ERROR] HTTP response {response.status_code}: {response.text}'
    return response.text


def get_aws_cognito_credentials(identity_id, user_pool, id_token, region):
    headers = {
        'Content-type': 'application/x-amz-json-1.1',
        'x-amz-api-version': '2016-04-18',
        'X-Amz-Target': 'AWSCognitoIdentityService.GetCredentialsForIdentity',
    }
    url = f'https://cognito-identity.{region}.amazonaws.com/'
    body = {
        'IdentityId': identity_id,
        'Logins': {
            f'cognito-idp.{region}.amazonaws.com/{user_pool}': id_token,
        },
    }
    body = json.dumps(body)
    response = requests.post(url, headers=headers, data=body)
    assert response.status_code == 200, f'[ERROR] HTTP response {response.status_code}: {response.text}'
    return response.text


if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('-i', '--identity-pool', type=str, help=f'Identity Pool ID (example: "us-east-1:00000000-0000-0000-0000-000000000000")', required=True)
    argparser.add_argument('-u', '--user-pool', type=str, help=f'User pool ID (example: "us-east-1_xxxxxxxxx")', required=True)
    argparser.add_argument('-r', '--region', type=str, help=f'AWS region (example: "us-east-1")', required=True)
    argparser.add_argument('-t', '--id-token', type=str, help=f'ID token. If not provided, the ${ID_TOKEN_ENV_VAR} environment variable is used. If ${ID_TOKEN_ENV_VAR} is not set, you will be prompted for the token.', default=None)
    argparser.add_argument('-v', '--verbose', help='Increase output verbosity', action='store_true')
    if '-h' in sys.argv or '--help' in sys.argv:
        argparser.print_help()
        print(f'\nTool to obtain AWS Cognito user credentials', file=sys.stderr)
        sys.exit(0)
    args = argparser.parse_args()

    if args.id_token is None:
        # Check for the identity token environment variable
        args.id_token = os.environ.get(ID_TOKEN_ENV_VAR)
    
    if args.id_token is None:
        # Prompt for identity token
        args.id_token = getpass.getpass('ID Token: ')
    
    # Get identity ID
    auth_response = get_aws_cognito_identity_id(args.identity_pool,
                                                args.user_pool,
                                                args.id_token,
                                                args.region)
    identity_id = json.loads(auth_response)['IdentityId']

    # Get credentials
    creds_response = get_aws_cognito_credentials(identity_id,
                                                 args.user_pool,
                                                 args.id_token,
                                                 args.region)
    credentials = json.dumps(json.loads(creds_response), indent=2)
    print(credentials)

aws_cognito_user_srp_auth.py

#!/usr/bin/env python3
# Author: Sean Pesce
#
# Obtain AWS credentials via AWS Cognito USER_SRP_AUTH authentication flow

import argparse
import getpass
import json
import os
import sys

from pycognito.aws_srp import AWSSRP

PW_ENV_VAR = 'AWS_COGNITO_PW'

def authenticate_USER_SRP_AUTH(username, password, pool_id, client_id, region):
    aws = AWSSRP(username=username,
                 password=password,
                 pool_id=pool_id,
                 client_id=client_id,
                 pool_region=region)
    response = aws.authenticate_user()
    return response


if __name__ == '__main__':
    argparser = argparse.ArgumentParser()
    argparser.add_argument('-u', '--user', type=str, help=f'Username', required=True)
    argparser.add_argument('-P', '--password', type=str, help=f'Account password. If not provided, the ${PW_ENV_VAR} environment variable is used. If ${PW_ENV_VAR} is not set, you will be prompted for a password.', default=None)
    argparser.add_argument('-p', '--pool', type=str, help=f'User pool ID (example: "us-east-1_xxxxxxxxx")', required=True)
    argparser.add_argument('-c', '--client-id', type=str, help=f'Client ID', required=True)
    argparser.add_argument('-r', '--region', type=str, help=f'AWS region (example: "us-east-1")', required=True)
    argparser.add_argument('-v', '--verbose', help='Increase output verbosity', action='store_true')
    if '-h' in sys.argv or '--help' in sys.argv:
        argparser.print_help()
        print(f'\nTool to obtain AWS credentials via AWS Cognito USER_SRP_AUTH authentication flow', file=sys.stderr)
        sys.exit(0)
    args = argparser.parse_args()

    if args.password is None:
        # Check for the password environment variable
        args.password = os.environ.get(PW_ENV_VAR)
    
    if args.password is None:
        # Prompt for password
        args.password = getpass.getpass('Password: ')
    
    auth_response = authenticate_USER_SRP_AUTH(args.user,
                                               args.password,
                                               args.pool,
                                               args.client_id,
                                               args.region)
    auth_data = json.dumps(auth_response, indent=2)
    print(auth_data)

Note: The ID token (required to obtain credentials) can be obtained by authenticating to AWS Cognito.