functions.py

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

import base64
import cPickle
from cStringIO import StringIO
import functools
import hashlib
import json
import smtplib
import time

from flask.globals import session
from PIL import Image, ImageDraw, ImageFont

import config
import utils

################################################################################
#    DEK

def get_dek_from_submissionid( submission_id ):
    """
        Get the Data Encryption Key related to a submission id folder.
    """
    try:
        sql = """
            SELECT donor_dek.dek
            FROM donor_dek
            LEFT JOIN users ON users.username = donor_dek.donor_name
            LEFT JOIN submissions ON submissions.donor_id = users.id
            WHERE submissions.uuid = %s
            LIMIT 1
        """
        dek = config.db.query_fetchone( sql, ( submission_id, ) )[ "dek" ]
        if dek == None:
            raise
        else:
            return dek
    
    except:
        try:
            return session[ "dek_{}".format( submission_id ) ]
        
        except:
            return None

def do_decrypt_dek( data, submission_id ):
    """
        AES encrypt the data with the Data Encryption Key related to the donor.
    """
    dek = get_dek_from_submissionid( submission_id )
    return utils.aes.do_decrypt( data, dek )

def do_encrypt_dek( data, submission_id ):
    """
        AES decrypt the data with the Data Encryption Key related to the donor.
    """
    dek = get_dek_from_submissionid( submission_id )
    return utils.aes.do_encrypt( data, dek )

def dek_check( submission_id ):
    if dek_exists( submission_id ):
        return True
    
    else:
        try:
            return dek_submitte_recreate_session( submission_id )
        except:
            return False

def dek_exists( submission_id ):
    if get_dek_from_submissionid( submission_id ) == None:
        return False
    else:
        return True

def dek_submitte_recreate_session( submission_id ):
    sql = """
        SELECT
            donor_dek.salt,
            donor_dek.dek_check,
            donor_dek.donor_name as username,
            submissions.email_aes as email
        FROM donor_dek
        LEFT JOIN users ON users.username = donor_dek.donor_name
        LEFT JOIN submissions ON submissions.donor_id = users.id
        WHERE submissions.uuid = %s
        LIMIT 1
    """
    user = config.db.query_fetchone( sql, ( submission_id, ) )
    
    username = user[ "username" ]
    email = do_decrypt_user_session( user[ "email" ] )
    dek_salt = user[ "salt" ]
    
    _, dek, _ = dek_generate( username = username, email = email, salt = dek_salt )
    
    to_check = utils.aes.do_decrypt( user[ "dek_check" ], dek )
    to_check = json.loads( to_check )
    
    if to_check[ "value" ] == "ok":
        session[ "dek_{}".format( submission_id ) ] = dek
        return True
    
    else:
        return False

def dek_generate( **kwargs ):
    if "email" in kwargs:
        email = kwargs[ "email" ]
        email = utils.hash.pbkdf2( email, "icnml_user_DEK" ).hash( True )
    elif "email_hash" in kwargs:
        email = kwargs[ "email_hash" ]
    else:
        raise Exception( "need the email or hashed_email" )
    
    if "username" in kwargs:
        username = kwargs[ "username" ]
    else:
        raise Exception( "need the username" )
    
    dek_salt = kwargs.get( "salt", utils.rand.random_data( config.DEK_SALT_LENGTH ) )
    dek = utils.hash.pbkdf2( 
        "{}:{}".format( username, email, ),
        dek_salt,
        iterations = config.DEK_NB_ITERATIONS,
        hash_name = "sha512"
    ).hash( True )
    
    check = {
        "value": "ok",
        "time": int( time.time() * 1000 ),
        "random": utils.rand.random_data( config.DEK_CHECK_SALT_LENGTH )
    }
    check = json.dumps( check )
    check = utils.aes.do_encrypt( check, dek )
    
    return dek_salt, dek, check

################################################################################
#    Decrypt/Encrypt user session

def do_decrypt_user_session( data ):
    return utils.aes.do_decrypt( data, session[ "password" ] )

def do_encrypt_user_session( data ):
    return utils.aes.do_encrypt( data, session[ "password" ] )

################################################################################
#    Redis Cache

def redis_cache( ttl = 3600 ):
    def decorator( func ):
        @functools.wraps( func )
        def wrapper_cache( *args, **kwargs ):
            lst = []
            lst.append( func.__name__ )
            lst.extend( args )
            lst = map( str, lst )
            index = "_".join( lst )
            index = hashlib.sha256( index ).hexdigest()
            
            d = config.redis_shared.get( index )
            
            if d != None:
                config.redis_shared.expire( index, ttl )
                
                buff = StringIO()
                buff.write( base64.b64decode( d ) )
                buff.seek( 0 )
                
                return cPickle.load( buff )
            
            else:
                d = func( *args, **kwargs )
                
                buff = StringIO()
                cPickle.dump( d, buff )
                buff.seek( 0 )
                d_cached = base64.b64encode( buff.getvalue() )
                
                config.redis_shared.set( index, d_cached, ex = ttl )
                
                return d
    
        return wrapper_cache
    return decorator

################################################################################
#    No preview image

def no_preview_image():
    img = Image.new( "L", ( 210, 297 ), 255 )
    draw = ImageDraw.Draw( img )
    font = ImageFont.truetype( "arial.ttf", 18 )
    draw.text( ( 0, 0 ), "No preview", 0, font = font )
     
    buff = utils.images.pil2buffer( img, "PNG" )
    return buff

################################################################################
#    SMTP context manager

class mySMTP( object ):
    def __init__( self ):
        self.host = config.smtpserver
        self.username = config.smtpuser
        self.password = config.smtppassword
        self.sender = config.sender
    
    def __enter__( self ):
        self.s = smtplib.SMTP( self.host )
        self.s.starttls()
        self.s.login( self.username, self.password )
        return self.s
    
    def __exit__( self, ty, value, traceback ):
        self.s.quit()