artFi/app/accounts.py

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 .models import User
from . import db

from .forms import LoginForm, RegForm, MFAForm

# MFA
import pyotp

accounts = Blueprint('accounts', __name__)

@accounts.route('/login', methods=['GET', 'POST'])
def login():
    form = LoginForm()

    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:
            if check_password_hash(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)
    
@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()
    
    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) < 8 or len(passwd_2) < 8:
            flash('Your Password must be longer than or equal to 8 characters.', category='error')
        else:
            if passwd_1 != passwd_2:
                flash('Your Passwords must match!', category='error')
            else:
                if len(pass_list) == 2:
                    new_user = User(email=email, 
                                    username=username,
                                    password=generate_password_hash(passwd_2, method='sha256')
                                    )
                    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)


@accounts.route('/mfa', methods=['GET', 'POST'])
def mfa():
    form = MFAForm()
    user_chal = request.args['user_chal']
    user = User.query.filter_by(id = user_chal).first()

    # 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: # create a new totphash
        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'))


    return render_template('mfa.html', secret = secret, form = form, user = user)