import sys
import struct
from Crypto import Random
from Crypto.Hash import SHA256
from Crypto.Cipher import AES

random_generator_file = Random.new()

# format
# -----header------
# iv     <- 16bytes
# chunksize   <- <Q
# -----------------
# ------chunk------
# size        <- <Q
# data <- chunksize
# -----------------

def encrypt(key, infile, outfile, chunksize=1024):
    iv = random_generator_file.read(16)
    encryptor = AES.new(key, AES.MODE_CBC, iv)
    outfile.write(iv)
    outfile.write(struct.pack('<Q', chunksize))

    while True:
        chunk = infile.read(chunksize)
        outfile.write(struct.pack('<Q', len(chunk)))
        if len(chunk) == 0:
            break
        elif len(chunk) % chunksize != 0:
            chunk += 'M' * (chunksize - len(chunk) % chunksize)
        outfile.write(encryptor.encrypt(chunk))

def decrypt(key, infile, outfile):
    iv = infile.read(16)
    decryptor = AES.new(key, AES.MODE_CBC, iv)
    chunksize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]

    while True:
        size = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]
        if size == 0:
            break
        chunk = infile.read(chunksize)
        outfile.write(decryptor.decrypt(chunk)[:size])

mode  = sys.argv[2]
hashkey = SHA256.new()
hashkey.update(sys.argv[1])
key = hashkey.digest()

if mode == "encrypt":
    encrypt(key, sys.stdin, sys.stdout)
elif mode == "decrypt":
    decrypt(key, sys.stdin, sys.stdout)