proxadmin/app/auth/routes.py

from flask import render_template, redirect, request, url_for, flash, session, abort, current_app
from flask_login import login_required, login_user, logout_user, current_user

from . import auth
from .. import db
from ..models import User, Transaction
from ..email import send_email
from .forms import LoginForm, TwoFAForm, RegistrationForm, ChangePasswordForm,PasswordResetRequestForm, PasswordResetForm
from ..decorators import admin_required, permission_required

from io import BytesIO
import pyqrcode

def get_google_auth(state=None, token=None):
    if token:
        return OAuth2Session(current_app.config['CLIENT_ID'], token=token)
    if state:
        return OAuth2Session(
            current_app.config['CLIENT_ID'],
            state=state,
            redirect_uri=current_app.config['REDIRECT_URI'])
    oauth = OAuth2Session(
        current_app.config['CLIENT_ID'],
        redirect_uri=current_app.config['REDIRECT_URI'],
        scope=current_app.config['SCOPE'])
    return oauth

@auth.before_app_request
def before_request():
    #print('session: %s' % str(session))
    if current_user.is_authenticated:
        current_user.ping()
        #print('request for {} from {}#{}'.format(request.endpoint, current_user.email, current_user.pid))
        if not current_user.confirmed and request.endpoint[:5] != 'auth.' and request.endpoint != 'static':
            print(request.endpoint)
            return redirect(url_for('auth.unconfirmed'))


@auth.route('/unconfirmed')
def unconfirmed():
    if current_user.is_anonymous or current_user.confirmed:
        return redirect(url_for('main.index'))
    return render_template('auth/unconfirmed.html')

@auth.route('/login', methods=['GET', 'POST'])
def login():
    page = { 'title': 'Login' }
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(email=form.email.data).first()

        if user is not None and user.verify_password(form.password.data):
            if user.active == False:
                flash('User disabled.')
                return redirect(url_for('main.index'))

            if user.twofactor:
                # redirect to the two-factor auth page, passing username in session
                session['email'] = user.email
                session['memberberry'] = form.remember_me.data
                return redirect(url_for('auth.twofactor'))

            #print('remember: ' + str(form.remember_me.data))
            login_user(user, form.remember_me.data)
            previp = user.last_ip
            if request.headers.getlist("X-Forwarded-For"):
                lastip = request.headers.getlist("X-Forwarded-For")[0]
            else:
                lastip = request.remote_addr
            user.last_ip = lastip
            db.session.add(user)
            db.session.commit()
            send_email(current_app.config['MAIL_USERNAME'], user.email + ' logged in.', 'auth/email/adm_loginnotify', user=user, ipaddr=lastip )
            #flash('Last Login: {} from {}'.format(user.last_seen.strftime("%a %d %B %Y %H:%M"), previp))
            flash('Last Login: {}'.format(user.last_seen.strftime("%a %d %B %Y %H:%M")))
            return redirect(request.args.get('next') or url_for('panel.dashboard'))
        else:
            flash('Invalid username or password.')

    return render_template('auth/login.html', page=page, form=form)


@auth.route('/twofactor', methods=['GET', 'POST'])
def twofactor():
    if 'email' not in session:
        abort(404)
    if 'memberberry' not in session:
        abort(404)
    page = { 'title': '2-Factor Login' }
    form = TwoFAForm()
    if form.validate_on_submit():
        user = User.query.filter_by(email=session['email']).first()
        del session['email']

        if user is not None and user.verify_totp(form.token.data):
            print('remember: ' + str(session['memberberry']))
            login_user(user, session['memberberry'])
            del session['memberberry']

            if request.headers.getlist("X-Forwarded-For"):
                lastip = request.headers.getlist("X-Forwarded-For")[0]
            else:
                lastip = request.remote_addr
            user.last_ip = lastip
            db.session.add(user)
            db.session.commit()
            #send_email(current_app.config['MAIL_USERNAME'], user.email + ' logged in.', 'auth/email/adm_loginnotify', user=user, ipaddr=lastip )
            return redirect(request.args.get('next') or url_for('panel.dashboard'))
        else:
            flash('Invalid token.')
    return render_template('auth/2fa.html', page=page, form=form)

@auth.route('/qrcode')
@login_required
def qrcode():
    #if 'email' not in session:
    #    abort(404)
    #user = User.query.filter_by(email=session['email']).first()
    #if user is None:
    #    abort(404)

    # for added security, remove username from session
    #del session['email']

    # render qrcode for FreeTOTP
    url = pyqrcode.create(current_user.get_totp_uri())
    stream = BytesIO()
    url.svg(stream, scale=6)
    return stream.getvalue(), 200, {
        'Content-Type': 'image/svg+xml',
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'Expires': '0'}

@auth.route("/logout", methods=['GET'])
@login_required
def logout():
    logout_user()
    flash('You have logged out')
    return redirect(url_for('main.index'))

@auth.route('/register', methods=['GET', 'POST'])
@admin_required
def register():
    #print(current_app.secret_key)
    page = { 'title': 'Register' }
    form = RegistrationForm()
    if form.validate_on_submit():
        user = User(email=form.email.data, password=form.password.data, wallet=current_app.config['REGISTER_BONUS'])
        db.session.add(user)
        db.session.commit()
        #transaction = Transaction(user_id=int(user.pid), description='Registered account bonus', value=current_app.config['REGISTER_BONUS'])
        #db.session.add(transaction)
        #db.session.commit()
        token = user.generate_confirmation_token()
        send_email(user.email, 'Потвърдете Вашата регистрация', 'auth/email/confirm', user=user, token=token)
        #notify admin
        newip = request.remote_addr
        if request.headers.getlist("X-Forwarded-For"):
            newip = request.headers.getlist("X-Forwarded-For")[0]
        else:
            newip = request.remote_addr
        send_email(current_app.config['MAIL_USERNAME'], user.email + ' registered!', 'auth/email/adm_regnotify', user=user, ipaddr=newip )
        flash('Благодарим за регистрацията! Моля проверете вашият email за потвърждение')
        return redirect(url_for('auth.login'))
    return render_template('auth/register.html', page=page, form=form)

@auth.route('/confirm/<token>')
@login_required
def confirm(token):
    if current_user.confirmed:
        return redirect(url_for('main.index'))
    if current_user.confirm(token):
        flash('Вашият акаунт е потвърден. Благодаря!')
    else:
        flash('Времето за потвърждение на вашият код изтече.')
    return redirect(url_for('main.index'))

@auth.route('/confirm')
@login_required
def resend_confirmation():
    token = current_user.generate_confirmation_token()
    send_email(current_user.email, 'Confirm_your_account',
               'auth/email/confirm', user=current_user, token=token)
    flash('New confirmation code was sent.')
    return redirect(url_for('main.index'))

@auth.route('/change-password', methods=['GET', 'POST'])
@login_required
def change_password():
    form = ChangePasswordForm()
    if form.validate_on_submit():
        if current_user.verify_password(form.old_password.data):
            current_user.password = form.password.data
            db.session.add(current_user)
            db.session.commit()
            flash('Your password was changed')
            return redirect(url_for('main.index'))
        else:
            flash('Wrong password.')
    return render_template("auth/change_password.html", form=form)

@auth.route('/reset', methods=['GET', 'POST'])
def password_reset_request():
    if not current_user.is_anonymous:
        return redirect(url_for('main.index'))
    form = PasswordResetRequestForm()
    if form.validate_on_submit():
        user = User.query.filter_by(email=form.email.data).first()
        if user:
            token = user.generate_reset_token()
            send_email(user.email, 'Reset Your Password',
                       'auth/email/reset_password',
                       user=user, token=token,
                       next=request.args.get('next'))
            flash('An email with instructions to reset your password has been sent to you.')
        return redirect(url_for('auth.login'))
    return render_template('auth/reset_password.html', form=form)

@auth.route('/reset/<token>', methods=['GET', 'POST'])
def password_reset(token):
    if not current_user.is_anonymous:
        return redirect(url_for('main.index'))
    form = PasswordResetForm()
    if form.validate_on_submit():
        user = User.query.filter_by(email=form.email.data).first()
        if user is None:
            return redirect(url_for('main.index'))
        if user.reset_password(token, form.password.data):
            flash('Your password has been updated.')
            return redirect(url_for('auth.login'))
        else:
            return redirect(url_for('main.index'))
    return render_template('auth/reset_password.html', form=form)