#!/usr/bin/env python
#
# Run it like this:
#  $ aescrypt <filename> => this will encrypt <filename> to <filename>.secret
#  $ aescrypt <filename.secret> => this will decrypt to <filename>
#
 
import getpass
import hashlib
import math
import os
import random
import shutil
import sys
from Crypto.Cipher import AES
 
 
HASH_FUNCTION = 'sha256'
 
 
class AESHandler(object):
    
    def __init__(self, password):
        self.__pass_hash = hashlib.new(HASH_FUNCTION, password)
        self.__aes = AES.new(self.pass_hash)
    
    def encrypt(self, data):
        length = len(data)
        padding_length = length % AES.block_size
        if padding_length:
            padding_length = AES.block_size - padding_length
        padded_data = data + padding(padding_length)
        encrypted = self.aes.encrypt(padded_data)
        return '%d:%s' % (length, encrypted)
    
    def decrypt(self, data):
        length, encrypted = data.split(':', 1)
        length = int(length)
        padded_data = self.aes.decrypt(encrypted)
        return padded_data[:length]
    
    def __get_pass_hash(self):
        return self.__pass_hash.digest()
    
    def __get_aes(self):
        return self.__aes
    
    pass_hash = property(__get_pass_hash)
    aes = property(__get_aes)
 
 
def encrypt(filename, handler):
    secret_fname = filename + '.secret'
    orig_file = open(filename)
    secret_file = open(secret_fname, 'w')
    try:
        data = orig_file.read()
        secret_file.write(handler.encrypt(data))
    except:
        try:
            os.remove(secret_fname)
        except:
            pass
        print 'Encryption failed:'
        raise
    else:
        os.remove(filename)
        print 'Encrypted %r to %r' % (filename, secret_fname)
    finally:
        orig_file.close()
        secret_file.close()
 
 
def decrypt(filename, handler):
    if filename.endswith('.secret'):
        new_fname = filename[:-7]
    else:
        new_fname = filename + '.unsecret'
    orig_file = open(filename)
    unsecret_file = open(new_fname, 'w')
    try:
        data = handler.decrypt(orig_file.read())
        unsecret_file.write(data)
    except:
        try:
            os.remove(new_fname)
        except:
            pass
        print 'Decryption failed:'
        raise
    else:
        os.remove(filename)
        print 'Decrypted %r to %r' % (filename, new_fname)
    finally:
        orig_file.close()
        unsecret_file.close()
 
 
def padding(length):
    return ''.join(chr(random.randint(0, 255)) for i in xrange(length))
 
 
def main():
    handlers = {}
    if sys.argv[1:]:
        for filename in sys.argv[1:]:
            if filename.endswith('.secret'):
                password = getpass.getpass(
                    prompt='Decryption password for "%s": ' % (filename,))
                if password not in handlers:
                    handlers[password] = AESHandler(password)
                decrypt(filename, handlers[password])
            else:
                password = getpass.getpass(
                    prompt='Encryption password for "%s": ' % (filename,))
                password2 = getpass.getpass(
                    prompt='Confirm encryption password: ')
                while password != password2:
                    print 'The two passwords do not match.'
                    password = getpass.getpass(
                        prompt='Encryption password for "%s": ' % (filename,))
                    password2 = getpass.getpass(
                        prompt='Confirm encryption password: ')
                if password not in handlers:
                    handlers[password] = AESHandler(password)
                encrypt(filename, handlers[password])
 
 
if __name__ == '__main__':
    main()