<filename>authentication/totp.py
import onetimepass
import pyqrcode
from io import BytesIO
from flask import Flask, render_template, session, request, redirect, url_for
from DAL import user_dal

def get_totp_uri(username, otp_secret):
    return 'otpauth://totp/Vul-Project:{0}?secret={1}&issuer=Vul-Project'.format(username, otp_secret)

def verify_totp():
    userid = session['userid']
    if user_dal.check_disabled_user_by_id(userid) == True:
        return render_template('login.html', error='This account has been disable due to many failure on login or verify token')
    token = request.form['token']
    otp_secret = session['user']['otp_secret']
    if token is None:
        return render_template('verify-token.html', error='Token is empty')
    verify = onetimepass.valid_totp(token, otp_secret)
    if verify == True:
        username = session['user']['username']
        session['verified'] = True
        user_dal.increase_number_of_failure(username, 0)
        return redirect(url_for('main'))
    else:
        username = session['user']['username']
        user_dal.increase_number_of_failure(username, 1)
        return render_template('verify-token.html', error='Failed to verify token')

def two_factor_setup():
    return render_template('two-factor-setup.html'), 200, {
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'Expires': '0'}

def qrcode():
    username = session['user']['username']
    otp_secret = session['user']['otp_secret']
    url = pyqrcode.create(get_totp_uri(username, otp_secret))
    stream = BytesIO()
    url.svg(stream, scale=3)
    userid = session['userid']
    user_dal.change_status_first_login(userid, False)
    session.modified = True
    session['user']['firstlogin'] = False
    return stream.getvalue(), 200, {
        'Content-Type': 'image/svg+xml',
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'Expires': '0'}

<filename>authentication/totp.py
import onetimepass
import pyqrcode
from io import BytesIO
from flask import Flask, render_template, session, request, redirect, url_for
from DAL import user_dal

def get_totp_uri(username, otp_secret):
    return 'otpauth://totp/Vul-Project:{0}?secret={1}&issuer=Vul-Project'.format(username, otp_secret)

def verify_totp():
    userid = session['userid']
    if user_dal.check_disabled_user_by_id(userid) == True:
        return render_template('login.html', error='This account has been disable due to many failure on login or verify token')
    token = request.form['token']
    otp_secret = session['user']['otp_secret']
    if token is None:
        return render_template('verify-token.html', error='Token is empty')
    verify = onetimepass.valid_totp(token, otp_secret)
    if verify == True:
        username = session['user']['username']
        session['verified'] = True
        user_dal.increase_number_of_failure(username, 0)
        return redirect(url_for('main'))
    else:
        username = session['user']['username']
        user_dal.increase_number_of_failure(username, 1)
        return render_template('verify-token.html', error='Failed to verify token')

def two_factor_setup():
    return render_template('two-factor-setup.html'), 200, {
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'Expires': '0'}

def qrcode():
    username = session['user']['username']
    otp_secret = session['user']['otp_secret']
    url = pyqrcode.create(get_totp_uri(username, otp_secret))
    stream = BytesIO()
    url.svg(stream, scale=3)
    userid = session['userid']
    user_dal.change_status_first_login(userid, False)
    session.modified = True
    session['user']['firstlogin'] = False
    return stream.getvalue(), 200, {
        'Content-Type': 'image/svg+xml',
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'Expires': '0'}
