Python: Simple AES with zlib

easy_aes.py

#!/data/local/bin/python
# AES encryption/decryption with zlib compression
from Crypto.Cipher import AES
from Crypto import Random
import hashlib
import zlib
import base64


_block_size = 16
_key_size = 32
_mode = AES.MODE_CBC

# MODE_B64 : encryption result as base64 encoded string (default)
# MODE_BIN : raw encryption result
MODE_B64, MODE_BIN = 0, 1


def encrypt(s, pass_string, mode=MODE_B64):
    h = hashlib.sha256()
    h.update(pass_string)
    key_bytes = h.digest()

    s = zlib.compress(s)
    pad = _block_size - len(s) % _block_size
    data = s + bytes(pad*[pad])

    iv_bytes = Random.get_random_bytes(_block_size)
    encrypted_bytes = iv_bytes + AES.new(key_bytes, _mode, iv_bytes).encrypt(data)
    if mode == MODE_B64:
        encrypted_bytes = base64.urlsafe_b64encode(encrypted_bytes)
    return encrypted_bytes


def decrypt(s, pass_string, mode=MODE_B64):
    h = hashlib.sha256()
    h.update(pass_string)
    key_bytes = h.digest()

    if mode == MODE_B64:
        s = base64.urlsafe_b64decode(s)
    encrypted_bytes = s
    iv_bytes = encrypted_bytes[:_block_size]
    encrypted_bytes = encrypted_bytes[_block_size:]
    plain_text = AES.new(key_bytes, _mode, iv_bytes).decrypt(encrypted_bytes)
    pad = plain_text[-1]
    return zlib.decompress(plain_text[:-pad])


if __name__ == "__main__":
    org_str = b"Some test data for simple AES encryption 1234567890"
    enc_str = encrypt(org_str, b'password')
    print("ENC STR[%d:%s]" % (len(enc_str), enc_str))
    dec_str = decrypt(enc_str, b'password')
    assert(org_str == dec_str)
    print("DEC STR[%d:%s]" % (len(dec_str), dec_str))

Updated for Python 3. Arguments s and password of encrypt and decrypt functions have to be bytes type. cc: @meson10, @balasmeh

Thank you so much. I was struggling entire day to compress and encrypt in python 3 and it was not working. Magic was hidden in the way you are adding padding in string to make sure its multiple of 16

 s = zlib.compress(s)    
pad = _block_size - len(s) % _block_size    
data = s + bytes(pad*[pad])

Thanks again.