flask_photo_scaling_app/app/auth/auth.py

73 lines
3.2 KiB
Python

#!/usr/bin/env python3
from flask import Blueprint, redirect, url_for, session, flash, render_template
from flask_login import current_user, login_user, logout_user
from app.forms import LoginForm, GetTotp
from app.models import Contributor
from pyotp.totp import TOTP
auths = Blueprint(
"auths", __name__, template_folder="templates"
)
@auths.route("/two-factor-input", methods=["GET", "POST"])
def two_factor_input():
if current_user.is_authenticated or 'id' not in session:
return redirect(url_for('proute.index'))
contributor = Contributor.query.get(session['id'])
if contributor is None:
return redirect(url_for('proute.index'))
form = GetTotp()
if form.validate_on_submit():
if TOTP(contributor.totp_key).verify(int(form.totp_code.data), valid_window=5):
login_user(contributor, remember=session['remember_me'])
flash("Congratulations, you are now logged in!")
return redirect(url_for('proute.index'))
else:
flash("Oops, the pin was wrong")
form.totp_code.data = None
return render_template('two_factor_input.html', form=form, inst="Code was wrong, try again?")
return render_template('two_factor_input.html', form=form, inst="Enter Auth Code")
@auths.route("/login", methods=["GET", "POST"])
def login():
if current_user.is_authenticated:
return redirect(url_for('proute.index'))
form = LoginForm()
if form.validate_on_submit():
contributor_by_name = Contributor.query.filter_by(name=form.username.data).first()
contributor_by_email = Contributor.query.filter_by(email=form.email.data).first()
if contributor_by_name is not None and contributor_by_name.check_password(form.password.data):
if contributor_by_name.use_totp:
session['id'] = contributor_by_name.id
session['remember_me'] = form.remember_me.data
return redirect(url_for('auths.two_factor_input'))
else:
login_user(contributor_by_name, remember=form.remember_me.data)
flash("Congratulations, you are now logged in!")
return redirect(url_for('proute.index'))
elif contributor_by_email is not None and contributor_by_email.check_password(form.password.data):
if contributor_by_email.use_totp:
session['id'] = contributor_by_email.id
session['remember_me'] = form.remember_me.data
return redirect(url_for('auths.two_factor_input'))
else:
login_user(contributor_by_email, remember=form.remember_me.data)
flash("Congratulations, you are now logged in!")
return redirect(url_for('proute.index'))
else:
flash("Error Invalid Contributor (Username or Email) or Password")
return(redirect(url_for('auths.login')))
return render_template('login.html', title='Sign In', form=form)
@auths.route("/logout")
def logout():
is_authenticated = current_user.is_authenticated
logout_user()
if is_authenticated:
flash("Congratulations, you are now logged out!")
return redirect(url_for('proute.index'))