lenticular_cloud2/lenticular_cloud/views/auth.py

from base64 import b64encode, b64decode, urlsafe_b64decode
import crypt
from datetime import datetime, timedelta
import jwt
from flask import request, url_for, jsonify, Blueprint, redirect, current_app, session
from flask.templating import render_template
from flask.typing import ResponseReturnValue
import logging
from ory_hydra_client import models as ory_hydra_m
from ory_hydra_client.api.o_auth_2 import get_o_auth_2_consent_request, accept_o_auth_2_consent_request, accept_o_auth_2_login_request, get_o_auth_2_login_request, accept_o_auth_2_login_request, accept_o_auth_2_logout_request, get_o_auth_2_login_request
from ory_hydra_client.models import TheRequestPayloadUsedToAcceptAConsentRequest, GenericError
from typing import Optional
from urllib.parse import urlparse
from uuid import uuid4, UUID
import webauthn
from webauthn.helpers.structs import (
    AuthenticatorSelectionCriteria,
    PublicKeyCredentialDescriptor,
    ResidentKeyRequirement,
    UserVerificationRequirement,
)

from ..model import db, User, PasskeyCredential
from ..form.auth import ConsentForm, LoginForm, RegistrationForm
from ..auth_providers import AUTH_PROVIDER_LIST
from ..hydra import hydra_service


logger = logging.getLogger(__name__)

auth_views = Blueprint('auth', __name__, url_prefix='/auth')


@auth_views.route('/consent', methods=['GET', 'POST'])
async def consent() -> ResponseReturnValue:
    """Always grant consent."""
    # DUMMPY ONLY

    form = ConsentForm()
    remember_for = 60*60*24*30  # remember for 30 days

    #try:
    consent_request = await get_o_auth_2_consent_request.asyncio(consent_challenge=request.args['consent_challenge'],_client=hydra_service.hydra_client)

    if consent_request is None or isinstance( consent_request, ory_hydra_m.OAuth20RedirectBrowserTo):
       return redirect(url_for('frontend.index'))

    requested_scope = consent_request.requested_scope
    requested_audiences = consent_request.requested_access_token_audience

    if form.validate_on_submit() or consent_request.skip:

        if type(consent_request.subject) != str:
            logger.error("not set subject `consent_request.subject`")
            return 'internal error', 500
        uid = UUID(consent_request.subject)
        user = User.query.get(uid)
        if user is None:
            logger.error("user not found, even if it should exist")
            return 'internal error', 500
        access_token = {
            'name': str(user.username),
            'preferred_username': str(user.username),
            'username': str(user.username),
            'email': str(user.email),
            'email_verified': True,
            #'given_name': str(user.username),
            #'family_name': '-',
            'groups': [group.name for group in user.groups]
        }
        id_token = {}
        if isinstance(requested_scope, list) and 'openid' in requested_scope:
            id_token = access_token
        body = TheRequestPayloadUsedToAcceptAConsentRequest(
                grant_scope= requested_scope,
                grant_access_token_audience= requested_audiences,
                remember= form.data['remember'],
                remember_for= remember_for,
                session= ory_hydra_m.PassSessionDataToAConsentRequest(
                    access_token= access_token,
                    id_token= id_token
                )
        )
        resp = await accept_o_auth_2_consent_request.asyncio(_client=hydra_service.hydra_client,
            json_body=body,
            consent_challenge=consent_request.challenge)
        if resp is None or isinstance( resp, GenericError):
            return 'internal error, could not forward request', 503
        return redirect(resp.redirect_to)
    return render_template(
            'auth/consent.html.j2',
            form=form,
            client=consent_request.client,
            requested_scope=requested_scope,
            requested_audiences=requested_audiences)


@auth_views.route('/login', methods=['GET', 'POST'])
async def login() -> ResponseReturnValue:
    secret_key = current_app.config['SECRET_KEY']
    public_url = urlparse(current_app.config['PUBLIC_URL'])
    
    login_challenge = request.args.get('login_challenge')
    if login_challenge is None:
        return 'login_challenge missing', 400
    login_request = await get_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client, login_challenge=login_challenge)
    if login_request is None or isinstance( login_request, ory_hydra_m.OAuth20RedirectBrowserTo):
        logger.exception("could not fetch login request")
        return redirect(url_for('frontend.index'))

    ## passkey
    options = webauthn.generate_authentication_options(
        rp_id = public_url.hostname,
        user_verification = UserVerificationRequirement.REQUIRED,
        challenge=webauthn.helpers.generate_challenge(32)
    )
    token = jwt.encode({
            'challenge': b64encode(options.challenge).decode(),
            'iat': datetime.utcnow() - timedelta(minutes=1),
            'nbf': datetime.utcnow(),
            'exp': datetime.utcnow() + timedelta(minutes=15),
        }, secret_key, algorithm="HS256"
    )

    ## 
    if login_request.skip:
        resp = await accept_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client,
            login_challenge=login_challenge,
            json_body=ory_hydra_m.HandledLoginRequestIsTheRequestPayloadUsedToAcceptALoginRequest(subject=login_request.subject))
        if resp is None or isinstance( resp, GenericError):
            return 'internal error, could not forward request', 503

        return redirect(resp.redirect_to)
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(username=form.data['name']).first() # type: Optional[User]
        if user:
            session['username'] = str(user.username)
        else:
            session['user'] = None
        session['auth_providers'] = []
        return redirect(
                url_for('auth.login_auth', login_challenge=login_challenge))
    return render_template(
        'auth/login.html.j2',
        form=form,
        options=webauthn.options_to_json(options),
        token=token,
        login_challenge=login_challenge,
    )


@auth_views.route('/login/auth', methods=['GET', 'POST'])
async def login_auth() -> ResponseReturnValue:
    login_challenge = request.args.get('login_challenge')
    if login_challenge is None:
        return 'missing login_challenge, bad request', 400
    login_request = await get_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client, login_challenge=login_challenge)
    if login_request is None:
        return redirect(url_for('frontend.index'))

    if 'username' not in session:
        return redirect(url_for('auth.login'))
    auth_forms = {}
    user = User.query.filter_by(username=session['username']).first_or_404()
    for auth_provider in AUTH_PROVIDER_LIST:
        form = auth_provider.get_form()
        if auth_provider.get_name() not in session['auth_providers'] and\
           auth_provider.check_auth(user, form):
            session['auth_providers'].append(auth_provider.get_name())
            session.modified = True

        if auth_provider.get_name() not in session['auth_providers']:
            auth_forms[auth_provider.get_name()]=form

    if len(session['auth_providers']) >= 1: # TODO should be to for TOTP working
        remember_me = True
#       if db_user is None:
#           db_user = User(username=session['username'])
#           db.session.add(db_user)
#           db.session.commit()

        subject = str(user.id)
        user.last_login = datetime.now()
        db.session.commit()
        resp = await accept_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client,
            login_challenge=login_challenge, json_body=ory_hydra_m.HandledLoginRequestIsTheRequestPayloadUsedToAcceptALoginRequest(
                subject=subject,
                remember=remember_me,
            ))
        if resp is None or isinstance( resp, GenericError):
            return 'internal error, could not forward request', 503
        return redirect(resp.redirect_to)
    return render_template('auth/login_auth.html.j2', forms=auth_forms)


@auth_views.route('/passkey/verify', methods=['POST'])
async def passkey_verify() -> ResponseReturnValue:
    secret_key = current_app.config['SECRET_KEY']
    public_url = urlparse(current_app.config['PUBLIC_URL'])


    data = request.get_json()
    token = jwt.decode(data['token'], secret_key, algorithms=['HS256'])
    challenge = urlsafe_b64decode(token['challenge'])
    credential = data['credential']
    credential_id = urlsafe_b64decode(credential['id'])

    login_challenge = data['login_challenge']
    if login_challenge is None:
        return 'missing login_challenge, bad request', 400
    login_request = await get_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client, login_challenge=login_challenge)
    if login_request is None:
        return redirect(url_for('frontend.index'))

    passkey = PasskeyCredential.query.filter(PasskeyCredential.credential_id == credential_id).first_or_404()
    
    result = webauthn.verify_authentication_response(
        credential = credential,
        expected_rp_id = public_url.hostname,
        expected_challenge = challenge,
        expected_origin = [ public_url.geturl() ],
        credential_public_key = passkey.credential_public_key,
        credential_current_sign_count = passkey.sign_count,
    )

    passkey.sign_count = result.new_sign_count
    passkey.last_used = datetime.utcnow()
    user = passkey.user
    user.last_login = datetime.now()

    subject = str(user.id)
    resp = await accept_o_auth_2_login_request.asyncio(_client=hydra_service.hydra_client,
        login_challenge=login_challenge, json_body=ory_hydra_m.HandledLoginRequestIsTheRequestPayloadUsedToAcceptALoginRequest(
            subject=subject,
            remember=True,
        ))
    if resp is None or isinstance( resp, GenericError):
        return 'internal error, could not forward request', 503

    db.session.commit()
    return jsonify({'redirect': resp.redirect_to})

@auth_views.route("/logout")
async def logout() -> ResponseReturnValue:
    logout_challenge = request.args.get('logout_challenge')
    if logout_challenge is None:
        return 'invalid request, logout_challenge not set', 400
    # TODO confirm
    resp = await accept_o_auth_2_logout_request.asyncio(_client=hydra_service.hydra_client, logout_challenge=logout_challenge)
    if resp is None or isinstance( resp, GenericError):
        return 'internal error, could not forward request', 503
    return redirect(resp.redirect_to)


@auth_views.route("/error", methods=["GET"])
def auth_error() -> ResponseReturnValue:
    error = request.args.get('error')
    error_description = request.args.get('error_description')

    return render_template('auth/error.html.j2', error=error, error_description=error_description)


@auth_views.route("/sign_up", methods=["GET"])
def sign_up():
    form = RegistrationForm()
    return render_template('auth/sign_up.html.j2', form=form)

@auth_views.route("/sign_up", methods=["POST"])
def sign_up_submit():
    form = RegistrationForm()
    if form.validate_on_submit():
        user = User()
        user.id = uuid4()
        user.username = form.data['username']
        user.password_hashed = crypt.crypt(form.data['password'])
        user.alternative_email = form.data['alternative_email']
        db.session.add(user)
        db.session.commit()
        return jsonify({})
    return jsonify({
            'status': 'error',
            'errors': form.errors
        })


@auth_views.route("/oob", methods=["GET"])
def oob_token():

    token_info = {
        'code': request.args.get('code', default="", type=str),
        'scope': request.args.get('scope', default="", type=str),
        'state': request.args.get('state', default="", type=str),
    }

    return render_template('auth/oob.html.j2', token_info=token_info)