| #!/usr/bin/python |
| # Copyright 2016 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Module for generating AES test vectors.""" |
| |
| from binascii import b2a_hex as b2a |
| from Crypto.Cipher import AES |
| from itertools import izip_longest |
| import os |
| |
| modes = { |
| AES.MODE_CBC: 'CBC', |
| AES.MODE_CFB: 'CFB', |
| AES.MODE_OFB: 'OFB', |
| } |
| |
| template = \ |
| ''' |
| <crypto_test name="AES:{mode}{key_bits} {test_num}"> |
| <clear_text format="hex"> |
| {pt} |
| </clear_text> |
| <key> |
| {key} |
| </key> |
| <cipher_text> |
| {ct} |
| </cipher_text> |
| <iv> |
| {iv} |
| </iv> |
| </crypto_test> |
| ''' |
| |
| def h2be(v): |
| # Convert input big-endian byte-string to 4-byte segmented |
| # little-endian words. Pad-bytes (if necessary) are the empty string. |
| word = [iter(v)] * 4 |
| return ''.join([ |
| ''.join(b[::-1]) for b in izip_longest(*word, fillvalue='') |
| ]) |
| |
| |
| for mode in [AES.MODE_CBC, AES.MODE_CFB, AES.MODE_OFB]: |
| for key_bytes in [16, 24, 32]: |
| test_num = 0 |
| for pt_len in [5, 16, 21, 32]: |
| # CBC mode requires block sized inputs. |
| if mode == AES.MODE_CBC and pt_len % 16: |
| continue |
| test_num += 1 |
| |
| actual_pt_len = pt_len |
| if pt_len % 16: |
| pt_len = 16 * ((pt_len / 16) + 1) |
| |
| key = os.urandom(key_bytes) |
| iv = os.urandom(16) |
| pt = os.urandom(pt_len) |
| |
| obj = AES.new(key, mode=mode, IV=iv, segment_size=128) |
| ct = obj.encrypt(pt) |
| obj = AES.new(key, mode=mode, IV=iv, segment_size=128) |
| |
| assert obj.decrypt(ct)[:pt_len] == pt |
| |
| print template.format(mode=modes[mode], |
| key_bits=str(key_bytes * 8), |
| test_num=str(test_num), |
| pt=b2a(h2be(pt[:actual_pt_len])), |
| key=b2a(h2be(key)), |
| ct=b2a(h2be(ct[:actual_pt_len])), |
| iv=b2a(h2be(iv))), |
| |
| |