Cognito userpool export

cognitoUserMigration.py

import os
import time
import sys
import boto3

REGION = "eu-central-1"
S3_CLIENT = boto3.client("s3", REGION)
COGNITO_CLIENT = boto3.client("cognito-idp", REGION)
USER_POOL_ID = os.environ['USER_POOL_ID']
EXPORTED_RECORDS_COUNT = 0
BUCKET_NAME = os.environ["BUCKET_NAME"]
PAGINATION_COUNTER = 0
CSV_FILE_NAME = "CognitoUsers.csv"
REQUIRED_ATTRIBUTES = [
    "name", "given_name", "family_name", "middle_name", "nickname",
    "preferred_username", "profile", "picture", "website", "email",
    "email_verified", "gender", "birthdate", "zoneinfo", "locale",
    "phone_number", "phone_number_verified", "address", "updated_at",
    "custom:joinedOn", "cognito:mfa_enabled", "cognito:username"
]
MAX_NUMBER_RECORDS = 0
LIMIT = 60
FILE_PATH = os.path.join(os.path.dirname(__file__),
                         f"../../tmp/{CSV_FILE_NAME}")


def get_list_cognito_users(cognito_idp_client,
                           next_pagination_token="",
                           limit=LIMIT):
    """
    First we get a list of all users in the pull.
    The Limit per page is 60. If there is next page we will move to it
    """
    return cognito_idp_client.list_users(
        UserPoolId=USER_POOL_ID,
        limit=limit,
        PaginationToken=next_pagination_token
    ) if next_pagination_token else cognito_idp_client.list_users(
        UserPoolId=USER_POOL_ID, limit=limit)


def check_next_pagination_token_existence(user_records):
    """
    We are checking there is a next page
    """
    if set(["PaginationToken", "NextToken"]).intersection(set(user_records)):
        pagination_token = user_records[
            "PaginationToken"] if "PaginationToken" in user_records else user_records[
                "NextToken"]
    return pagination_token


def write_to_csv_file(user_records, csv_file, csv_new_line):
    """
     We will save each user attribute as csv
     """
    csv_lines = []
    for user in user_records["Users"]:
        csv_line = csv_new_line.copy()
        for required_attribute in REQUIRED_ATTRIBUTES:
            csv_line[required_attribute] = ""
            if required_attribute in user.keys():
                csv_line[required_attribute] = str(user[required_attribute])
                continue
            if required_attribute in ("phone_number_verified",
                                      "cognito:mfa_enabled"):
                csv_line[required_attribute] = "false"
            if required_attribute == "email_verified" and required_attribute not in user.keys(
            ):
                csv_line[required_attribute] = "false"
            for user_attribute in user["Attributes"]:
                if user_attribute["Name"] == required_attribute:
                    csv_line[required_attribute] = str(user_attribute["Value"])
        csv_line["cognito:username"] = user["Username"]
        if not csv_line["email_verified"] == "false":
            csv_lines.append(",".join(csv_line.values()) + "\n")
    print(csv_lines)
    csv_file.writelines(csv_lines)
    global EXPORTED_RECORDS_COUNT
    EXPORTED_RECORDS_COUNT += len(csv_lines)
    global PAGINATION_COUNTER
    print("Page: #{} \n Total Exported Records: #{} \n".format(
        str(PAGINATION_COUNTER), str(EXPORTED_RECORDS_COUNT)))


def open_csv_file(csv_new_line):
    """
    Create a csv file in the tmp directory of the lambda function
    """
    try:
        csv_file = open(FILE_PATH, "w")
        csv_file.write(",".join(csv_new_line.keys()) + "\n")
    except Exception as err:
        error_message = repr(err)
        print("\nERROR: Can not create file: " + FILE_PATH)
        print("\tError Reason: " + error_message)
        sys.exit()
    return csv_file


def cooldown_before_next_batch():
    """
    Wait before next batch
    """
    time.sleep(0.15)


def save_file():
    """
    Upload file to S3
    """
    with open(FILE_PATH, "rb") as file:
        S3_CLIENT.upload_fileobj(file, BUCKET_NAME, CSV_FILE_NAME)


def lambda_handler(event, context):
    """
    lambda_handler is called, when the lambda function was invoked
    """
    pagination_token = ""
    csv_new_line = {
        REQUIRED_ATTRIBUTES[i]: ""
        for i in range(len(REQUIRED_ATTRIBUTES))
    }
    csv_file = open_csv_file(csv_new_line)
    global PAGINATION_COUNTER
    while pagination_token is not None:
        try:
            user_records = get_list_cognito_users(
                cognito_idp_client=COGNITO_CLIENT,
                next_pagination_token=pagination_token,
                limit=LIMIT
                if LIMIT < MAX_NUMBER_RECORDS else MAX_NUMBER_RECORDS)
        except COGNITO_CLIENT.exceptions.ClientError as client_error:
            error_message = client_error.response["Error"]["Message"]
            print("PLEASE CHECK YOUR COGNITO CONFIGS")
            print("Error Reason: " + error_message)
            csv_file.close()
            sys.exit()
        except Exception as unknown_exception:
            print("Error Reason: " + str(unknown_exception))
            csv_file.close()
            sys.exit()
        pagination_token = check_next_pagination_token_existence(user_records)
        write_to_csv_file(user_records, csv_file, csv_new_line)
        PAGINATION_COUNTER += 1
        cooldown_before_next_batch()
    csv_file.close()
    save_file()