module.py

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

from cStringIO import StringIO
from datetime import timedelta
from email.mime.text import MIMEText
from threading import Thread
from uuid import uuid4
import base64
import hashlib
import json
import os
import smtplib

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
from functions import random_data
from functions import render_jinja_html
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,
        } )

@app.route( baseurl + '/reset_password' )
def password_reset():
    return render_template( 
        "password_reset.html",
        baseurl = baseurl,
        js = config.cdnjs,
        css = config.cdncss
    )

@app.route( baseurl + '/do_reset_password', methods = [ 'POST' ] )
def do_password_reset():
    email = request.form.get( "email", None )
    
    Thread( target = do_password_reset_thread, args = ( email, ) ).start()
    
    return jsonify( {
        'error': False,
        'message': 'OK'
    } )

def do_password_reset_thread( email ):
    q = config.db.query( 'SELECT id, username, email FROM users' )
    users = q.fetchall()
    
    for user in users:
        if not user[ 'email' ].startswith( "pbkdf2$" ):
            continue
        
        elif pbkdf2( email, user[ 'email' ] ):
            id = hashlib.sha512( random_data( 100 ) ).hexdigest()
            
            ####################################################################
            
            data = {
                'process': 'password_reset',
                'process_id': id,
                'user_id': user[ 'id' ]
            }
            data = json.dumps( data )
            data = base64.b64encode( data )
            
            config.redis_shared.set( "reset_" + id, data, ex = 24 * 3600 )
            
            ####################################################################
            
            email_content = render_jinja_html( 
                "templates", "email.html",
                id = id,
                url = config.domain + baseurl + "/reset_password_stage2"
            )
            
            msg = MIMEText( email_content, "html" )
            
            msg[ 'Subject' ] = 'ICNML - User password reset'
            msg[ 'From' ] = config.sender
            msg[ 'To' ] = email
            
            s = smtplib.SMTP( config.smtpserver )
            s.sendmail( config.sender, [ email ], msg.as_string() )
            s.quit()
            
            break
    
    else:
        print "email not found"

@app.route( baseurl + '/reset_password_stage2/<id>', methods = [ 'GET', 'POST' ] )
def password_reset_stage2( id ):
    id = str( id )
    data = config.redis_shared.get( "reset_" + id )
    
    if data != None:
        data = base64.b64decode( data )
        data = json.loads( data )
        
        password = request.form.get( "password", None )
        
        userid = data.get( "user_id", None )
        
        if password != None:
            password = pbkdf2( password, random_data( 50 ), 50000 )
            config.db.query( "UPDATE users SET password = %s WHERE id = %s", ( password, userid ) )
            config.db.commit()
            
            config.redis_shared.delete( "reset_" + id )
            
            return jsonify( {
                'error': False,
                'password_updated': True
            } )
            
        else:
            return render_template( 
                "password_reset_stage2.html",
                baseurl = baseurl,
                id = id,
                js = config.cdnjs,
                css = config.cdncss
            ) 
        
    else:
        return jsonify( {
            'error': True,
            'message': 'Reset procedure not found/expired'
        } )

################################################################################
#    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 )