#!/usr/bin/env python

from Crypto.Cipher import Blowfish
from Crypto.Hash import SHA
from Crypto import Random
from optparse import OptionParser
import sys
import struct
import os.path
import os

_MAGIC = bytearray([0, 2, 144, 4, 3, 4, 5, 6, 3, 12, 144, 134, 5, 6, 24, 2, 0, 3, 4, 3, 4, 3, 88, 4, 5])

def check_file(f):
    h = open(f, 'r')
    test = h.read(len(_MAGIC))
    h.close()
    return test == _MAGIC

def encrypt_file(inf, outf, password):
    infile = open(inf, 'r')
    outfile = open(outf, 'w')

    iv = Random.get_random_bytes(Blowfish.block_size)
    bf = Blowfish.new(SHA.new(password).hexdigest(), Blowfish.MODE_CBC, iv)

    filesize = os.path.getsize(inf)

    outfile.write(str(_MAGIC))
    outfile.write(struct.pack('<Q', filesize))
    outfile.write(iv)
    outfile.write(bf.encrypt(str(_MAGIC)))

    while True:
        chunk = infile.read(8192)
        if len(chunk) == 0:
            break
        elif len(chunk) % 16 != 0:
            chunk += ' ' * (16 - len(chunk) % 16)

        outfile.write(bf.encrypt(chunk))

    outfile.close()
    infile.close()

def decrypt_file(inf, outf, password):
    infile = open(inf, 'r')
    outfile = open(outf, 'w')

    test = infile.read(len(_MAGIC))
    if test != _MAGIC:
        raise Exception("This is not an encrypted file")

    origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
    iv = infile.read(Blowfish.block_size)
    bf = Blowfish.new(SHA.new(password).hexdigest(), Blowfish.MODE_CBC, iv)
    testdecrypt = bf.decrypt(infile.read(len(_MAGIC)))
    if testdecrypt != _MAGIC:
        raise Exception("Failure decrypting file.")

    while True:
        chunk = infile.read(8192)
        if len(chunk) == 0:
            break

        outfile.write(bf.decrypt(chunk))

    outfile.truncate(origsize)

    outfile.close()
    infile.close()

def main():
    usage = "Usage: %prog [options] file"
    parser = OptionParser(usage)
    parser.add_option("-i", "--input", dest="infile")
    parser.add_option("-o", "--output", dest="outfile")

    (options, args) = parser.parse_args()

    if not options.infile:
        if len(args) != 1:
            parser.error("Incorrect number of arguments")
        infile = args[0]
    else:
        if args:
            parser.error("Incorrect number of arguments")
        infile = options.infile

    outfile = options.outfile 
    
    if check_file(infile):
        print "File is to be decrypted."
        if not outfile: outfile = "%s.dec" % infile
        decrypt = True
    else:
        print "File is to be encrypted."
        if not outfile: outfile = "%s.enc" % infile
        decrypt = False

    os.system("stty -echo")
    print "Please enter your password: ",
    password = sys.stdin.readline().strip()
    os.system("stty echo")
    print ""
    if decrypt:
        decrypt_file(infile, outfile, password)
    else:
        encrypt_file(infile, outfile, password)

    print "Done.\n"

if __name__ == '__main__':
    main()