AES256 Encryption / Decryption using Python

gistfile1.py

# Copyright 2015 Bluecore, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from Crypto.Cipher import AES
from Crypto import Random
from StringIO import StringIO
import base64

def pad(s):
    """ Simple padding, since the block cipher expects clear text to be multiples of the blocksize.
    """
    return s + b"\0" * (AES.block_size - len(s) % AES.block_size)

def unpad(s):
    return s.rstrip("\0").lstrip("\0")

def encrypt_file(input_file, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, IV=iv)
    buf = StringIO()
    with open(input_file, "r") as fh:
        # Read a line as a multiple of the AES block size.
        for line in fh.xreadlines():
            # pad (to match block size) -> encrypt -> b64encode (to make it suitable for urls)
            encrypted = cipher.encrypt(pad(line))
            # Use the URL safe version of the b64 encoding
            b64_encoded = base64.urlsafe_b64encode(encrypted)
            buf.write(b64_encoded)
            buf.write("\n")
    return buf

def decrypt_file(input_file, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, IV=iv)
    buf = StringIO()
    with open(input_file, "r") as fh:
        for line in fh.xreadlines():
            # Reverse the process
            # b64decode -> decrypt -> unpad
            # Use the URL safe version of the b64 decoding
            decrypted = unpad(cipher.decrypt(base64.urlsafe_b64decode(line)))
            buf.write(decrypted)
    return buf

if __name__ == "__main__":
    # key = "BYOUwqYyMgWfEIjSHhmVHBoJgobVUJbR"
    # Instead of a static key, create a key at random
    key = Random.new().read(256//8)

    # iv = "I6V8HN5DMUPM4AES"
    # Instead of a static IV, create a random IV
    # The IV needs to be the block size
    iv = Random.new().read(AES.block_size)

    print "Key is %s" % base64.b64encode(key)
    print "IV is %s" % base64.b64encode(iv)

    print "\n"

    # Encrypt the file
    encrypted = encrypt_file("clear_text_email_list.txt", key, iv)
    print "Encrypted lines:"
    print encrypted.getvalue()

    # Write the encrypted lines into a file
    with open("encrypted.txt", "wb") as fh:
        fh.write(encrypted.getvalue())

    # Decrypt the encrypted file
    decrypted = decrypt_file("encrypted.txt", key, iv)
    print "\n"
    print "Decrypted lines are:"
    print decrypted.getvalue()

    # Compare the two
    with open("clear_text_email_list.txt", "r") as fh:
        clear_lines = [line.strip() for line in fh.readlines()]
    decrypted_lines = decrypted.getvalue().strip().split("\n")

    # ... sanity check assertions
    assert len(clear_lines) == len(decrypted_lines), "%d vs %d" % (len(clear_lines), len(decrypted_lines))


    for i in xrange(0, len(clear_lines)):
        assert str(clear_lines[i]) == str(decrypted_lines[i]), "%s vs %s" % (clear_lines[i], decrypted_lines[i])

    print "Encrypted then decrypted and files match!"