Electrum 2.8 password checking with a partial wallet file

electrum28-test.py

#!/usr/bin/env python

from Crypto.Cipher.AES import new as new_aes, MODE_CBC
import zlib, re, hashlib, base64, random


# The decypted and decompressed wallet should start with one of these two:
EXPECTED_BYTES_1 = b'{\n    "'
EXPECTED_BYTES_2 = b'{\r\n    "'
EXPECTED_BYTES_LEN = max(len(EXPECTED_BYTES_1), len(EXPECTED_BYTES_2))

# After the initial string above, it should continue with:
# one or more printable non-doublequotes, then doubleqoute, colon, space
EXPECTED_RE = re.compile(b'[\x23-\x7E!]+": ')

# Returns True if the key correctly starts decrypting the ciphertext (which may be truncated)
def check_encrypted_zlib(ciphertext, key, iv):
    assert ciphertext and len(ciphertext) % 16 == 0
    assert len(key) == 16
    assert len(iv)  == 16

    aes = new_aes(key, MODE_CBC, iv)
    plaintext = aes.decrypt(ciphertext[:16])  # decrypt the first block

    # This is an unnecessary check, but zlib is slow, and this speeds things up; YMMV
    if not (plaintext.startswith(b'\x78\x9c') and ord(plaintext[2]) & 0x7 == 0x5):
        return False

    decompressor = zlib.decompressobj(15)  # calls deflateInit2 with windowBits==15, same as Electrum
    uncompressed = b''  # the entire uncompressed string so far
    pos = found_valid = 0
    while pos < len(ciphertext):
        try:
            uncompressed += decompressor.decompress(plaintext)  # calls deflate, raises on error (wrong password)

            # Look for the expected initial bytes
            if uncompressed.startswith(EXPECTED_BYTES_1):
                found_valid = len(EXPECTED_BYTES_1)
            elif uncompressed.startswith(EXPECTED_BYTES_2):
                found_valid = len(EXPECTED_BYTES_2)

            # If found, look for the byte pattern which follows
            # the static bytes above to minimize false positives
            if found_valid:
                plaintext = aes.decrypt(ciphertext[pos+16:])        # decrypt and then
                uncompressed += decompressor.decompress(plaintext)  # uncompress the rest
                return bool(EXPECTED_RE.match(uncompressed[found_valid:]))

        except zlib.error as e:
            return False  # wrong passwords usually end up here

        if len(uncompressed) >= EXPECTED_BYTES_LEN:
            return False  # didn't find the expected initial bytes

        # Decrypt the next block
        pos += 16
        plaintext = aes.decrypt(ciphertext[pos:pos+16])

    assert False, 'at least some text decompressed from 512 input bytes'


# Read in up to 512 bytes of ciphertext -- with 3ish bytes of zlib & deflate header and up to
# 289ish bytes of Huffman codes, this should be plenty to find the beginning of the JSON data.
# https://stackoverflow.com/questions/25431160
max_data_len = 512
max_data_len += 37  # add the 4-byte magic & 33-byte ephemeral compressed pubkey
max_data_len = (max_data_len + 2) // 3 * 4  # convert to its base64 length (rounding up)
with open(r'btcrecover\test\test-wallets\electrum28-wallet', 'rb') as wallet_file:
    base64_data = wallet_file.read(max_data_len)
ciphertext = base64.b64decode(base64_data)[37:]       # remove the 4-byte magic & 33-byte pubkey, then
ciphertext = ciphertext[:len(ciphertext) // 16 * 16]  # truncate to a multiple of 16 (AES block len)

# ECC code elided; these are the results w/ password 'btcr-test-password' and the test file
shared_pubkey = b'\x02-mdy<]\xe7\x89-1\xdc\x13)\xb2\xa3\x90!\xf9U\x7f\x18\xb4NC\x03g\xf4; \xf2\x90\x8f'

# Should print True
keys = hashlib.sha512(shared_pubkey).digest()
print check_encrypted_zlib(ciphertext, keys[16:32], keys[:16])  # ciphertext, key, iv

# Should print False
keys = 48 * b'\0'
print check_encrypted_zlib(ciphertext, keys[16:32], keys[:16])  # ciphertext, key, iv

# Not-very-good false positive checking
for i in xrange(1000000):
    keys = b''.join(chr(random.randrange(256)) for j in xrange(48))
    assert check_encrypted_zlib(ciphertext, keys[16:32], keys[:16]) == False