1from __future__ import division
2
3import hmac
4import hashlib
5import sys
6
7if sys.version_info[0] == 3:
8	buffer = lambda x: x
9
10def hkdf_extract(salt, input_key_material, hash=hashlib.sha512):
11	'''
12	Extract a pseudorandom key suitable for use with hkdf_expand
13	from the input_key_material and a salt using HMAC with the
14	provided hash (default SHA-512).
15
16	salt should be a random, application-specific byte string. If
17	salt is None or the empty string, an all-zeros string of the same
18	length as the hash's block size will be used instead per the RFC.
19
20	See the HKDF draft RFC and paper for usage notes.
21	'''
22	hash_len = hash().digest_size
23	if salt == None or len(salt) == 0:
24		salt = bytearray((0,) * hash_len)
25	return hmac.new(bytes(salt), buffer(input_key_material), hash).digest()
26
27def hkdf_expand(pseudo_random_key, info=b"", length=32, hash=hashlib.sha512):
28	'''
29	Expand `pseudo_random_key` and `info` into a key of length `bytes` using
30	HKDF's expand function based on HMAC with the provided hash (default
31	SHA-512). See the HKDF draft RFC and paper for usage notes.
32	'''
33	hash_len = hash().digest_size
34	length = int(length)
35	if length > 255 * hash_len:
36		raise Exception("Cannot expand to more than 255 * %d = %d bytes using the specified hash function" %\
37			(hash_len, 255 * hash_len))
38	blocks_needed = length // hash_len + (0 if length % hash_len == 0 else 1) # ceil
39	okm = b""
40	output_block = b""
41	for counter in range(blocks_needed):
42		output_block = hmac.new(pseudo_random_key, buffer(output_block + info + bytearray((counter + 1,))),\
43			hash).digest()
44		okm += output_block
45	return okm[:length]
46
47class Hkdf(object):
48	'''
49	Wrapper class for HKDF extract and expand functions
50	'''
51	def __init__(self, salt, input_key_material, hash=hashlib.sha256):
52		'''
53		Extract a pseudorandom key from `salt` and `input_key_material` arguments.
54
55		See the HKDF draft RFC for guidance on setting these values. The constructor
56		optionally takes a `hash` arugment defining the hash function use,
57		defaulting to hashlib.sha256.
58		'''
59		self._hash = hash
60		self._prk = hkdf_extract(salt, input_key_material, self._hash)
61	def expand(self, info=b"", length=32):
62		'''
63		Generate output key material based on an `info` value
64
65		Arguments:
66		- info - context to generate the OKM
67		- length - length in bytes of the key to generate
68
69		See the HKDF draft RFC for guidance.
70		'''
71		return hkdf_expand(self._prk, info, length, self._hash)
72
73