Skip to content

Instantly share code, notes, and snippets.

@mimoo
Last active September 7, 2019 11:11
Show Gist options
  • Save mimoo/11383475 to your computer and use it in GitHub Desktop.
Save mimoo/11383475 to your computer and use it in GitHub Desktop.
encrypt with AES in python using pycrypto lib
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)
@awarmanf
Copy link

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()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment