import sys
import getpass
import os.path
from Crypto.Cipher import AES
from Crypto.Hash import SHA512
from subprocess import call

#pycrypto

def getFileNameFromInput (argv, maxSize):
	if len(argv) < 3:
		print "No input file specified."
		return None

	filename = argv[2]

	if not os.path.isfile(filename):
		print "This is not a file."
		return None

	if os.stat(filename).st_size > maxSize*1024*1024:
		print "too big. max size is:", maxSize, "mb."
		return None

	return filename

def getPasswordFromStdIn(text):
	password = getpass.getpass(text)

	#if len(password) < 1:
	#	print "too short"
	#	return None
	return password

def isHelpMode(argv):
	for i in xrange(1, len(argv)):
		if argv[i] == "-h" or argv[i] == "--help":
			return True
	return False

def helpMode():
	print """Usage:
	Syntax:
		python enc OPERATION FILE
	---------------------------------------------------
		-h
			prints usage
		-e | --encrypt
			encrypts FILE
		-d | --decrypt
			decrypts FILE
		-r | --read
			shows the contents of FILE and re-encrypts
			with the same password after exit
		-w | --write
			opens the FILE for r/w and re-encrypts
			withe the same password after exit
	----------------------------------------------------
			 """

def enum(**enums):
    return type('Enum', (), enums)

def getMode(argv, mode):
	if isHelpMode(argv):
		return mode.HELP
	for i in xrange(1, len(argv)):
		if argv[i] == "-e" or argv[i] == "--encrypt":
			return mode.ENCRYPT
		if argv[i] == "-d" or argv[i] == "--decrypt":
			return mode.DECRYPT
		if argv[i] == "-r" or argv[i] == "--read":
			return mode.READ
		if argv[i] == "-w" or argv[i] == "--write":
			return mode.WRITE

	print "no valid mode. use -h to print usage section"
	return None

def pad(s, blocksize, padding):
	return s + (blocksize - len(s) % blocksize) * padding

def unpad(s, padding):
	while s[len(s)-1] == padding:
		s = s[:len(s)-2]
	return s

def createNewFileName(cryptMode, mode, filename, encryptExt, decryptExt, readWriteExt):
	if mode == cryptMode.ENCRYPT and filename[-3:] == decryptExt:
		return filename[:-3] + encryptExt
	if mode == cryptMode.ENCRYPT:
		return filename + encryptExt
	if mode == cryptMode.DECRYPT and filename[-3:] == encryptExt:
		return filename[:-3] + decryptExt
	if mode == cryptMode.READ or mode == cryptMode.WRITE:
		return filename + readWriteExt
	return None

def calcHash(pw):
	h = SHA512.new()
	h.update(pw)
	return h.hexdigest()

def encryptFile(fileContent, password, newFileName, blocksize, padding):
	cipher = pad(fileContent, blocksize, padding)
	for i in xrange(4):
		encrypter = AES.new(password[i*32:i*32+32], AES.MODE_CBC, password[i*32:i*32+16])
		cipher = encrypter.encrypt(cipher)

	open(newFileName, 'w').close()
	newFile = open(newFileName, 'w')
	newFile.write(cipher)
	newFile.close()

def decryptFile(plaintext, password, newFileName, padding):
	for i in xrange(3, -1, -1):
		decrypter = AES.new(password[i*32:i*32+32], AES.MODE_CBC, password[i*32:i*32+16])
		plaintext = decrypter.decrypt(plaintext)

	plaintext = unpad(plaintext, padding)
	open(newFileName, 'w').close()
	newFile = open(newFileName, 'w')
	newFile.write(plaintext)
	newFile.close()

# ======================================================================

cryptMode = enum(ENCRYPT=1, DECRYPT=2, READ=3, WRITE=4, HELP=5)

maxFileSize = 100
fileExtensionEncrypt = "_en"
fileExtensionDecrypt = "_de"
fileExtensionTmp	 = "_tmp_DELETE_THIS_FILE"
padding = '{'
blocksize = 32

filename = getFileNameFromInput(sys.argv, maxFileSize)
if filename == None:
	helpMode()
	exit(0)

fileContent = open(filename).read()

password = getPasswordFromStdIn("enter password: ")
if password == None:
	exit(0)

password = calcHash(password)

mode = getMode(sys.argv, cryptMode)
if mode == cryptMode.HELP or mode == None:
	helpMode()
	exit(0)

newFileName = createNewFileName(cryptMode, mode, filename, fileExtensionEncrypt, fileExtensionDecrypt, fileExtensionTmp)
if newFileName == None:
	print "not an encrypted file. use a file ending in '_en'"
	exit(0)

if mode == cryptMode.ENCRYPT:
	encryptFile(fileContent, password, newFileName, blocksize, padding)
	exit(0)

if mode == cryptMode.DECRYPT:
	decryptFile(fileContent, password, newFileName, padding)
	exit(0)

if mode == cryptMode.READ:
	decryptFile(fileContent, password, newFileName, padding)
	call(["less", newFileName])
	call(["rm", "-rf", newFileName])
	exit(0)

if mode == cryptMode.WRITE:
	decryptFile(fileContent, password, newFileName, padding)
	call(["nano", newFileName])
	newEncryptFileName = createNewFileName(cryptMode, cryptMode.ENCRYPT, filename, fileExtensionEncrypt, fileExtensionDecrypt, fileExtensionTmp)
	newFileContent = open(newFileName).read()
	call(["rm", "-rf", newFileName])
	encryptFile(newFileContent, password, newEncryptFileName, blocksize, padding)
	exit(0)