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

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

from flask.globals import session

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( 100 ) )
    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( 10 )
    }
    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