#!/usr/bin/env python

# This script will encrypt/decrypt the AES ciphertext file.

import argparse
from Crypto import Random
from Crypto.Cipher import AES
import base64
import os.path
import sys

parser = argparse.ArgumentParser(description='Encrypt/decrypt using ciphertext/plaintext and an encryption key with AES256.')
parser.add_argument("-m", "--mode", help="Operation mode.", choices=["encrypt","decrypt"])
parser.add_argument("-i", "--input", help="Specify the input ciphertext/plaintext filepath.")
parser.add_argument("-k", "--key", help="Specify the encryption key filepath.")
parser.add_argument("-o", "--output", help="Specify the output ciphertext/plaintext filepath.")
parser.add_argument("-v", "--verbose", action="store_true", dest="verbose", help="Output the ciphertext/plaintext to stdout.")
args = parser.parse_args()


if args.input and os.path.isfile(args.input):
  with open(args.input, 'r') as input_file:
    input_rawtext = input_file.read()
else:
  print "--input option required and must be valid file."
  sys.exit(1)

if args.key and os.path.isfile(args.key):
  with open(args.key, 'r') as key_file:
    key_rawtext = key_file.read()
else:
  print "--key option required and must be valid file."
  sys.exit(1)

if args.mode == 'encrypt':
  key_text = base64.b64decode(key_rawtext)
  plain_text = input_rawtext
  cipher = AES.new(key_text, AES.MODE_ECB)
  # Insert padding if necessary to ciphertext (must be a multiple of 16 bytes)
  length = 16 - (len(plain_text) % 16)
  plain_text += chr(length)*length
  cipher_text = base64.b64encode(cipher.encrypt(plain_text))
  if args.output:
    with open(args.output, 'w') as output_file:
      output_file.truncate
      output_file.write(cipher_text)
  if args.verbose:
    print cipher_text

if args.mode == 'decrypt':
  key_text = base64.b64decode(key_rawtext)
  cipher_text = base64.b64decode(input_rawtext)
  cipher = AES.new(key_text, AES.MODE_ECB)
  plain_text = cipher.decrypt(cipher_text)
  if args.output:
    with open(args.output, 'w') as output_file:
      output_file.truncate
      output_file.write(plain_text)
  if args.verbose:
    print plain_text