AWS Cognito MFA

import boto3, json, pyotp
import string, random
import sys
import hmac, hashlib, base64

class CognitoMFA:
    def __init__(self, USERNAME, EMAIL, TEMP_PASSWORD, CLIENT_ID, CLIENT_SECRET, USER_POOL_ID):
        self.USERNAME = USERNAME
        self.TEMP_PASSWORD = TEMP_PASSWORD
        self.CLIENT_ID = CLIENT_ID
        self.CLIENT_SECRET = CLIENT_SECRET
        self.USER_POOL_ID = USER_POOL_ID
        self.EMAIL = EMAIL
        self.SECRET_HASH = ''
        self.ACCESS_TOKEN = ''
        self.SECRET_TOKEN = ''
        # Create boto3 client for cognito to use
        self.client = boto3.client('cognito-idp')
    
    def __str__(self):
        return self.SECRET_HASH
    
    # Get the mysterious secret hash
    def GetSecretHash(self):
        message = bytes(self.USERNAME+self.CLIENT_ID,'utf-8')
        key = bytes(self.CLIENT_SECRET,'utf-8')
        self.SECRET_HASH = base64.b64encode(hmac.new(key, message, digestmod=hashlib.sha256).digest()).decode()
        return self.SECRET_HASH

    # Get a password that meets compliance
    def RandomPassword(self):
        chars = string.ascii_uppercase + string.ascii_lowercase + string.digits + string.punctuation
        size = random.randint(16, 20)
        return 'T3m9' + ''.join(random.choice(chars) for x in range(size))

    # Create the user account
    def CreateUser(self):
        response = self.client.admin_create_user(
            UserPoolId = self.USER_POOL_ID,
            Username = self.USERNAME,
            UserAttributes=[
                {
                    'Name': 'email',
                    'Value': self.EMAIL
                },
                {
                    'Name': 'email_verified',
                    'Value': 'true'
                }
            ],
            TemporaryPassword = self.TEMP_PASSWORD,
            MessageAction = 'SUPPRESS',
            DesiredDeliveryMediums = [
                'EMAIL',
            ]
        )
    
    # Change the users password for enabling mfa
    def SetUserPassword(self, permState, random):
        if (random == 'y'):
            password = self.RandomPassword()
        else:
            password = self.TEMP_PASSWORD
        response = self.client.admin_set_user_password(
            UserPoolId = self.USER_POOL_ID,
            Username = self.USERNAME,
            Password = password,
            Permanent = permState
        )
        return password
    
    # Get the user token
    def GetUserToken(self):
        initiateAuth = self.client.initiate_auth(
            AuthFlow = "USER_PASSWORD_AUTH",
            AuthParameters = {
                'USERNAME': self.USERNAME,
                'PASSWORD': self.TEMP_PASSWORD,
                'SECRET_HASH': self.SECRET_HASH
            },
            ClientId = self.CLIENT_ID
        )
        self.ACCESS_TOKEN = initiateAuth["AuthenticationResult"]["AccessToken"]
        return self.ACCESS_TOKEN
    
    # Get the mfa token
    def GetMFAToken(self):
        associateSoftwareToken = self.client.associate_software_token(
            AccessToken = self.ACCESS_TOKEN
        )
        self.SECRET_TOKEN = associateSoftwareToken["SecretCode"]
        return self.SECRET_TOKEN

    # Verify the token
    def VerifyToken(self):
        totp = pyotp.TOTP(self.SECRET_TOKEN)
        response = self.client.verify_software_token(
            AccessToken = self.ACCESS_TOKEN,
            UserCode = totp.now()
        )

    # Enable on mfa for the user account
    def EnableUserMFA(self):
        response = self.client.admin_set_user_mfa_preference(
            SoftwareTokenMfaSettings = {
                'Enabled': True,
                'PreferredMfa': True
            },
            Username = self.USERNAME,
            UserPoolId = self.USER_POOL_ID
        )

if __name__ == '__main__':
    # Check for all of the arguments
    if len(sys.argv) != 6:
        print("[!] Usage python3 CognitoMFA.py USERNAME EMAIL CLIENT_ID CLIENT_SECRET USER_POOL_ID")
        sys.exit(0)

    # Get arguments from the command line
    USERNAME = sys.argv[1]
    EMAIL = sys.argv[2]
    TEMP_PASSWORD = "TempPass123!"
    CLIENT_ID = sys.argv[3]
    CLIENT_SECRET = sys.argv[4]
    USER_POOL_ID = sys.argv[5]

    # Run the appropriate python commands from the class above
    R = CognitoMFA(USERNAME, EMAIL, TEMP_PASSWORD, CLIENT_ID, CLIENT_SECRET, USER_POOL_ID)
    R.GetSecretHash()
    R.CreateUser()
    R.SetUserPassword(True, 'n')
    R.GetUserToken()
    MFA_TOKEM = R.GetMFAToken()
    print(f"Your MFA token: {MFA_TOKEM}")
    R.VerifyToken()
    R.EnableUserMFA()
    Password = R.SetUserPassword(False, 'y')
    print(f"Your password: {Password}")

Usage python3 CognitoMFA.py USERNAME EMAIL CLIENT_ID CLIENT_SECRET USER_POOL_ID