126 lines
3.4 KiB
Python
126 lines
3.4 KiB
Python
#!/usr/bin/env python
|
|
import os
|
|
from base64 import b64encode
|
|
from base64 import b64decode
|
|
|
|
from Cryptodome import Random
|
|
from Cryptodome.Cipher import AES
|
|
|
|
__author__ = 'rolandh'
|
|
|
|
POSTFIX_MODE = {
|
|
"cbc": AES.MODE_CBC,
|
|
"cfb": AES.MODE_CFB,
|
|
"ecb": AES.MODE_CFB,
|
|
}
|
|
|
|
BLOCK_SIZE = 16
|
|
|
|
|
|
class AESCipher(object):
|
|
def __init__(self, key, iv=""):
|
|
"""
|
|
|
|
:param key: The encryption key
|
|
:param iv: Init vector
|
|
:return: AESCipher instance
|
|
"""
|
|
self.key = key
|
|
self.iv = iv
|
|
|
|
def build_cipher(self, iv="", alg="aes_128_cbc"):
|
|
"""
|
|
:param iv: init vector
|
|
:param alg: cipher algorithm
|
|
:return: A Cipher instance
|
|
"""
|
|
typ, bits, cmode = alg.split("_")
|
|
|
|
if not iv:
|
|
if self.iv:
|
|
iv = self.iv
|
|
else:
|
|
iv = Random.new().read(AES.block_size)
|
|
else:
|
|
assert len(iv) == AES.block_size
|
|
|
|
if bits not in ["128", "192", "256"]:
|
|
raise Exception("Unsupported key length")
|
|
try:
|
|
assert len(self.key) == int(bits) >> 3
|
|
except AssertionError:
|
|
raise Exception("Wrong Key length")
|
|
|
|
try:
|
|
return AES.new(self.key, POSTFIX_MODE[cmode], iv), iv
|
|
except KeyError:
|
|
raise Exception("Unsupported chaining mode")
|
|
|
|
|
|
def encrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7",
|
|
b64enc=True, block_size=BLOCK_SIZE):
|
|
"""
|
|
:param key: The encryption key
|
|
:param iv: init vector
|
|
:param msg: Message to be encrypted
|
|
:param padding: Which padding that should be used
|
|
:param b64enc: Whether the result should be base64encoded
|
|
:param block_size: If PKCS#7 padding which block size to use
|
|
:return: The encrypted message
|
|
"""
|
|
|
|
if padding == "PKCS#7":
|
|
_block_size = block_size
|
|
elif padding == "PKCS#5":
|
|
_block_size = 8
|
|
else:
|
|
_block_size = 0
|
|
|
|
if _block_size:
|
|
plen = _block_size - (len(msg) % _block_size)
|
|
c = chr(plen)
|
|
msg += c*plen
|
|
|
|
cipher, iv = self.build_cipher(iv, alg)
|
|
cmsg = iv + cipher.encrypt(msg)
|
|
if b64enc:
|
|
return b64encode(cmsg)
|
|
else:
|
|
return cmsg
|
|
|
|
|
|
def decrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7", b64dec=True):
|
|
"""
|
|
:param key: The encryption key
|
|
:param iv: init vector
|
|
:param msg: Base64 encoded message to be decrypted
|
|
:return: The decrypted message
|
|
"""
|
|
if b64dec:
|
|
data = b64decode(msg)
|
|
else:
|
|
data = msg
|
|
|
|
_iv = data[:AES.block_size]
|
|
if iv:
|
|
assert iv == _iv
|
|
cipher, iv = self.build_cipher(iv, alg=alg)
|
|
res = cipher.decrypt(data)[AES.block_size:]
|
|
if padding in ["PKCS#5", "PKCS#7"]:
|
|
res = res[:-ord(res[-1])]
|
|
return res
|
|
|
|
if __name__ == "__main__":
|
|
key_ = "1234523451234545" # 16 byte key
|
|
# Iff padded, the message doesn't have to be multiple of 16 in length
|
|
msg_ = "ToBeOrNotTobe W.S."
|
|
aes = AESCipher(key_)
|
|
iv_ = os.urandom(16)
|
|
encrypted_msg = aes.encrypt(key_, msg_, iv_)
|
|
txt = aes.decrypt(key_, encrypted_msg, iv_)
|
|
assert txt == msg_
|
|
|
|
encrypted_msg = aes.encrypt(key_, msg_, 0)
|
|
txt = aes.decrypt(key_, encrypted_msg, 0)
|
|
assert txt == msg_
|