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