Compare data lenghts for different encryption methods and mnemonic formats

python_encrypt_methods_tests.py

from io import StringIO
import hashlib
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import base64
from qrcode import QRCode
from embit.wordlists.bip39 import WORDLIST
from embit import bip39

MNEMONIC_ID = "test_ID"
TEST_KEY = "test_KEY"
TEST_12W_MNEMONIC = "domain liberty need select pledge orient isolate fade drift fragile round axis"
TEST_24W_MNEMONIC = "yard genius riot ball obey churn twin decrease unaware glide grunt typical surge age uncover used spatial pioneer nice swap math hood twin inquiry"

class AESCipher(object):
    """Helper for AES encrypt/decrypt"""

    def __init__(self, key, salt):
        self.key = hashlib.pbkdf2_hmac(
            'sha256',
            key.encode(),
            salt.encode(),
            100000
        )

    def encrypt(self, raw, cbc=False):
        """Encrypt using AES MODE_ECB and return the value encoded as base64"""
        data_bytes = raw if isinstance(raw, bytes) else raw.encode()
        print("Data bytes:", len(data_bytes))
        padded_len = len(data_bytes + b"\x00" * ((16- (len(data_bytes) % 16)) % 16))
        print("Padded bytes:", padded_len)
        print("Blocks:", padded_len//16)
        if cbc:
            iv = get_random_bytes(AES.block_size)
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
            encrypted = cipher.encrypt(
                data_bytes + b"\x00" * ((16- (len(data_bytes) % 16)) % 16)
            )
            payload = iv+encrypted
            print("Encrypted:", base64.b64encode(payload).decode("utf-8"))
            print("Encrypted lenght:", len(payload))
        else:
            cipher = AES.new(self.key, AES.MODE_ECB)
            encrypted = cipher.encrypt(
                data_bytes + b"\x00" * ((16- (len(data_bytes) % 16)) % 16)
            )
            print("Encrypted:", base64.b64encode(encrypted).decode("utf-8"))
            print("Encrypted lenght:", len(encrypted))
        # Prints ascii QR code
        qr_code = QRCode()
        qr_code.add_data(encrypted)
        qr_string = StringIO()
        qr_code.print_ascii(out=qr_string, invert=True)
        print("QR Code:")
        print(qr_string.getvalue())
        # Decription test
        if cbc and len(payload) == 48:
            print("Testing pure bytes decription and mnemonic recriation")
            cipher = AES.new(self.key, AES.MODE_CBC, payload[:AES.block_size])
            decrypted = cipher.decrypt(payload[AES.block_size:])
            print(bip39.mnemonic_from_bytes(decrypted))
            # What if wrong key is used on encrypted bytes as mnemonic?
            print("Testing pure bytes decription and mnemonic recriation with wrong key, salt")
            wrong_key = hashlib.pbkdf2_hmac(
                'sha256',
                "wrong key".encode(),
                "wrong salt".encode(),
                100000
            )
            cipher = AES.new(wrong_key, AES.MODE_CBC, payload[:AES.block_size])
            decrypted = cipher.decrypt(payload[AES.block_size:])
            print(bip39.mnemonic_from_bytes(decrypted))
        return base64.b64encode(encrypted)
    


    def decrypt(self, enc):
        """Decrypt a base64 using AES MODE_ECB and return the value decoded as string"""
        encrypted = base64.b64decode(enc)
        decryptor = AES.new(self.key, AES.MODE_ECB)
        encrypted = decryptor.decrypt(encrypted).decode("utf-8")
        return encrypted.replace("\x00", "")
    

def report(mnemonic, cbc=False):
    print("Mnemonic:", mnemonic)
    print("Encrypting words:")
    encryptor = AESCipher(TEST_KEY, MNEMONIC_ID)
    _ = encryptor.encrypt(mnemonic, cbc).decode("utf-8")

    print("\nEncrypting numbers:")
    numbers = ""
    for word in mnemonic.split():
        word_list_index = WORDLIST.index(word) + 1
        numbers += format(word_list_index, '04d')
    encryptor = AESCipher(TEST_KEY, MNEMONIC_ID)
    _ = encryptor.encrypt(numbers, cbc).decode("utf-8")

    print("\nEncrypting hex numbers:")
    numbers = ""
    for word in mnemonic.split():
        word_list_index = WORDLIST.index(word) + 1
        numbers += f'{word_list_index:0>3X}'
    encryptor = AESCipher(TEST_KEY, MNEMONIC_ID)
    _ = encryptor.encrypt(numbers, cbc).decode("utf-8")
    print("\nEncrypting bytes:")
    words = mnemonic.split(" ")
    checksum_bits = 8 if len(words) == 24 else 4
    indexes = [WORDLIST.index(word) for word in words]
    bitstring = "".join([f"{bin(index)[2:]:0>11}" for index in indexes])[
        :-checksum_bits
    ]
    qr_data = int(bitstring, 2).to_bytes((len(bitstring) + 7) // 8, "big")
    encryptor = AESCipher(TEST_KEY, MNEMONIC_ID)
    _ = encryptor.encrypt(qr_data, cbc).decode("utf-8")

    print("\n")

print("12 words AES-ECB")
report(TEST_12W_MNEMONIC)
print("12 words AES-CBC")
report(TEST_12W_MNEMONIC,cbc=True)
print("24 words AES-ECB")
report(TEST_24W_MNEMONIC)
print("24 words AES-CBC")
report(TEST_24W_MNEMONIC,cbc=True)