from hmac import new from unicodedata import category from flask import Blueprint, render_template, request, flash, redirect, url_for from flask_login import login_user, login_required, logout_user, current_user from werkzeug.security import generate_password_hash, check_password_hash from argon2 import PasswordHasher from .models import User from . import db from .forms import LoginForm, RegForm, MFAForm, MFAForm2, SearchForm from app.lib import email # MFA import pyotp accounts = Blueprint('accounts', __name__) @accounts.route('/login', methods=['GET', 'POST']) def login(): form = LoginForm() seform = SearchForm() if form.validate_on_submit(): email = form.email.data challenge_passwd = form.passwd.data # Check fo user in User table user = User.query.filter_by(email=email).first() # If there's a user if user: ##: ph = PasswordHasher() if ph.verify(user.password, challenge_passwd): return redirect(url_for('accounts.mfa', user_chal = user.id)) # passes user to mfa else: flash('Unsucessful Login!', category='error') else: flash('Unsucessful Login!', category='error') return render_template('login.html', user = current_user, form = form, seform = SearchForm()) @accounts.route('/logout') @login_required def logout(): logout_user() return redirect(url_for('accounts.login')) @accounts.route('/register', methods=['GET', 'POST']) def register(): pass_list = list() form = RegForm() seform = SearchForm() if form.validate_on_submit(): email = form.email.data username = form.username.data passwd_1 = form.passwd_1.data passwd_2 = form.passwd_2.data # Basic User Input Checks email_check = User.query.filter_by(email=email).first() if len(email) < 1: flash('Your Email must be longer than 0 characters.', category='error') elif email_check: flash('This Email is already taken', category='error') else: pass_list.append('p') if len(username) < 1: flash('Username must be something', category='error') else: pass_list.append('p') if len(passwd_1) < 12 or len(passwd_2) < 12: flash('Your Password must be longer than or equal to 12 characters.', category='error') else: if passwd_1 != passwd_2: flash('Your Passwords must match!', category='error') else: if len(pass_list) == 2: ##: ph = PasswordHasher() npasswd = ph.hash(passwd_2) new_user = User(email=email, username=username, password=npasswd ) db.session.add(new_user) db.session.commit() flash('Account Registration Successful!', category='success') return redirect(url_for('dashboards.market')) else: flash('Registration Failed', category='error') return render_template("register.html", user = current_user, form = form, seform = seform) @accounts.route('/mfa', methods=['GET', 'POST']) def mfa(): form = MFAForm() form2 = MFAForm2() user_chal = request.args['user_chal'] user = User.query.filter_by(id = user_chal).first() seform = SearchForm() # check for existing totphash if not user.totphash: # generate random secret key for auth secret = pyotp.random_base32() # add to User table and show this secret next time. dbcall = User.query.filter_by(id = user.id).first() dbcall.totphash = secret db.session.commit() flash('Generated new TOTP Secret', category='success') else: # use existing hash secret = user.totphash challenge_answer = int(pyotp.TOTP(secret).now()) if form.validate_on_submit(): otp = int(form.otp.data) # checks MFA if challenge_answer == otp: flash('Login Successful!', category='sucess') login_user(user, remember=True) return redirect(url_for('dashboards.market')) else: flash('Login Unsuccessful!', category='error') return redirect(url_for('accounts.mfa')) if form2.validate_on_submit(): if form2.send_key.data: email.send_email(user.email, user.totphash) elif form2.reset_key.data: # generate random secret key for auth secret = pyotp.random_base32() # add to User table and show this secret next time. dbcall = User.query.filter_by(id = user.id).first() dbcall.totphash = secret db.session.commit() email.send_email(user.email, user.totphash) return render_template('mfa.html', secret = secret, form = form, user = user, seform = seform, form2 = form2)