functions.py

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

from cStringIO import StringIO
import base64
import binascii
import hashlib
import random
import re
import string

from Crypto import Random
from Crypto.Cipher import AES
from PIL import ExifTags
import jinja2

def pbkdf2( word, salt = None, iterations = 20000, hash_name = 'sha512' ):
    if salt == None:
        salt = random_data( 100 )
    
    if salt.startswith( "pbkdf2$" ):
        stored_hash = salt
        _, hash_name, salt, iterations, h = salt.split( "$" )
        iterations = int( iterations )
        tested_hash = pbkdf2( word, salt, iterations, hash_name )
        
        return tested_hash == stored_hash
    
    else:
        h = hashlib.pbkdf2_hmac( hash_name, word, salt, iterations )
        h = binascii.hexlify( h )
        r = "$".join( map( str, [ "pbkdf2", hash_name, salt, iterations, h ] ) )
        
        return r

def random_data( N ):
    return ''.join( random.choice( string.ascii_uppercase + string.digits ) for _ in range( N ) )

def urlsplit( url ):
    data = re.match( '((?P<protocol>[^:]+)://)((?P<user>[^:]+)?(:(?P<password>[^@]+))?)?@(?P<host>[^:/]+)(:(?P<port>\d+))?(/(?P<database>[^&]+))?', url )
    return dict( [ ( key, data.group( key ) ) for key in [ 'user', 'password', 'host', 'port', 'database' ] ] )

def render_jinja_html( template_loc, file_name, **context ):
    return jinja2.Environment( 
        loader = jinja2.FileSystemLoader( template_loc + '/' )
    ).get_template( file_name ).render( context )

def rotate_image_upon_exif( img ):
    try:
        for orientation in ExifTags.TAGS.keys():
            if ExifTags.TAGS[ orientation ] == 'Orientation':
                break
        
        exif = dict( img._getexif().items() )
    
        if exif[orientation] == 3:
            img = img.rotate( 180, expand = True )
        
        elif exif[orientation] == 6:
            img = img.rotate( 270, expand = True )
        
        elif exif[orientation] == 8:
            img = img.rotate( 90, expand = True )
    
    except ( AttributeError, KeyError, IndexError ):
        pass
    
    return img

class AESCipher( object ):
    def __init__( self, key ): 
        self.key = hashlib.sha256( key.encode() ).digest()
    
    def encrypt( self, data ):
        data = self._pad( data )
        iv = Random.new().read( AES.block_size )
        cipher = AES.new( self.key, AES.MODE_CBC, iv )
        
        iv = base64.b64encode( iv )
        encrypted = base64.b64encode( cipher.encrypt( data ) )
        
        return "$".join( map( str, [ "AES256", iv, encrypted ] ) )
    
    def decrypt( self, data ):
        _, iv, data = data.split( "$" )
        iv = base64.b64decode( iv )
        data = base64.b64decode( data )
        cipher = AES.new( self.key, AES.MODE_CBC, iv )
        return self._unpad( cipher.decrypt( data ) ).decode( 'utf-8' )
    
    def _pad( self, s ):
        diff = AES.block_size - len( s ) % AES.block_size
        return s + diff * chr( diff )
    
    @staticmethod
    def _unpad( s ):
        return s[ :-ord( s[ len( s ) - 1: ] ) ]

def float_or_null( v ):
    try:
        return float( v )
    except:
        return None

def pil2buffer( img, format ):
    buff = StringIO()
    img.save( buff, format = format )
    buff.seek( 0 )
    
    return buff