Skip to content

Instantly share code, notes, and snippets.

@andrew-aiken
Last active June 30, 2023 10:39
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save andrew-aiken/f46b29c6eeeb51b30c2df618bff6f01c to your computer and use it in GitHub Desktop.
Save andrew-aiken/f46b29c6eeeb51b30c2df618bff6f01c to your computer and use it in GitHub Desktop.
AWS Cognito MFA Setup Script
# AWS Cognito MFA
# MFA is not configured by default when using the AWS Cognito web UI.
# The following script will setup a user account, setup MFA for the user, and return a temporary password.
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}")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment