From 27f805e2379b2fce87883d1405518b586055d1a3 Mon Sep 17 00:00:00 2001
From: Rich Wareham <rjw57@cam.ac.uk>
Date: Fri, 17 Sep 2010 13:25:00 +0100
Subject: [PATCH] add a simple self-test

Added a simple self test mode where a source and sink key are generated and the
source/sink shared-secret key is generated using the source keys and the sink
keys. The test makes sure those keys are identical.
---
 README          |  4 ++++
 generate_key.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/README b/README
index e505e65..2e124bc 100644
--- a/README
+++ b/README
@@ -12,6 +12,7 @@ Options:
   -k, --sink            generate a sink key rather than a source key
   --ksv=KSV             use a specific KSV expressed in hexadecimal
   -j, --json            output key and KSV as JSON
+  -t, --test            generate source and sink keys and test they work
 
 Examples:
 
@@ -22,4 +23,7 @@ Examples:
  # human-readable form
  ./generate_key.py
 
+ # Run a self-test to make sure the source a sink key generation is consistent
+ ./generate_key.py -t
+
 % vim:spelllang=en_gb:spell
diff --git a/generate_key.py b/generate_key.py
index 92531c4..f03d5c5 100755
--- a/generate_key.py
+++ b/generate_key.py
@@ -50,12 +50,21 @@ def main():
 	parser.add_option('-j', '--json', action='store_true', 
 		dest='output_json', default=False, 
 		help='output key and KSV as JSON')
+	
+	parser.add_option('-t', '--test', action='store_true', 
+		dest='do_test', default=False, 
+		help='generate source and sink keys and test they work')
 
 	(options, args) = parser.parse_args()
 
 	# read the master key file
 	key_matrix = read_key_file(open(options.master_key_file))
 
+	# if asked to do a test, do one and exit
+	if options.do_test:
+		do_test(key_matrix)
+		return
+
 	# generate a ksv if necessary
 	if options.ksv is not None:
 		ksv = int(options.ksv, 16)
@@ -74,6 +83,39 @@ def main():
 	else:
 		output_human_readable(ksv, key, options.gen_sink)
 
+def do_test(key_matrix):
+	"""Perform a self-test.
+
+	Generate both a source key and sink key with random (different) KSVs
+	and test that the HDCP shared-key system works."""
+
+	print('Performing self test.')
+
+	# generate source key
+	src_ksv = gen_ksv() 
+	src_key = gen_source_key(src_ksv, key_matrix)
+	output_human_readable(src_ksv, src_key, False)
+
+	# generate sink key
+	snk_ksv = gen_ksv() 
+	snk_key = gen_sink_key(snk_ksv, key_matrix)
+	output_human_readable(snk_ksv, snk_key, True)
+
+	# add sink keys together according to src ksv
+	key1 = reduce(lambda x, y: (x+y) & 0xffffffffffffff, 
+		map(lambda x: x[1], filter(lambda x: src_ksv & (1<<x[0]), zip(range(40), snk_key))))
+
+	# add source keys together according to sink ksv
+	key2 = reduce(lambda x, y: (x+y) & 0xffffffffffffff, 
+		map(lambda x: x[1], filter(lambda x: snk_ksv & (1<<x[0]), zip(range(40), src_key))))
+
+	print('\nGenerated keys: sink = %014x, source = %014x' % (key1, key2))
+
+	if key1 == key2:
+		print('Test PASSED')
+	else:
+		print('Test FAILED')
+
 def read_key_file(filelike):
 	"""Read a HDCP master key from a key file.