1 files changed, 66 insertions, 0 deletions
diff --git a/generate-test-cases b/generate-test-cases
new file mode 100644
index 0000000..2e40d5f
--- /dev/null
+++ b/generate-test-cases
@@ -0,0 +1,66 @@
+#!/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")