#!/usr/bin/python
#
# by Erik Osheim
import os
from subprocess import Popen, PIPE

class Cipher(object):
    '''Cipher represents a particular hashing strategy (password, seed, and type). Cipher can encrypt or decrypt data.'''
    def __init__(self, password, seed='aes.py', hashtype='rmd160'):
        self.password = password
        self.seed     = seed
        self.hashtype = hashtype
    def encrypt(self, data):
        return encrypt_data(data, self.password, self.seed, self.hashtype)
    def decrypt(self, encrypted):
        return decrypt_data(encrypted, self.password, self.seed, self.hashtype)

def _check_aespipe():
    '''This function checks if we have the "aespipe" binary in $PATH'''
    result = os.system('which aespipe > /dev/null')
    if result != 0:
        raise Exception, "Could not find aespipe; is it installed?"

def encrypt_data(data, password, seed='aes.py', hashtype='rmd160'):
    '''Uses password to encrypt data'''
    _check_aespipe()
    args = ["aespipe", "-S", seed, "-H", hashtype, "-p", "0"]
    p = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    (stdout, stdin, stderr) = (p.stdout, p.stdin, p.stderr)
    stdin.write(password + '\n')
    stdin.write(data)
    stdin.close()
    encrypted = stdout.read()
    err = stderr.read()
    if err:
        raise Exception, "Problem: %s" % err
    return encrypted

def encrypt_path(path, data, password, seed='aes.py', hashtype='rmd160'):
    '''Uses password to encrypt data and writes result to path'''
    encrypted = encrypt_data(data, password, seed, hashtype)
    f = open(path, 'w')
    f.write(encrypted)
    f.close()

def decrypt_data(encrypted, password, seed='aes.py', hashtype='rmd160'):
    '''Uses password to decrypt data'''
    _check_aespipe()
    args = ["aespipe", "-d", "-S", seed, "-H", hashtype, "-p", "0"]
    p = Popen(args, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    (stdout, stdin, stderr) = (p.stdout, p.stdin, p.stderr)
    stdin.write(password + '\n')
    stdin.write(encrypted)
    stdin.close()
    data = stdout.read()
    err = stderr.read()
    if err:
        raise Exception, "Problem: %s" % err

    # data is null-padded at the end to align on 16 or 512 bytes boundaries
    i = len(data)
    while i > 1:
        if data[i-1] == '\x00':
            i -= 1
        else:
            break
    return data[:i]

def decrypt_path(path, password, seed='aes.py', hashtype='rmd160'):
    '''Uses password to decrypt data from path'''
    f = open(path, 'r')
    encrypted = f.read()
    f.close()
    data = decrypt_data(encrypted, password, seed, hashtype)
    return data

if __name__ == "__main__":
    import optparse, sys

    parser = optparse.OptionParser()
    parser.set_defaults(mode='decrypt')
    parser.set_defaults(password='insecure1@3$5^')
    parser.set_defaults(filename='output.aes')

    parser.add_option('-e', dest='mode', action='store_const', const='encrypt',
                      help="perform encryption on data from stdin")
    parser.add_option('-d', dest='mode', action='store_const', const='decrypt',
                      help="perform decryption on data from stdin")
    parser.add_option('-f', dest='filename', action='store', metavar='FILENAME',
                      help="encrypt to/from FILENAME (default: output.aes)")
    parser.add_option('-p', dest='password', action='store', metavar='PASSWORD',
                      default="insecure1@3$5^",
                      help="use password PASSWORD (default: insecure1@3$5^)")

    (opts, args) = parser.parse_args()

    c = Cipher(opts.password)

    if opts.mode == 'encrypt':
        data = sys.stdin.read()
        encrypted = c.encrypt(data)
        f = open(opts.filename, 'w')
        f.write(encrypted)
        f.close()
        print "data written to %r." % opts.filename
    else:
        print "reading data from %r:" % opts.filename
        f = open(opts.filename, 'r')
        encrypted = f.read()
        f.close()
        data = c.decrypt(encrypted)
        print data