HackingAES.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import logging
import random
import os
from Crypto.Cipher import AES

def _make_binary(x):
    for c in list(x):
        assert c in '0123456789ABCDEFabcdef'
    return ''.join([chr(int(x[2*i:2*i+2], 16)) for i in range(len(x) / 2)])


def pkcs7_pad(s):
    n = len(s) % AES.block_size
    assert 0 <= n < AES.block_size
    n = AES.block_size - n
    assert 0 < n <= AES.block_size
    return s + n * chr(n)


def pkcs7_unpad(s):
    n = ord(s[-1])
    x, y = s[:-n], s[-n:]
    assert y == n * chr(n), repr((y, n * chr(n)))
    return x


class AesCrypt256(object):
    def __init__(self, key, iv):
        assert len(key) == 64
        assert len(iv) == 32
        self.key = _make_binary(key)
        self.iv = _make_binary(iv)

    def encrypt_string(self, R):
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        return cipher.encrypt(pkcs7_pad(R))

    def decrypt_string(self, R):
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        return pkcs7_unpad(cipher.decrypt(R))

    def encrypt_file(self, path):
        logging.debug("Encrypt " + path)
        R = open(path, 'rb').read()
        X = self.encrypt_string(R)
        open(path + '.aes', 'wb').write(X)

    def decrypt_file(self, path):
        logging.debug("Decrypt " + path)
        assert path.endswith('.aes')
        R = open(path, 'rb').read()
        X = self.decrypt_string(R)
        open(path[:-4], 'wb').write(X)

    def encrypt_tree(self, tree):
        for root, _, filenames in os.walk(tree):
            for f in filenames:
                self.encrypt_file(os.path.join(root, f))

    def decrypt_tree(self, tree):
        for root, _, filenames in os.walk(tree):
            for f in filenames:
                self.decrypt_file(os.path.join(root, f))


if __name__ == "__main__":
    key = ''.join([random.choice("0123456789ABCDEF") for i in range(64)])
    assert len(key) == 2 * len(_make_binary(key))
    iv = ''.join([random.choice("0123456789ABCDEF") for i in range(32)])
    assert len(iv) == 2 * len(_make_binary(iv))
    aes = AesCrypt256(key, iv)

    R = ''.join([chr(int(256 * random.random())) for i in range(50000 * AES.block_size - 3)])
    assert pkcs7_unpad(pkcs7_pad(R)) == R

    x = aes.encrypt_string(R)
    y = aes.decrypt_string(x)
    assert R == y