WebAuthn Support

class WebAuthnRegisterForm(Form):
    """
    WebAuthn credential registration initiation form.
    
    Fields:
    - usage: Usage type for credential ('first', 'secondary')
    - name: Human-readable name for the credential
    """
    usage: SelectField
    name: StringField

class WebAuthnRegisterResponseForm(Form):
    """
    WebAuthn credential registration completion form.
    
    Fields:
    - credential: JSON credential response from authenticator
    - name: Human-readable name for the credential
    """
    credential: HiddenField
    name: StringField

class WebAuthnSigninForm(Form):
    """
    WebAuthn signin initiation form.
    
    Fields:
    - identity: User identity (email/username) for credential lookup
    """
    identity: StringField

class WebAuthnSigninResponseForm(Form):
    """
    WebAuthn signin completion form.
    
    Fields:
    - credential: JSON credential assertion from authenticator
    """
    credential: HiddenField

class WebAuthnDeleteForm(Form):
    """
    WebAuthn credential deletion form.
    
    Fields:
    - credential_id: ID of credential to delete
    """
    credential_id: HiddenField

class WebAuthnVerifyForm(Form):
    """
    WebAuthn verification form for sensitive operations.
    
    Fields:
    - credential: JSON credential assertion for verification
    """
    credential: HiddenField

class WebAuthnMixin:
    """
    Mixin for user models providing WebAuthn credential interface.
    """
    
    @property
    def webauthn_credentials(self):
        """
        Return list of WebAuthn credentials for user.
        
        Returns:
        List of WebAuthnCredential objects associated with user
        """
    
    def get_webauthn_credential_by_id(self, credential_id):
        """
        Get WebAuthn credential by credential ID.
        
        Parameters:
        - credential_id: Base64-encoded credential ID to search for
        
        Returns:
        WebAuthnCredential object if found, None otherwise
        """
    
    def get_webauthn_credential_options(self):
        """
        Get options for WebAuthn credential creation.
        
        Returns:
        Dictionary containing credential creation options including:
        - rp: Relying party information
        - user: User information for credential
        - challenge: Random challenge bytes
        - pubKeyCredParams: Supported algorithms
        - excludeCredentials: Existing credentials to exclude
        """

class WebauthnUtil:
    """
    WebAuthn operations and credential management utility class.
    """
    
    def __init__(self, app=None):
        """
        Initialize WebAuthn utility.
        
        Parameters:
        - app: Flask application instance (optional)
        """
    
    def init_app(self, app):
        """
        Initialize WebAuthn utility with Flask app.
        
        Parameters:
        - app: Flask application instance
        """
    
    def generate_challenge(self):
        """
        Generate cryptographic challenge for WebAuthn operations.
        
        Returns:
        Base64-encoded challenge bytes
        """
    
    def get_credential_creation_options(self, user):
        """
        Get credential creation options for user registration.
        
        Parameters:
        - user: User object for credential creation
        
        Returns:
        Dictionary with credential creation options
        """
    
    def get_credential_request_options(self, user=None):
        """
        Get credential request options for authentication.
        
        Parameters:
        - user: User object (optional, for usernameless authentication)
        
        Returns:
        Dictionary with credential request options
        """
    
    def register_credential(self, user, credential_response, name=None):
        """
        Register new WebAuthn credential for user.
        
        Parameters:
        - user: User object to associate credential with
        - credential_response: Credential response from authenticator
        - name: Human-readable name for credential (optional)
        
        Returns:
        WebAuthnCredential object if registration successful
        
        Raises:
        ValidationError if credential registration fails
        """
    
    def authenticate_credential(self, credential_response, user=None):
        """
        Authenticate using WebAuthn credential assertion.
        
        Parameters:
        - credential_response: Credential assertion from authenticator
        - user: User object for verification (optional)
        
        Returns:
        Tuple of (user, credential) if authentication successful, 
        (None, None) otherwise
        """
    
    def delete_credential(self, user, credential_id):
        """
        Delete WebAuthn credential for user.
        
        Parameters:
        - user: User object owning the credential
        - credential_id: ID of credential to delete
        
        Returns:
        True if credential deleted successfully, False otherwise
        """

class WebAuthnTfPlugin(TfPluginBase):
    """
    WebAuthn two-factor authentication plugin.
    """
    
    def get_setup_form(self):
        """
        Get form for WebAuthn 2FA setup.
        
        Returns:
        WebAuthnRegisterForm class for credential registration
        """
    
    def get_verify_form(self):
        """
        Get form for WebAuthn 2FA verification.
        
        Returns:
        WebAuthnVerifyForm class for credential verification
        """
    
    def setup_validate(self, form):
        """
        Validate WebAuthn 2FA setup.
        
        Parameters:
        - form: WebAuthnRegisterResponseForm with credential data
        
        Returns:
        True if setup validation successful, False otherwise
        """
    
    def verify(self, form):
        """
        Verify WebAuthn 2FA credential.
        
        Parameters:
        - form: WebAuthnVerifyForm with credential assertion
        
        Returns:
        True if verification successful, False otherwise
        """
    
    def delete(self, totp_secret):
        """
        Delete WebAuthn 2FA credential.
        
        Parameters:
        - totp_secret: Credential identifier for deletion
        
        Returns:
        True if deletion successful, False otherwise
        """

# Enable WebAuthn support
app.config['SECURITY_WEBAUTHN'] = True

# WebAuthn relying party configuration
app.config['SECURITY_WEBAUTHN_RP_NAME'] = 'Your App Name'
app.config['SECURITY_WEBAUTHN_RP_ID'] = 'example.com'

# Challenge configuration
app.config['SECURITY_WEBAUTHN_CHALLENGE_TTL'] = 300  # 5 minutes

# Credential algorithms (default includes ES256, PS256, RS256)
app.config['SECURITY_WEBAUTHN_ALGORITHMS'] = ['ES256', 'PS256', 'RS256']

# Authenticator attachment preference
app.config['SECURITY_WEBAUTHN_AUTHENTICATOR_ATTACHMENT'] = 'cross-platform'

# User verification requirement
app.config['SECURITY_WEBAUTHN_USER_VERIFICATION'] = 'preferred'

from flask import Flask
from flask_security import Security, WebAuthnMixin, UserMixin
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)
app.config['SECURITY_WEBAUTHN'] = True
app.config['SECURITY_WEBAUTHN_RP_NAME'] = 'My Secure App'
app.config['SECURITY_WEBAUTHN_RP_ID'] = 'myapp.com'

db = SQLAlchemy(app)

class User(db.Model, UserMixin, WebAuthnMixin):
    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(255), unique=True)
    password = db.Column(db.String(255))
    active = db.Column(db.Boolean(), default=True)

class WebAuthnCredential(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
    credential_id = db.Column(db.Text, unique=True)
    public_key = db.Column(db.Text)
    name = db.Column(db.String(100))
    usage = db.Column(db.String(20))
    created_at = db.Column(db.DateTime)

security = Security(app, user_datastore)

from flask_security import current_user, login_required

@app.route('/register-webauthn')
@login_required
def register_webauthn():
    # Display WebAuthn registration form
    return render_template('webauthn_register.html')

@app.route('/webauthn/register/begin', methods=['POST'])
@login_required  
def webauthn_register_begin():
    # Get credential creation options
    options = current_user.get_webauthn_credential_options()
    session['webauthn_challenge'] = options['challenge']
    return jsonify(options)

@app.route('/webauthn/register/complete', methods=['POST'])
@login_required
def webauthn_register_complete():
    credential_response = request.json
    challenge = session.get('webauthn_challenge')
    
    try:
        # Register the credential
        webauthn_util = current_app.extensions['security'].webauthn_util
        credential = webauthn_util.register_credential(
            current_user, 
            credential_response,
            name=request.form.get('name', 'Security Key')
        )
        
        return jsonify({'success': True, 'credential_id': credential.id})
    except ValidationError as e:
        return jsonify({'success': False, 'error': str(e)}), 400

@app.route('/webauthn-signin')
def webauthn_signin():
    return render_template('webauthn_signin.html')

@app.route('/webauthn/authenticate/begin', methods=['POST'])
def webauthn_authenticate_begin():
    identity = request.form.get('identity')
    user = lookup_identity(identity) if identity else None
    
    # Get authentication options
    webauthn_util = current_app.extensions['security'].webauthn_util
    options = webauthn_util.get_credential_request_options(user)
    session['webauthn_challenge'] = options['challenge']
    
    return jsonify(options)

@app.route('/webauthn/authenticate/complete', methods=['POST'])
def webauthn_authenticate_complete():
    credential_response = request.json
    challenge = session.get('webauthn_challenge')
    
    try:
        webauthn_util = current_app.extensions['security'].webauthn_util
        user, credential = webauthn_util.authenticate_credential(credential_response)
        
        if user:
            login_user(user)
            return jsonify({'success': True, 'redirect': '/dashboard'})
        else:
            return jsonify({'success': False, 'error': 'Authentication failed'}), 401
            
    except ValidationError as e:
        return jsonify({'success': False, 'error': str(e)}), 400

from flask_security import TwoFactorSetupForm

# Enable WebAuthn 2FA plugin
app.config['SECURITY_TWO_FACTOR_ENABLED_METHODS'] = ['authenticator', 'webauthn']

@app.route('/setup-2fa-webauthn')
@login_required
def setup_2fa_webauthn():
    """Setup WebAuthn as second factor."""
    if not current_user.tf_totp_secret:
        # User needs to complete initial 2FA setup first
        return redirect(url_for_security('two_factor_setup'))
    
    return render_template('setup_webauthn_2fa.html')

@app.route('/verify-2fa-webauthn')
@login_required
def verify_2fa_webauthn():
    """Verify WebAuthn credential for 2FA."""
    return render_template('verify_webauthn_2fa.html')

@app.route('/manage-webauthn')
@login_required
def manage_webauthn_credentials():
    """Display user's WebAuthn credentials."""
    credentials = current_user.webauthn_credentials
    return render_template('manage_webauthn.html', credentials=credentials)

@app.route('/delete-webauthn/<credential_id>', methods=['POST'])
@login_required
def delete_webauthn_credential(credential_id):
    """Delete a WebAuthn credential."""
    webauthn_util = current_app.extensions['security'].webauthn_util
    
    if webauthn_util.delete_credential(current_user, credential_id):
        flash('Security key removed successfully')
    else:
        flash('Failed to remove security key')
    
    return redirect(url_for('manage_webauthn_credentials'))

// WebAuthn credential registration
async function registerWebAuthnCredential() {
    try {
        // Get registration options from server
        const optionsResponse = await fetch('/webauthn/register/begin', {
            method: 'POST',
            headers: {'Content-Type': 'application/json'},
        });
        const options = await optionsResponse.json();
        
        // Convert base64 strings to ArrayBuffers
        options.challenge = base64ToArrayBuffer(options.challenge);
        options.user.id = base64ToArrayBuffer(options.user.id);
        
        // Create credential
        const credential = await navigator.credentials.create({
            publicKey: options
        });
        
        // Send credential to server
        const registrationResponse = await fetch('/webauthn/register/complete', {
            method: 'POST',
            headers: {'Content-Type': 'application/json'},
            body: JSON.stringify({
                id: credential.id,
                rawId: arrayBufferToBase64(credential.rawId),
                response: {
                    clientDataJSON: arrayBufferToBase64(credential.response.clientDataJSON),
                    attestationObject: arrayBufferToBase64(credential.response.attestationObject)
                }
            })
        });
        
        const result = await registrationResponse.json();
        if (result.success) {
            alert('Security key registered successfully!');
        } else {
            alert('Registration failed: ' + result.error);
        }
        
    } catch (error) {
        console.error('WebAuthn registration error:', error);
        alert('WebAuthn registration failed');
    }
}

// WebAuthn authentication
async function authenticateWithWebAuthn() {
    try {
        // Get authentication options
        const optionsResponse = await fetch('/webauthn/authenticate/begin', {
            method: 'POST',
            headers: {'Content-Type': 'application/json'},
        });
        const options = await optionsResponse.json();
        
        // Convert challenge to ArrayBuffer
        options.challenge = base64ToArrayBuffer(options.challenge);
        
        // Get assertion
        const assertion = await navigator.credentials.get({
            publicKey: options
        });
        
        // Send assertion to server
        const authResponse = await fetch('/webauthn/authenticate/complete', {
            method: 'POST',
            headers: {'Content-Type': 'application/json'},
            body: JSON.stringify({
                id: assertion.id,
                rawId: arrayBufferToBase64(assertion.rawId),
                response: {
                    clientDataJSON: arrayBufferToBase64(assertion.response.clientDataJSON),
                    authenticatorData: arrayBufferToBase64(assertion.response.authenticatorData),
                    signature: arrayBufferToBase64(assertion.response.signature),
                    userHandle: assertion.response.userHandle ? 
                        arrayBufferToBase64(assertion.response.userHandle) : null
                }
            })
        });
        
        const result = await authResponse.json();
        if (result.success) {
            window.location.href = result.redirect || '/dashboard';
        } else {
            alert('Authentication failed: ' + result.error);
        }
        
    } catch (error) {
        console.error('WebAuthn authentication error:', error);
        alert('WebAuthn authentication failed');
    }
}

// Utility functions for base64/ArrayBuffer conversion
function base64ToArrayBuffer(base64) {
    const binaryString = atob(base64);
    const bytes = new Uint8Array(binaryString.length);
    for (let i = 0; i < binaryString.length; i++) {
        bytes[i] = binaryString.charCodeAt(i);
    }
    return bytes.buffer;
}

function arrayBufferToBase64(buffer) {
    const bytes = new Uint8Array(buffer);
    let binary = '';
    for (let i = 0; i < bytes.byteLength; i++) {
        binary += String.fromCharCode(bytes[i]);
    }
    return btoa(binary);
}