Device Authentication

def confirm_device(self, tokens: dict, device_name: str = None) -> tuple:
    """
    Confirm a device for device tracking and authentication.
    
    Args:
        tokens (dict): Authentication result tokens containing device metadata
        device_name (str, optional): Friendly name for device (defaults to hostname)
        
    Returns:
        tuple: (response, device_password)
            - response: HTTP response from device confirmation
            - device_password: Generated device password for future authentication
            
    Actions:
        - Registers device with Cognito
        - Generates device-specific credentials
        - Sets up device for remembered authentication
        - Stores device keys and metadata
        
    Sets instance attributes:
        - access_token: Token from authentication result
        - device_key: Unique device identifier
        - device_group_key: Device group identifier
        - device_name: Friendly device name
    """

from pycognito.aws_srp import AWSSRP

# Authenticate user and get tokens with device metadata
aws_srp = AWSSRP(
    username='user@example.com',
    password='password',
    pool_id='us-east-1_example',
    client_id='your-client-id'
)

# Authenticate and get tokens
tokens = aws_srp.authenticate_user()

# Check if device confirmation is needed
if 'NewDeviceMetadata' in tokens.get('AuthenticationResult', {}):
    print("New device detected - confirming device")
    
    # Confirm device with friendly name
    response, device_password = aws_srp.confirm_device(
        tokens, 
        device_name="John's iPhone"
    )
    
    print(f"Device confirmed! Device password: {device_password}")
    print(f"Device key: {aws_srp.device_key}")
    
    # Store device credentials securely for future use
    store_device_credentials(aws_srp.device_key, device_password)

def update_device_status(self, is_remembered: bool, access_token: str, device_key: str) -> str:
    """
    Update device remembered status.
    
    Args:
        is_remembered (bool): Whether device should be remembered
        access_token (str): Valid access token
        device_key (str): Device key from device confirmation
        
    Returns:
        str: Response status from update operation
        
    Device status options:
        - remembered: Device is trusted and remembers user
        - not_remembered: Device requires full authentication each time
        
    Use cases:
        - Mark personal devices as remembered
        - Mark public/shared devices as not remembered
        - Revoke device trust due to security concerns
    """

# Mark device as remembered (trusted)
response = aws_srp.update_device_status(
    is_remembered=True,
    access_token=user.access_token,
    device_key=device_key
)
print(f"Device marked as remembered: {response}")

# Later, mark device as not remembered (untrusted)
response = aws_srp.update_device_status(
    is_remembered=False,
    access_token=user.access_token,
    device_key=device_key
)
print(f"Device marked as not remembered: {response}")

def forget_device(self, access_token: str, device_key: str) -> str:
    """
    Forget/remove a device from trusted devices.
    
    Args:
        access_token (str): Valid access token
        device_key (str): Device key to forget
        
    Returns:
        str: Response status from forget operation
        
    Actions:
        - Removes device from trusted device list
        - Invalidates device-specific credentials
        - Forces full authentication on next login
        - Clears device memory and tracking
        
    Use cases:
        - Device is lost or stolen
        - User no longer uses the device
        - Security incident requires device removal
        - User explicitly removes device trust
    """

# Forget a lost or stolen device
response = aws_srp.forget_device(
    access_token=user.access_token,
    device_key=lost_device_key
)
print(f"Device forgotten: {response}")

# Device will require full authentication next time

class AWSSRP:
    def __init__(self, username: str, password: str, pool_id: str, client_id: str,
                 device_key: str = None, device_group_key: str = None, 
                 device_password: str = None, **kwargs):
        """
        Initialize AWSSRP with device authentication support.
        
        Args:
            username (str): Username
            password (str): User password
            pool_id (str): User pool ID
            client_id (str): Client ID
            device_key (str, optional): Device key from device confirmation
            device_group_key (str, optional): Device group key
            device_password (str, optional): Device password
            
        Note:
            For device authentication, either all device parameters must be provided
            or none at all. Partial device parameters will raise ValueError.
        """
    
    def authenticate_user(self, client=None, client_metadata: dict = None) -> dict:
        """
        Authenticate user with optional device authentication.
        
        Returns:
            dict: Authentication tokens and device metadata
            
        Flow:
            1. Standard SRP authentication
            2. If device keys provided, performs device authentication
            3. Returns tokens with device confirmation if new device
            4. Handles device challenges automatically
        """

# Authenticate with remembered device
aws_srp = AWSSRP(
    username='user@example.com',
    password='password',
    pool_id='us-east-1_example',
    client_id='your-client-id',
    device_key='stored-device-key',
    device_group_key='stored-device-group-key',
    device_password='stored-device-password'
)

try:
    # Authentication with device credentials
    tokens = aws_srp.authenticate_user()
    print("Authentication successful with device credentials!")
    
except Exception as e:
    print(f"Device authentication failed: {e}")
    # Fall back to standard authentication
    aws_srp_fallback = AWSSRP(
        username='user@example.com',
        password='password',
        pool_id='us-east-1_example',
        client_id='your-client-id'
    )
    tokens = aws_srp_fallback.authenticate_user()

def register_device(username, password, device_name):
    """Complete device registration and confirmation flow."""
    
    # Step 1: Initial authentication
    aws_srp = AWSSRP(
        username=username,
        password=password,
        pool_id='your-pool-id',
        client_id='your-client-id'
    )
    
    # Step 2: Authenticate and check for new device
    tokens = aws_srp.authenticate_user()
    
    auth_result = tokens.get('AuthenticationResult', {})
    
    if 'NewDeviceMetadata' in auth_result:
        print("New device detected - setting up device authentication")
        
        # Step 3: Confirm device
        response, device_password = aws_srp.confirm_device(tokens, device_name)
        
        if response.status_code == 200:
            print("Device confirmed successfully!")
            
            # Step 4: Mark as remembered
            aws_srp.update_device_status(
                is_remembered=True,
                access_token=auth_result['AccessToken'],
                device_key=aws_srp.device_key
            )
            
            # Step 5: Store device credentials securely
            device_credentials = {
                'device_key': aws_srp.device_key,
                'device_group_key': aws_srp.device_group_key,
                'device_password': device_password,
                'device_name': device_name
            }
            
            # Store in secure storage (keychain, encrypted file, etc.)
            save_device_credentials(username, device_credentials)
            
            print(f"Device '{device_name}' registered and remembered")
            return device_credentials
            
        else:
            print(f"Device confirmation failed: {response}")
            return None
    else:
        print("No new device metadata - device may already be registered")
        return None

# Usage
device_creds = register_device('user@example.com', 'password', 'My Laptop')

class DeviceAwareAuth:
    """Authentication manager with device support."""
    
    def __init__(self, pool_id, client_id):
        self.pool_id = pool_id
        self.client_id = client_id
    
    def authenticate(self, username, password, device_credentials=None):
        """Authenticate with optional device credentials."""
        
        if device_credentials:
            # Try device-enhanced authentication first
            try:
                aws_srp = AWSSRP(
                    username=username,
                    password=password,
                    pool_id=self.pool_id,
                    client_id=self.client_id,
                    device_key=device_credentials['device_key'],
                    device_group_key=device_credentials['device_group_key'],
                    device_password=device_credentials['device_password']
                )
                
                tokens = aws_srp.authenticate_user()
                return tokens, "DEVICE_AUTH_SUCCESS"
                
            except Exception as e:
                print(f"Device authentication failed: {e}")
                # Fall through to standard authentication
        
        # Standard authentication
        aws_srp = AWSSRP(
            username=username,
            password=password,
            pool_id=self.pool_id,
            client_id=self.client_id
        )
        
        tokens = aws_srp.authenticate_user()
        
        # Check if device registration is needed
        if 'NewDeviceMetadata' in tokens.get('AuthenticationResult', {}):
            return tokens, "NEW_DEVICE_DETECTED"
        else:
            return tokens, "STANDARD_AUTH_SUCCESS"
    
    def setup_device(self, tokens, device_name, remember=True):
        """Set up device after authentication."""
        
        # This would typically be called by the same AWSSRP instance
        # that performed authentication
        aws_srp = AWSSRP(
            username="", password="",  # Not needed for device operations
            pool_id=self.pool_id,
            client_id=self.client_id
        )
        
        response, device_password = aws_srp.confirm_device(tokens, device_name)
        
        if response.status_code == 200 and remember:
            aws_srp.update_device_status(
                is_remembered=True,
                access_token=tokens['AuthenticationResult']['AccessToken'],
                device_key=aws_srp.device_key
            )
        
        return {
            'device_key': aws_srp.device_key,
            'device_group_key': aws_srp.device_group_key,
            'device_password': device_password,
            'device_name': device_name
        }

# Usage
auth_manager = DeviceAwareAuth('pool-id', 'client-id')

# Load stored device credentials
device_creds = load_device_credentials('user@example.com')

# Authenticate
tokens, status = auth_manager.authenticate(
    'user@example.com', 
    'password', 
    device_creds
)

if status == "NEW_DEVICE_DETECTED":
    # Set up new device
    new_device_creds = auth_manager.setup_device(
        tokens, 
        "My New Device", 
        remember=True
    )
    save_device_credentials('user@example.com', new_device_creds)
    print("New device set up successfully!")

elif status == "DEVICE_AUTH_SUCCESS":
    print("Authenticated with trusted device!")

elif status == "STANDARD_AUTH_SUCCESS":
    print("Standard authentication successful!")

def device_management_interface(username, access_token):
    """Interactive device management interface."""
    
    print(f"\n=== Device Management for {username} ===")
    print("1. List trusted devices")
    print("2. Forget a device")
    print("3. Update device status")
    print("4. Register new device")
    
    choice = input("Select option (1-4): ")
    
    if choice == "1":
        # List devices (would require additional API calls)
        print("Trusted devices:")
        devices = get_user_devices(username)  # Custom function
        for device in devices:
            print(f"- {device['name']} ({device['key']})")
    
    elif choice == "2":
        # Forget device
        device_key = input("Enter device key to forget: ")
        
        aws_srp = AWSSRP(
            username="", password="",
            pool_id='pool-id', client_id='client-id'
        )
        
        response = aws_srp.forget_device(access_token, device_key)
        print(f"Device forgotten: {response}")
    
    elif choice == "3":
        # Update device status
        device_key = input("Enter device key: ")
        remember = input("Remember device? (y/n): ").lower() == 'y'
        
        aws_srp = AWSSRP(
            username="", password="",
            pool_id='pool-id', client_id='client-id'
        )
        
        response = aws_srp.update_device_status(
            is_remembered=remember,
            access_token=access_token,
            device_key=device_key
        )
        print(f"Device status updated: {response}")
    
    elif choice == "4":
        print("Please authenticate on the new device to register it")

# Usage after authentication
device_management_interface('user@example.com', user.access_token)

import json
import keyring
from cryptography.fernet import Fernet

class SecureDeviceStorage:
    """Secure storage for device credentials."""
    
    def __init__(self, service_name="MyApp"):
        self.service_name = service_name
    
    def _get_encryption_key(self, username):
        """Get or create encryption key for user."""
        key_name = f"{username}_device_key"
        
        # Try to get existing key
        key = keyring.get_password(self.service_name, key_name)
        
        if not key:
            # Generate new key
            key = Fernet.generate_key().decode()
            keyring.set_password(self.service_name, key_name, key)
        
        return key.encode()
    
    def store_device_credentials(self, username, credentials):
        """Store device credentials securely."""
        encryption_key = self._get_encryption_key(username)
        fernet = Fernet(encryption_key)
        
        # Encrypt credentials
        credentials_json = json.dumps(credentials)
        encrypted_credentials = fernet.encrypt(credentials_json.encode())
        
        # Store in keyring
        cred_name = f"{username}_device_credentials"
        keyring.set_password(
            self.service_name, 
            cred_name, 
            encrypted_credentials.decode()
        )
    
    def load_device_credentials(self, username):
        """Load device credentials securely."""
        try:
            encryption_key = self._get_encryption_key(username)
            fernet = Fernet(encryption_key)
            
            # Get encrypted credentials
            cred_name = f"{username}_device_credentials"
            encrypted_credentials = keyring.get_password(
                self.service_name, 
                cred_name
            )
            
            if not encrypted_credentials:
                return None
            
            # Decrypt credentials
            decrypted_data = fernet.decrypt(encrypted_credentials.encode())
            credentials = json.loads(decrypted_data.decode())
            
            return credentials
            
        except Exception as e:
            print(f"Failed to load device credentials: {e}")
            return None
    
    def remove_device_credentials(self, username):
        """Remove stored device credentials."""
        cred_name = f"{username}_device_credentials"
        key_name = f"{username}_device_key"
        
        try:
            keyring.delete_password(self.service_name, cred_name)
            keyring.delete_password(self.service_name, key_name)
            print("Device credentials removed")
        except Exception as e:
            print(f"Failed to remove credentials: {e}")

# Usage
storage = SecureDeviceStorage("MyCognitoApp")

# Store credentials after device registration
device_creds = {
    'device_key': 'device-key-123',
    'device_group_key': 'group-key-456',
    'device_password': 'device-password-789',
    'device_name': 'My Laptop'
}

storage.store_device_credentials('user@example.com', device_creds)

# Load credentials for authentication
loaded_creds = storage.load_device_credentials('user@example.com')
if loaded_creds:
    print(f"Loaded device: {loaded_creds['device_name']}")