encrypt with AES in python using pycrypto lib

tool.py

import argparse, os, sys

from Crypto.Cipher import AES
from Crypto.Hash import HMAC
from Crypto.Protocol.KDF import PBKDF2

# check arguments
parser = argparse.ArgumentParser()
parser.add_argument("file", help="the file we want to encrypt/decrypt")
parser.add_argument("key", help="your key")
parser.add_argument("-e", "--encrypt", action="store_true")
parser.add_argument("-d", "--decrypt", action="store_true")
args = parser.parse_args()

# input file
try:
    inputfile = open(args.file, "rb")
except IOError:
    sys.exit("Could not open the input file")

# output file
try:
    output = open("a.out", "wb")
except IOError:
    sys.exit("Could not create the output file")

# make 256bits keys for encryption and mac
salt = "this is a salt"
kdf = PBKDF2(args.key, salt, 64, 1000)
key = kdf[:32]
key_mac = kdf[32:]

# create HMAC
mac = HMAC.new(key_mac) # default is MD5

# encryption
if args.encrypt:
    iv = os.urandom(16)
    cipher = AES.new(key, AES.MODE_CFB, iv)    

    encrypted = cipher.encrypt(inputfile.read())
    mac.update(iv + encrypted)

    # output
    output.write(mac.hexdigest())
    output.write(iv)
    output.write(encrypted)

# decryption
else:
    data = inputfile.read()
    # check for MAC first
    verify = data[0:32]
    mac.update(data[32:])

    if mac.hexdigest() != verify:
        sys.exit("message was modified, aborting decryption")

    # decrypt
    iv = data[32:48]
    cipher = AES.new(key, AES.MODE_CFB, iv)    

    decrypted = cipher.decrypt(data[48:])
    
    #output
    output.write(decrypted)

Hi, I am using your code and conver it into class so it's easy to use it as module.

import os, sys
from Crypto.Cipher import AES
from Crypto.Hash import HMAC
from Crypto.Protocol.KDF import PBKDF2

class aesPBKDF2(object):
	
	def __init__(self, key, salt="this is a salt", dklen=64,count=1000):
	    self.key = key
	    self.dklen = dklen
	    self.count = count
	    # make 256bits keys for encryption and mac
	    kdf = PBKDF2(self.key, salt, self.dklen, self.count)
	    self.key = kdf[:32]
	    key_mac = kdf[32:]
	    # create HMAC
	    self.mac = HMAC.new(key_mac) # default is MD5


	def encrypt_file(self, in_filename, out_filename):

	    try:
	        infile = open(in_filename)
	    except IOError:
	        sys.exit("Could not open the input file")
	    try:
	        outfile = open(out_filename, 'wb')
	    except IOError:
	        sys.exit("Could not create the output file")
	    
	    iv = os.urandom(16)
	    cipher = AES.new(self.key, AES.MODE_CFB, iv)    
	    
	    encrypted = cipher.encrypt(infile.read())
	    self.mac.update(iv + encrypted)

	    # outfile
	    outfile.write(self.mac.hexdigest())
	    outfile.write(iv)
	    outfile.write(encrypted)
	    
	    infile.close()
	    outfile.close()


	def decrypt_file(self, in_filename, out_filename):
	    
	    try:
	        infile = open(in_filename,'rb')
	    except IOError:
	        sys.exit("Could not open the input file")
	    try:
	        outfile = open(out_filename,'w')
	    except IOError:
	        sys.exit("Could not create the output file")    
	    
	    data = infile.read()
	    # check for MAC first
	    verify = data[0:32]
	    self.mac.update(data[32:])

	    if self.mac.hexdigest() != verify:
	        outfile.close()
	        os.unlink(out_filename)
	        sys.exit("Wrong password !")

	    # decrypt
	    iv = data[32:48]
	    cipher = AES.new(self.key, AES.MODE_CFB, iv)    

	    decrypted = cipher.decrypt(data[48:])
	    
	    #outfile
	    outfile.write(decrypted)

	    infile.close()
	    outfile.close()