Authentication and Token Management

def authenticate(self, password: str, client_metadata: dict = None) -> None:
    """
    Authenticate user using the SRP protocol.
    
    Args:
        password (str): The user's password
        client_metadata (dict, optional): Metadata for custom workflows
        
    Raises:
        MFAChallengeException: When MFA is required
        ForceChangePasswordException: When password change is required
        
    Sets instance attributes:
        - id_token: JWT ID token
        - access_token: JWT access token  
        - refresh_token: Token for refreshing access
        - token_type: Type of token (usually "Bearer")
        - id_claims: Verified claims from ID token
        - access_claims: Verified claims from access token
    """

from pycognito import Cognito
from pycognito.exceptions import MFAChallengeException

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

try:
    u.authenticate(password='bobs-password')
    print("Authentication successful!")
    print(f"Access token: {u.access_token}")
except MFAChallengeException as e:
    # Handle MFA challenge
    mfa_tokens = e.get_tokens()
    # Proceed with MFA flow

def admin_authenticate(self, password: str) -> None:
    """
    Authenticate user using admin super privileges.
    
    Args:
        password (str): The user's password
        
    Sets same instance attributes as authenticate()
    """

u = Cognito('your-user-pool-id', 'your-client-id', username='bob')
u.admin_authenticate(password='bobs-password')

def new_password_challenge(self, password: str, new_password: str) -> None:
    """
    Handle NEW_PASSWORD_REQUIRED challenge by setting a new password.
    
    Args:
        password (str): The user's current/temporary password
        new_password (str): The new password to set
        
    Raises:
        Exception: If the challenge response fails
        
    Sets instance attributes:
        - id_token: JWT ID token
        - access_token: JWT access token  
        - refresh_token: Token for refreshing access
        - token_type: Type of token (usually "Bearer")
        - id_claims: Verified claims from ID token
        - access_claims: Verified claims from access token
    """

from pycognito import Cognito
from pycognito.exceptions import ForceChangePasswordException

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

try:
    u.authenticate(password='temporary-password')
except ForceChangePasswordException:
    # User must change their password
    u.new_password_challenge(
        password='temporary-password',
        new_password='new-secure-password'
    )
    print("Password changed successfully!")
    print(f"Access token: {u.access_token}")

def verify_tokens(self) -> None:
    """
    Verify the current id_token and access_token.
    
    Raises:
        TokenVerificationException: If tokens fail verification
        
    Verifies:
        - Token signatures using JWK keys
        - Token expiration times
        - Token audience and issuer
        - Token use claims
    """

def verify_token(self, token: str, id_name: str, token_use: str) -> dict:
    """
    Verify a specific token.
    
    Args:
        token (str): The JWT token to verify
        id_name (str): Token identifier ("id_token" or "access_token")
        token_use (str): Expected token use ("id" or "access")
        
    Returns:
        dict: Verified token payload
        
    Raises:
        TokenVerificationException: If token verification fails
    """

# Verify tokens from a previous session
u = Cognito(
    'your-user-pool-id', 
    'your-client-id',
    id_token='existing-id-token',
    access_token='existing-access-token',
    refresh_token='existing-refresh-token'
)

try:
    u.verify_tokens()  # Throws exception if invalid
    print("Tokens are valid!")
    print(f"User ID: {u.id_claims['sub']}")
    print(f"Username: {u.access_claims['username']}")
except TokenVerificationException:
    print("Tokens are invalid or expired")

def check_token(self, renew: bool = True) -> bool:
    """
    Check if access token has expired and optionally renew it.
    
    Args:
        renew (bool): Whether to automatically renew expired tokens
        
    Returns:
        bool: True if token was expired, False if still valid
        
    Raises:
        AttributeError: If no access token is available
    """

def renew_access_token(self) -> None:
    """
    Set a new access token using the cached refresh token.
    
    Raises:
        Exception: If refresh token is invalid or expired
        
    Updates:
        - access_token: New access token
        - id_token: New ID token
        - token_type: Token type
        - access_claims: New access token claims
        - id_claims: New ID token claims
    """

# Check and renew token automatically
if u.check_token(renew=True):
    print("Token was expired and has been renewed")
else:
    print("Token is still valid")

# Manual token renewal
u.renew_access_token()
print(f"New access token: {u.access_token}")

def logout(self) -> None:
    """
    Log user out of all clients and remove token attributes.
    
    Actions:
        - Calls AWS Cognito global sign out
        - Clears id_token, refresh_token, access_token, token_type
        - Invalidates tokens on all devices/clients
    """

# Logout user
u.logout()
print("User logged out successfully")

# Tokens are now None
assert u.access_token is None
assert u.id_token is None
assert u.refresh_token is None

def get_keys(self) -> dict:
    """
    Get JWK keys for token verification.
    
    Returns:
        dict: JWK key set from AWS Cognito or environment
        
    Checks:
        1. Environment variable COGNITO_JWKS
        2. AWS Cognito /.well-known/jwks.json endpoint
    """

def get_key(self, kid: str) -> dict:
    """
    Get specific JWK key by key ID.
    
    Args:
        kid (str): Key ID from JWT header
        
    Returns:
        dict: JWK key data
    """

def switch_session(self, session) -> None:
    """
    Switch the boto3 session used for Cognito operations.
    
    Args:
        session: boto3 session object
        
    Used primarily for:
        - Unit testing with placebo library
        - Different AWS credentials/regions
        - Custom boto3 configurations
    """

from pycognito import Cognito

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

# Authenticate
u.authenticate(password='password')

# Use tokens
headers = {'Authorization': f'Bearer {u.access_token}'}

# Check token before use
if not u.check_token():
    # Token is still valid
    make_api_call(headers)

# Save tokens after authentication
session_data = {
    'id_token': u.id_token,
    'access_token': u.access_token, 
    'refresh_token': u.refresh_token
}

# Later, restore session
u = Cognito(
    'pool-id', 'client-id',
    id_token=session_data['id_token'],
    access_token=session_data['access_token'],
    refresh_token=session_data['refresh_token']
)

# Verify and potentially renew
try:
    u.verify_tokens()
except TokenVerificationException:
    # Handle invalid tokens
    redirect_to_login()

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

try:
    u.authenticate(password='password')
except SoftwareTokenMFAChallengeException as e:
    # Handle TOTP MFA
    code = input("Enter TOTP code: ")
    u.respond_to_software_token_mfa_challenge(code, e.get_tokens())
except SMSMFAChallengeException as e:
    # Handle SMS MFA
    code = input("Enter SMS code: ")
    u.respond_to_sms_mfa_challenge(code, e.get_tokens())