notifications-admin/app/main/views/two_factor.py

from flask import (
    render_template,
    redirect,
    session,
    url_for,
    request
)

from flask_login import login_user, current_user
from app.main import main
from app.main.dao import users_dao, services_dao
from app.main.forms import TwoFactorForm


@main.route('/two-factor', methods=['GET', 'POST'])
def two_factor():
    # TODO handle user_email not in session
    try:
        user_id = session['user_details']['id']
    except KeyError:
        return redirect('main.sign_in')

    def _check_code(code):
        return users_dao.check_verify_code(user_id, code, "sms")

    form = TwoFactorForm(_check_code)

    if form.validate_on_submit():
        try:
            user = users_dao.get_user_by_id(user_id)
            services = services_dao.get_services(user_id).get('data', [])
            # Check if coming from new password page
            if 'password' in session['user_details']:
                user.set_password(session['user_details']['password'])
                users_dao.update_user(user)
            login_user(user, remember=True)
        finally:
            del session['user_details']

        next_url = request.args.get('next')
        if next_url and _is_safe_redirect_url(next_url):
            return redirect(next_url)

        if current_user.platform_admin:
            return redirect(url_for('main.show_all_services'))
        if len(services) == 1:
            return redirect(url_for('main.service_dashboard', service_id=services[0]['id']))
        else:
            return redirect(url_for('main.choose_service'))

    return render_template('views/two-factor.html', form=form)


# see http://flask.pocoo.org/snippets/62/
def _is_safe_redirect_url(target):
    from urllib.parse import urlparse, urljoin
    host_url = urlparse(request.host_url)
    redirect_url = urlparse(urljoin(request.host_url, target))
    return redirect_url.scheme in ('http', 'https') and \
        host_url.netloc == redirect_url.netloc