app-MAIL-temp/app/auth/views/proton.py

import requests
from flask import request, session, redirect, flash, url_for
from flask_limiter.util import get_remote_address
from flask_login import current_user
from requests_oauthlib import OAuth2Session

from app.auth.base import auth_bp
from app.auth.views.login_utils import after_login
from app.config import (
    PROTON_BASE_URL,
    PROTON_CLIENT_ID,
    PROTON_CLIENT_SECRET,
    PROTON_EXTRA_HEADER_NAME,
    PROTON_EXTRA_HEADER_VALUE,
    PROTON_VALIDATE_CERTS,
    URL,
)
from app.log import LOG
from app.models import ApiKey, User
from app.proton.proton_client import HttpProtonClient, convert_access_token
from app.proton.proton_callback_handler import (
    ProtonCallbackHandler,
    Action,
)
from app.proton.utils import get_proton_partner
from app.utils import sanitize_next_url

_authorization_base_url = PROTON_BASE_URL + "/oauth/authorize"
_token_url = PROTON_BASE_URL + "/oauth/token"

# need to set explicitly redirect_uri instead of leaving the lib to pre-fill redirect_uri
# when served behind nginx, the redirect_uri is localhost... and not the real url
_redirect_uri = URL + "/auth/proton/callback"

SESSION_ACTION_KEY = "oauth_action"
SESSION_STATE_KEY = "oauth_state"


def get_api_key_for_user(user: User) -> str:
    ak = ApiKey.create(
        user_id=user.id,
        name="Created via Login with Proton on mobile app",
        commit=True,
    )
    return ak.code


def extract_action() -> Action:
    action = request.args.get("action")
    if action is not None:
        if action == "link":
            return Action.Link
        else:
            raise Exception(f"Unknown action: {action}")
    return Action.Login


def get_action_from_state() -> Action:
    oauth_action = session[SESSION_ACTION_KEY]
    if oauth_action == Action.Login.value:
        return Action.Login
    elif oauth_action == Action.Link.value:
        return Action.Link
    raise Exception(f"Unknown action in state: {oauth_action}")


@auth_bp.route("/proton/login")
def proton_login():
    if PROTON_CLIENT_ID is None or PROTON_CLIENT_SECRET is None:
        return redirect(url_for("auth.login"))

    next_url = sanitize_next_url(request.args.get("next"))
    if next_url:
        session["oauth_next"] = next_url
    elif "oauth_next" in session:
        del session["oauth_next"]

    mode = request.args.get("mode", "session")
    if mode == "apikey":
        session["oauth_mode"] = "apikey"
    else:
        session["oauth_mode"] = "session"

    proton = OAuth2Session(PROTON_CLIENT_ID, redirect_uri=_redirect_uri)
    authorization_url, state = proton.authorization_url(_authorization_base_url)

    # State is used to prevent CSRF, keep this for later.
    session[SESSION_STATE_KEY] = state
    session[SESSION_ACTION_KEY] = extract_action().value
    return redirect(authorization_url)


@auth_bp.route("/proton/callback")
def proton_callback():
    if SESSION_STATE_KEY not in session or SESSION_STATE_KEY not in session:
        flash("Invalid state, please retry", "error")
        return redirect(url_for("auth.login"))
    if PROTON_CLIENT_ID is None or PROTON_CLIENT_SECRET is None:
        return redirect(url_for("auth.login"))

    # user clicks on cancel
    if "error" in request.args:
        flash("Please use another sign in method then", "warning")
        return redirect("/")

    proton = OAuth2Session(
        PROTON_CLIENT_ID,
        state=session[SESSION_STATE_KEY],
        redirect_uri=_redirect_uri,
    )

    def check_status_code(response: requests.Response) -> requests.Response:
        if response.status_code != 200:
            raise Exception(
                f"Bad Proton API response [status={response.status_code}]: {response.json()}"
            )
        return response

    proton.register_compliance_hook("access_token_response", check_status_code)

    headers = None
    if PROTON_EXTRA_HEADER_NAME and PROTON_EXTRA_HEADER_VALUE:
        headers = {PROTON_EXTRA_HEADER_NAME: PROTON_EXTRA_HEADER_VALUE}

    try:
        token = proton.fetch_token(
            _token_url,
            client_secret=PROTON_CLIENT_SECRET,
            authorization_response=request.url,
            verify=PROTON_VALIDATE_CERTS,
            method="GET",
            include_client_id=True,
            headers=headers,
        )
    except Exception as e:
        LOG.warning(f"Error fetching Proton token: {e}")
        flash("There was an error in the login process", "error")
        return redirect(url_for("auth.login"))

    credentials = convert_access_token(token["access_token"])
    action = get_action_from_state()

    proton_client = HttpProtonClient(
        PROTON_BASE_URL, credentials, get_remote_address(), verify=PROTON_VALIDATE_CERTS
    )
    handler = ProtonCallbackHandler(proton_client)
    proton_partner = get_proton_partner()

    if action == Action.Login:
        res = handler.handle_login(proton_partner)
    elif action == Action.Link:
        res = handler.handle_link(current_user, proton_partner)
    else:
        raise Exception(f"Unknown Action: {action.name}")

    if res.flash_message is not None:
        flash(res.flash_message, res.flash_category)

    if session.get("oauth_mode", "session") == "apikey":
        apikey = get_api_key_for_user(res.user)
        return redirect(f"auth.simplelogin://callback?apikey={apikey}")

    if res.redirect_to_login:
        return redirect(url_for("auth.login"))

    if res.redirect:
        return after_login(res.user, res.redirect, login_from_proton=True)

    next_url = session.get("oauth_next")
    return after_login(res.user, next_url, login_from_proton=True)