#!/usr/bin/env python

"""
Generate test cases for AES Key Wrap, RFC 5649, for all key lengths
rather than just the 192-bit test vectors given in the RFC.
"""

# Motivation: The test vectors in the RFC happen to be for the one
# size of AES key which our AES core doesn't support (sigh), so
# they're useless for testing our code.  An external implementation
# would be ideal, but failing that, generating test vectors with an
# implementation that supports all key lengths and which got the right
# answers for the 192-bit case seems like the best bet.  We do this
# using our Python test implementation, because it's easier, but, more
# importantly, because it shares no code (just an author) with the C
# code we want to test.

from aes_keywrap import (KEK, hex2bin,
                         block_wrap_key, block_unwrap_key,
                         array_wrap_key, array_unwrap_key)
from argparse import ArgumentParser, FileType
from textwrap import TextWrapper
from cryptlib_py import *
import sys, os, atexit

wrapper = TextWrapper(width = 78, initial_indent = " " * 2, subsequent_indent = " " * 2)

def dump_hex(name, value, comment):
  args.output.write("\nstatic const uint8_t %s[] = { /* %s, %d bytes */\n" % (name, comment, len(value)))
  args.output.write(wrapper.fill(", ".join("0x%02x" % ord(v) for v in value)))
  args.output.write("\n};\n")

parser = ArgumentParser(description = __doc__)
parser.add_argument("--plaintext", help = "specify plaintext to be encrypted")
parser.add_argument("--kek-128", help = "specify 128-bit KEK to use")
parser.add_argument("--kek-256", help = "specify 256-bit KEK to use")
parser.add_argument("output", nargs = "?", type = FileType("w"), default = sys.stdout, help = "output file")
args = parser.parse_args()

def generate_and_test(K):
  C = array_wrap_key(Q, K)
  if block_wrap_key(Q, K) != C:
    raise RuntimeError
  if array_unwrap_key(C, K) != Q:
    raise RuntimeError
  if block_unwrap_key(C, K) != Q:
    raise RuntimeError
  return C

cryptInit()
atexit.register(cryptEnd)

Q = args.plaintext or "Hello!  My name is Inigo Montoya.  You broke my AES key wrapper.  Prepare to die."

K_128 = hex2bin(args.kek_128) if args.kek_128 else os.urandom(128 / 8)
K_256 = hex2bin(args.kek_256) if args.kek_256 else os.urandom(256 / 8)

C_128 = generate_and_test(KEK(K_128))
C_256 = generate_and_test(KEK(K_256))

args.output.write("/* Automated AES Key Wrap test cases generated by %s */\n" % os.path.basename(sys.argv[0]))
dump_hex("Q", Q, "Plaintext")
dump_hex("K_128", K_128, "128-bit KEK")
dump_hex("K_256", K_256, "256-bit KEK")
dump_hex("C_128", C_128, "Plaintext wrapped by 128-bit KEK")
dump_hex("C_256", C_256, "Plaintext wrapped by 256-bit KEK")