#!/usr/bin/env python

'''
Requirements
------------
* Install hkdf globally.
    sudo pip install hkdf
* Copy file `hkdf-expander` to /usr/local/bin
* Make `hkdf-expander` executable
    sudo chmod +x /usr/local/bin/hkdf-expander

Help
----
    hkdf-expander --help


Testme
------
    echo -n hFtWoAyknKOorIBXUjgOptpVwCfmOFUN | ./hkdf-expander --info "foo" --length 64 | openssl sha -sha512 -binary | base64 -w 0 && echo
'''

import sys
import argparse
import hkdf
import hashlib

parser = argparse.ArgumentParser(description='HKDF expander')
parser.add_argument('--info', default="",
                    help='HKDF info parameter')
parser.add_argument("--length", default=32, type=int,
                    help="Output length in bytes")
parser.add_argument("--hash", default='sha512',
                    help="Hash algorithm (sha1/sha256/sha512; default: sha512)")
parser.add_argument("--hex", action='store_true',
                    help="output hexadecimal (default: raw)")
parser.add_argument("-n", action='store_true',
                    help="do not output the trailing newline in hexmode")
args = parser.parse_args()

hash_algos = { 'sha1':   hashlib.sha1,
               'sha256': hashlib.sha256,
               'sha512': hashlib.sha512 }

infile = sys.stdin
data = bytearray("")
for line in infile:
    data += bytearray(line)

okm = hkdf.hkdf_expand(data, args.info, args.length, hash_algos[args.hash])

if args.hex:
    sys.stdout.write(okm.encode('hex'))
    if not args.n:
        sys.stdout.write("\n")
else:
    sys.stdout.write(okm)

