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

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

def sql_insert_generate( table, fields, returning = None ):
    if isinstance( fields, ( str ) ):
        fields = [ fields ]
    
    if not isinstance( fields, ( list, tuple, ) ):
        raise Exception( "list or tuple needed as fields" )
    
    f = ",".join( fields )
    place_holder = ",".join( [ "%s" ] * len( fields ) )
    sql = "INSERT INTO {} ({}) VALUES ({})".format( table, f, place_holder )
    
    if returning != None:
        sql += " RETURNING {}".format( returning )
        
    return sql

################################################################################
#     AES
################################################################################

encryption_prefix = "icnml$"

def do_decrypt( data, password ):
    try:
        data = AESCipher( password ).decrypt( data )
        
        if data.startswith( encryption_prefix ):
            return data[ len( encryption_prefix ): ]
        else:
            return "-"
    
    except:
        return "-"

def do_encrypt( data, password ):
    return AESCipher( password ).encrypt( encryption_prefix + data )

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: ] ) ]

################################################################################
#    PBKDF2
################################################################################

class pbkdf2( object ):
    def __init__( self, word, salt = None, iterations = 20000, hash_name = "sha512" ):
        if salt != None and salt.startswith( "pbkdf2$" ):
            self.word = word
            self.stored_hash = salt
            _, self.hash_name, self.salt, self.iterations, self.h = salt.split( "$" )
            self.iterations = int( self.iterations )
        
        else:
            self.word = word
            self.salt = salt or random_data( 100 )
            self.iterations = int( iterations )
            self.hash_name = hash_name
            self.stored_hash = None
            self.h = None
    
    def hash( self, hash_only = False ):
        h = hashlib.pbkdf2_hmac( self.hash_name, self.word, self.salt, self.iterations )
        h = binascii.hexlify( h )
        
        if hash_only:
            return h
        else:
            return "$".join( map( str, [ "pbkdf2", self.hash_name, self.salt, self.iterations, h ] ) )
    
    def verify( self ):
        return self.stored_hash == self.hash()