tessl/pypi-pycognito

Python class to integrate Boto3's Cognito client so it is easy to login users with SRP support.

—

Pending

Overview

Eval results

Files

Multi-Factor Authentication (MFA)

Name: tessl/pypi-pycognito
Author: tessl

Complete multi-factor authentication support for AWS Cognito User Pools, including software token (TOTP) and SMS MFA setup, verification, preference management, and challenge responses. Provides enhanced security through two-factor authentication methods.

Capabilities

Software Token MFA (TOTP)

Set up and manage Time-based One-Time Password (TOTP) authentication using authenticator apps like Google Authenticator, Authy, or Microsoft Authenticator.

def associate_software_token(self) -> str:
    """
    Get the secret code for Software Token MFA setup.
    
    Returns:
        str: Base32-encoded secret code for TOTP setup
        
    Usage:
        - Display as QR code for authenticator app scanning
        - Provide as manual entry code for authenticator apps
        - Secret is unique per user and doesn't change
        
    Requirements:
        - User must be authenticated (valid access token)
        - User pool must have MFA enabled
    """

def verify_software_token(self, code: str, device_name: str = "") -> bool:
    """
    Verify TOTP code to complete Software Token MFA registration.
    
    Args:
        code (str): 6-digit TOTP code from authenticator app
        device_name (str, optional): Friendly name for the device
        
    Returns:
        bool: True if verification successful, False otherwise
        
    Actions:
        - Validates TOTP code against user's secret
        - Registers the software token for the user
        - Enables software token MFA for the account
        
    Note:
        Must be called after associate_software_token() to complete setup
    """

Usage Example:

from pycognito import Cognito
import qrcode
import io

# User must be authenticated
u = Cognito('pool-id', 'client-id', username='user')
u.authenticate(password='password')

# Get TOTP secret
secret = u.associate_software_token()
print(f"Secret code: {secret}")

# Generate QR code for easy setup
totp_uri = f"otpauth://totp/{u.username}?secret={secret}&issuer=MyApp"
qr = qrcode.QRCode()
qr.add_data(totp_uri)
qr.make()

print("Scan QR code with authenticator app:")
qr.print_ascii()

# User sets up authenticator and provides first code
code = input("Enter 6-digit code from authenticator: ")

# Verify and complete setup
if u.verify_software_token(code, "My Phone"):
    print("Software Token MFA enabled successfully!")
else:
    print("Verification failed. Please try again.")

MFA Preference Management

Configure MFA preferences for users, including which methods are enabled and which is preferred.

def set_user_mfa_preference(self, sms_mfa: bool, software_token_mfa: bool, preferred: str = None) -> None:
    """
    Set MFA preferences for the authenticated user.
    
    Args:
        sms_mfa (bool): Enable/disable SMS MFA
        software_token_mfa (bool): Enable/disable Software Token MFA
        preferred (str, optional): Preferred method ("SMS", "SOFTWARE_TOKEN", or None)
        
    Valid combinations:
        - Both False: Disable MFA entirely
        - One True: Enable that method as preferred
        - Both True: Enable both, use preferred parameter to choose default
        
    Requirements:
        - User must be authenticated
        - If enabling SMS MFA, phone number must be verified
        - If enabling Software Token MFA, must call verify_software_token() first
        
    Raises:
        ValueError: If preferred value is invalid
    """

Usage Example:

# Enable only Software Token MFA
u.set_user_mfa_preference(
    sms_mfa=False,
    software_token_mfa=True,
    preferred="SOFTWARE_TOKEN"
)

# Enable both methods with SMS as preferred
u.set_user_mfa_preference(
    sms_mfa=True,
    software_token_mfa=True,
    preferred="SMS"
)

# Disable MFA entirely
u.set_user_mfa_preference(
    sms_mfa=False,
    software_token_mfa=False
)

print("MFA preferences updated!")

MFA Challenge Responses

Handle MFA challenges during authentication by responding with appropriate codes.

def respond_to_software_token_mfa_challenge(self, code: str, mfa_tokens: dict = None) -> None:
    """
    Respond to Software Token MFA challenge during authentication.
    
    Args:
        code (str): 6-digit TOTP code from authenticator app
        mfa_tokens (dict, optional): MFA tokens from exception (uses instance tokens if not provided)
        
    Actions:
        - Validates TOTP code
        - Completes authentication flow
        - Sets access, ID, and refresh tokens on success
        
    Usage:
        Called after catching SoftwareTokenMFAChallengeException during authenticate()
    """

def respond_to_sms_mfa_challenge(self, code: str, mfa_tokens: dict = None) -> None:
    """
    Respond to SMS MFA challenge during authentication.
    
    Args:
        code (str): SMS verification code
        mfa_tokens (dict, optional): MFA tokens from exception (uses instance tokens if not provided)
        
    Actions:
        - Validates SMS code
        - Completes authentication flow  
        - Sets access, ID, and refresh tokens on success
        
    Usage:
        Called after catching SMSMFAChallengeException during authenticate()
    """

Usage Example:

from pycognito import Cognito
from pycognito.exceptions import (
    SoftwareTokenMFAChallengeException,
    SMSMFAChallengeException
)

u = Cognito('pool-id', 'client-id', username='user')

try:
    u.authenticate(password='password')
    print("Authentication successful!")
    
except SoftwareTokenMFAChallengeException as e:
    print("Software Token MFA required")
    code = input("Enter 6-digit code from authenticator: ")
    u.respond_to_software_token_mfa_challenge(code, e.get_tokens())
    print("MFA authentication successful!")
    
except SMSMFAChallengeException as e:
    print("SMS MFA required")
    code = input("Enter SMS verification code: ")
    u.respond_to_sms_mfa_challenge(code, e.get_tokens())
    print("MFA authentication successful!")

Exception Handling

MFA operations use specialized exceptions to handle different challenge types.

class MFAChallengeException(WarrantException):
    """Base class for MFA challenge exceptions."""
    
    def __init__(self, message: str, tokens: dict, *args, **kwargs):
        """
        Args:
            message (str): Exception message
            tokens (dict): MFA challenge tokens needed for response
        """
    
    def get_tokens(self) -> dict:
        """
        Get MFA challenge tokens.
        
        Returns:
            dict: Tokens needed for MFA challenge response
        """

class SoftwareTokenMFAChallengeException(MFAChallengeException):
    """Raised when Software Token (TOTP) MFA is required."""

class SMSMFAChallengeException(MFAChallengeException):
    """Raised when SMS MFA is required."""

Usage Patterns

Complete TOTP Setup Flow

def setup_totp_mfa(username, password):
    """Complete TOTP MFA setup flow."""
    u = Cognito('pool-id', 'client-id', username=username)
    
    # Authenticate user
    u.authenticate(password=password)
    
    # Get TOTP secret
    secret = u.associate_software_token()
    
    print("Set up your authenticator app:")
    print(f"Manual entry code: {secret}")
    print(f"Or scan QR code for: otpauth://totp/{username}?secret={secret}&issuer=MyApp")
    
    # Wait for user to set up authenticator
    input("Press Enter after setting up authenticator app...")
    
    # Verify setup
    while True:
        code = input("Enter 6-digit code from authenticator: ")
        
        if u.verify_software_token(code, "Primary Device"):
            print("TOTP setup successful!")
            break
        else:
            print("Invalid code. Please try again.")
    
    # Set as preferred MFA method
    u.set_user_mfa_preference(
        sms_mfa=False,
        software_token_mfa=True, 
        preferred="SOFTWARE_TOKEN"
    )
    
    print("TOTP MFA is now enabled and preferred!")

# Usage
setup_totp_mfa('user@example.com', 'password')

MFA-Aware Authentication

def authenticate_with_mfa(username, password):
    """Authenticate user with MFA support."""
    u = Cognito('pool-id', 'client-id', username=username)
    
    try:
        # Attempt primary authentication
        u.authenticate(password=password)
        print("Authentication successful (no MFA required)")
        return u
        
    except SoftwareTokenMFAChallengeException as e:
        print("TOTP MFA required")
        
        # Get TOTP code from user
        code = input("Enter 6-digit code from authenticator: ")
        
        # Respond to challenge
        u.respond_to_software_token_mfa_challenge(code, e.get_tokens())
        print("TOTP MFA authentication successful!")
        return u
        
    except SMSMFAChallengeException as e:
        print("SMS MFA required")
        print("SMS code sent to your registered phone number")
        
        # Get SMS code from user
        code = input("Enter SMS verification code: ")
        
        # Respond to challenge
        u.respond_to_sms_mfa_challenge(code, e.get_tokens())
        print("SMS MFA authentication successful!")
        return u
        
    except Exception as e:
        print(f"Authentication failed: {e}")
        return None

# Usage
user = authenticate_with_mfa('user@example.com', 'password')
if user:
    print(f"Access token: {user.access_token}")

MFA Configuration Management

def manage_mfa_settings(u):
    """Interactive MFA settings management."""
    
    print("\nCurrent MFA Settings:")
    print("1. Enable TOTP MFA only")
    print("2. Enable SMS MFA only") 
    print("3. Enable both (TOTP preferred)")
    print("4. Enable both (SMS preferred)")
    print("5. Disable MFA")
    
    choice = input("Select option (1-5): ")
    
    if choice == "1":
        # Setup TOTP if not already done
        secret = u.associate_software_token()
        print(f"Setup authenticator with: {secret}")
        code = input("Enter verification code: ")
        
        if u.verify_software_token(code):
            u.set_user_mfa_preference(
                sms_mfa=False,
                software_token_mfa=True,
                preferred="SOFTWARE_TOKEN"
            )
            print("TOTP MFA enabled!")
            
    elif choice == "2":
        u.set_user_mfa_preference(
            sms_mfa=True,
            software_token_mfa=False,
            preferred="SMS"
        )
        print("SMS MFA enabled!")
        
    elif choice == "3":
        # Ensure TOTP is set up
        secret = u.associate_software_token()
        print(f"Setup authenticator with: {secret}")
        code = input("Enter verification code: ")
        
        if u.verify_software_token(code):
            u.set_user_mfa_preference(
                sms_mfa=True,
                software_token_mfa=True,
                preferred="SOFTWARE_TOKEN"
            )
            print("Both MFA methods enabled (TOTP preferred)!")
            
    elif choice == "4":
        # Ensure TOTP is set up
        secret = u.associate_software_token()
        print(f"Setup authenticator with: {secret}")
        code = input("Enter verification code: ")
        
        if u.verify_software_token(code):
            u.set_user_mfa_preference(
                sms_mfa=True,
                software_token_mfa=True,
                preferred="SMS"
            )
            print("Both MFA methods enabled (SMS preferred)!")
            
    elif choice == "5":
        u.set_user_mfa_preference(
            sms_mfa=False,
            software_token_mfa=False
        )
        print("MFA disabled!")
        
    else:
        print("Invalid choice")

# Usage
u = Cognito('pool-id', 'client-id', username='user')
u.authenticate(password='password')
manage_mfa_settings(u)

Production MFA Helper

class MFAHelper:
    """Helper class for production MFA handling."""
    
    def __init__(self, user_pool_id, client_id):
        self.user_pool_id = user_pool_id
        self.client_id = client_id
    
    def authenticate_user(self, username, password, mfa_callback=None):
        """
        Authenticate user with MFA callback support.
        
        Args:
            username: Username
            password: Password
            mfa_callback: Function to get MFA code (type, tokens) -> code
        """
        u = Cognito(self.user_pool_id, self.client_id, username=username)
        
        try:
            u.authenticate(password=password)
            return u, "SUCCESS"
            
        except SoftwareTokenMFAChallengeException as e:
            if mfa_callback:
                code = mfa_callback("TOTP", e.get_tokens())
                u.respond_to_software_token_mfa_challenge(code, e.get_tokens())
                return u, "MFA_SUCCESS"
            return None, "MFA_REQUIRED_TOTP"
            
        except SMSMFAChallengeException as e:
            if mfa_callback:
                code = mfa_callback("SMS", e.get_tokens())
                u.respond_to_sms_mfa_challenge(code, e.get_tokens())
                return u, "MFA_SUCCESS"
            return None, "MFA_REQUIRED_SMS"
            
        except Exception as e:
            return None, f"AUTH_FAILED: {e}"
    
    def setup_totp(self, username, password, device_name="Device"):
        """Setup TOTP MFA for user."""
        u = Cognito(self.user_pool_id, self.client_id, username=username)
        u.authenticate(password=password)
        
        secret = u.associate_software_token()
        return secret, u

# Usage
def get_mfa_code(mfa_type, tokens):
    """Callback to get MFA code from user."""
    if mfa_type == "TOTP":
        return input("Enter TOTP code: ")
    elif mfa_type == "SMS":
        return input("Enter SMS code: ")

helper = MFAHelper('pool-id', 'client-id')
user, status = helper.authenticate_user('username', 'password', get_mfa_code)

if status == "SUCCESS":
    print("Authenticated without MFA")
elif status == "MFA_SUCCESS":
    print("Authenticated with MFA")
else:
    print(f"Authentication failed: {status}")

Install with Tessl CLI