module.py

#!/usr/bin/python
# -*- coding: UTF-8 -*-

from cStringIO import StringIO
from datetime import timedelta
from uuid import uuid4
import os

from flask import Flask
from flask import jsonify
from flask import render_template, send_from_directory 
from flask import request
from flask import send_file
from flask import session
from flask import url_for
from flask_compress import Compress
from flask_session import Session
from pyotp import random_base32
from werkzeug import redirect
import pyotp
import qrcode

from functions import pbkdf2
import config

################################################################################

app = Flask( __name__ )
app.config.from_pyfile( 'config.py' )

Compress( app )
Session( app )

debug = os.environ.get( "DEBUG", False )
baseurl = os.environ.get( "BASEURL", "" )

################################################################################
#    Generic routing

@app.route( baseurl + '/ping' )
def ping():
    return "pong"

################################################################################
#    App serving

@app.route( baseurl + '/app/<path>' )
def send_app_files( path ):
    return send_from_directory( 'app', path )

################################################################################
#    Sessions

@app.before_request
def renew_session():
    session.permanent = True
    app.permanent_session_lifetime = timedelta( seconds = config.session_timeout )

@app.route( baseurl + '/logout' )
def logout():
    session.clear()
    return redirect( url_for( 'home' ) )

@app.route( baseurl + '/login' )
def login():
    session.clear()
    session[ 'stage' ] = 'password'
    
    return render_template( 
        "login.html",
        baseurl = baseurl,
        js = config.cdnjs,
        css = config.cdncss
    )

@app.route( baseurl + '/do_login', methods = [ 'POST' ] )
def do_login():
    if session[ 'stage' ] == 'password' or not 'stage' in session:
        q = config.db.query( 'SELECT * FROM users WHERE username = %s', ( request.form.get( "username" ), ) )
        user = q.fetchone()
        
        if user == None:
            return jsonify( {
                'error': False,
                'logged': False
            } )
        
        form_password = request.form.get( "password", None )
        
        if form_password == None or not pbkdf2( form_password, user[ 'password' ] ):
            return jsonify( {
                'error': False,
                'logged': False,
            } )
        
        elif not user[ 'active' ]:
            return jsonify( {
                'error': False,
                'logged': False,
                'message': 'Your account is not activated. Please contact an administrator.'
            } )
        
        else:
            session[ 'session_id' ] = str( uuid4() )
            session[ 'username' ] = user[ 'username' ]
            session[ 'password_check' ] = 'ok'
        
        if user[ 'must_use_totp' ]:
            session[ 'stage' ] = 'totp'
            
            return jsonify( {
                'error': False,
                'must_use_totp': True
            } )
        else:
            session[ 'logged' ] = True
    
    elif session[ 'stage' ] == 'totp':
        q = config.db.query( 'SELECT username, totp FROM users WHERE username = %s', ( session[ 'username' ], ) )
        user = q.fetchone()
        
        if not pyotp.TOTP( user[ 'totp' ] ).verify( request.form[ "totp" ], valid_window = 1 ):
            return jsonify( {
                'error': False,
                'logged': False,
                'message': 'Wrong TOTP'
            } )
        
        else:
            session[ 'logged' ] = True
    
    ############################################################################
    
    if session.get( 'logged', False ) and session.get( 'username', False ):
        return jsonify( {
            'error': False,
            'logged': True,
        } )
    
    else:
        session.clear()
        return jsonify( {
            'error': False,
            'logged': False,
        } )

################################################################################
#    QR Code generation

def renew_secret():
    secret = random_base32( 40 )
    session[ 'secret' ] = secret
    
    return secret

def get_secret():
    secret = session.get( "secret", None )
    if secret == None:
        secret = renew_secret()
    
    return secret

@app.route( baseurl + '/set_secret' )
def set_secret():
    config.db.query( "UPDATE users SET totp = %s WHERE username = %s", ( session[ 'secret' ], session[ 'username' ], ) )
    config.db.commit()
    
    return jsonify( {
        'error': False
    } )

@app.route( baseurl + '/secret' )
def request_secret():
    get_secret()
    
    return jsonify( {
        'error': False,
        'secret': session[ 'secret' ]
    } )

@app.route( baseurl + '/new_secret' )
def request_renew_secret():
    renew_secret()
    
    return jsonify( {
        'error': False,
        'secret': session[ 'secret' ]
    } )

@app.route( baseurl + '/qrcode' )
def send_qrcode():
    img = qrcode.make( 'otpauth://totp/ICNML%20' + session[ 'username' ] + '?secret=' + get_secret() )
    
    temp = StringIO()
    img.save( temp, format = "png" )
    temp.seek( 0 )
    
    return send_file( temp, mimetype = 'image/png' )
    
@app.route( baseurl + '/user_qrcode' )
def serve_qrcode():
    if session.get( "logged", False ):
        return render_template( 
            "qrcode.html",
            baseurl = baseurl,
            secret = get_secret(),
            js = config.cdnjs,
            css = config.cdncss
        )
    
    else:
        return redirect( url_for( 'home' ) )

################################################################################
#    Home page

@app.route( baseurl + '/' )
def home():
    if not session.get( "logged", False ):
        return redirect( url_for( 'login' ) )
    
    else:
        return render_template( 
            "index.html",
            baseurl = baseurl,
            js = config.cdnjs,
            css = config.cdncss,
            session_timeout = config.session_timeout
        )

################################################################################
#    Main startup

if __name__ == '__main__':
    app.run( debug = debug, host = "0.0.0.0", threaded = True )