element-web/test/utils/generate-megolm-test-vector...

118 lines
3.6 KiB
Python
Executable File

#!/usr/bin/env python
from __future__ import print_function
import base64
import json
import struct
from cryptography.hazmat import backends
from cryptography.hazmat.primitives import ciphers, hashes, hmac
from cryptography.hazmat.primitives.kdf import pbkdf2
from cryptography.hazmat.primitives.ciphers import algorithms, modes
backend = backends.default_backend()
def parse_u128(s):
a, b = struct.unpack(">QQ", s)
return (a << 64) | b
def encrypt_ctr(key, iv, plaintext, counter_bits=64):
alg = algorithms.AES(key)
# Some AES-CTR implementations treat some parts of the IV as a nonce (which
# remains constant throughought encryption), and some as a counter (which
# increments every block, ie 16 bytes, and wraps after a while). Different
# implmententations use different amounts of the IV for each part.
#
# The python cryptography library uses the whole IV as a counter; to make
# it match other implementations with a given counter size, we manually
# implement wrapping the counter.
# number of AES blocks between each counter wrap
limit = 1 << counter_bits
# parse IV as a 128-bit int
parsed_iv = parse_u128(iv)
# split IV into counter and nonce
counter = parsed_iv & (limit - 1)
nonce = parsed_iv & ~(limit - 1)
# encrypt up to the first counter wraparound
size = 16 * (limit - counter)
encryptor = ciphers.Cipher(
alg,
modes.CTR(iv),
backend=backend
).encryptor()
input = plaintext[:size]
result = encryptor.update(input) + encryptor.finalize()
offset = size
# do remaining data starting with a counter of zero
iv = struct.pack(">QQ", nonce >> 64, nonce & ((1 << 64) - 1))
size = 16 * limit
while offset < len(plaintext):
encryptor = ciphers.Cipher(
alg,
modes.CTR(iv),
backend=backend
).encryptor()
input = plaintext[offset:offset+size]
result += encryptor.update(input) + encryptor.finalize()
offset += size
return result
def hmac_sha256(key, message):
h = hmac.HMAC(key, hashes.SHA256(), backend=backend)
h.update(message)
return h.finalize()
def encrypt(key, iv, salt, plaintext, iterations=1000):
"""
Returns:
(bytes) ciphertext
"""
if len(salt) != 16:
raise Exception("Expected 128 bits of salt - got %i bits" % len((salt) * 8))
if len(iv) != 16:
raise Exception("Expected 128 bits of IV - got %i bits" % (len(iv) * 8))
sha = hashes.SHA512()
kdf = pbkdf2.PBKDF2HMAC(sha, 64, salt, iterations, backend)
k = kdf.derive(key)
aes_key = k[0:32]
sha_key = k[32:]
packed_file = (
b"\x01" # version
+ salt
+ iv
+ struct.pack(">L", iterations)
+ encrypt_ctr(aes_key, iv, plaintext)
)
packed_file += hmac_sha256(sha_key, packed_file)
return (
b"-----BEGIN MEGOLM SESSION DATA-----\n" +
base64.encodestring(packed_file) +
b"-----END MEGOLM SESSION DATA-----"
)
def gen(password, iv, salt, plaintext, iterations=1000):
ciphertext = encrypt(
password.encode('utf-8'), iv, salt, plaintext.encode('utf-8'), iterations
)
return (plaintext, password, ciphertext.decode('utf-8'))
print (json.dumps([
gen("password", b"\x88"*16, b"saltsaltsaltsalt", "plain", 10),
gen("betterpassword", b"\xFF"*8 + b"\x00"*8, b"moresaltmoresalt", "Hello, World"),
gen("SWORDFISH", b"\xFF"*8 + b"\x00"*8, b"yessaltygoodness", "alphanumerically" * 4),
gen("password"*32, b"\xFF"*16, b"\xFF"*16, "alphanumerically" * 4),
], indent=4))