#!/usr/bin/env python
#
"""
Tool to generate YubiKey secret keys using YubiHSM.

This tool is currently somewhat of a dead end. It generates secrets for
one or more YubiHSMs, but that effectively means they are not available
in clear text anywhere. One of the following pieces are missing :

 1) a "soft" YubiHSM that can decrypt the AEAD and program a YubiKey

 2) export to plain text file as well as AEADs, to then use the plain
    text file to program the YubiKey(s)

 3) the ability for this script to also program attached YubiKeys,
    either directly (through python-yubico) or by calling ykpersonalize
    for example.

Because of this, it is currently in the examples/ directory (although
I can't really tell what it would be an example of).
"""
#
# Copyright (c) 2011 Yubico AB
# See the file COPYING for licence statement.
#

import os
import sys
import argparse
sys.path.append('Lib')
import pyhsm
import pyhsm.yubikey

default_device = "/dev/ttyACM0"
default_dir = "/var/cache/yubikey-ksm/aeads"

def parse_args():
    """
    Parse the command line arguments
    """
    parser = argparse.ArgumentParser(description = "Generate secrets for YubiKeys using YubiHSM",
                                     add_help=True,
                                     formatter_class = argparse.ArgumentDefaultsHelpFormatter,
                                     )
    parser.add_argument('-D', '--device',
                        dest='device',
                        default=default_device,
                        required=False,
                        help='YubiHSM device',
                        )
    parser.add_argument('-O', '--output-dir', '--aead-dir',
                        dest='output_dir',
                        default=default_dir,
                        required=False,
                        help='Output directory (AEAD base dir)',
                        )
    parser.add_argument('-c', '--count',
                        dest='count',
                        type=int, default=1,
                        required=False,
                        help='Number of secrets to generate',
                        )
    parser.add_argument('-v', '--verbose',
                        dest='verbose',
                        action='store_true', default=False,
                        help='Enable verbose operation',
                        )
    parser.add_argument('--public-id-chars',
                        dest='public_id_chars',
                        type=int, default=12,
                        required=False,
                        help='Number of chars in generated public ids',
                        )
    parser.add_argument('--key-handles',
                        dest='key_handles',
                        nargs='+',
                        required=True,
                        help='Key handles to encrypt the generated secrets with',
                        )
    parser.add_argument('--start-public-id',
                        dest='start_id',
                        required=True,
                        help='The first public id to generate AEAD for',
                        )

    return parser.parse_args()

def args_fixup(args):
    if not os.path.isdir(args.output_dir):
        sys.stderr.write("Output directory '%s' does not exist.\n" % (args.output_dir))
        sys.exit(1)

    keyhandles_fixup(args)

    try:
        n = int(args.start_id)
    except ValueError:
        hexstr = pyhsm.yubikey.modhex_decode(args.start_id)
        n = int(hexstr, 16)

    args.start_id = n


def keyhandles_fixup(args):
    """
    Walk through the supplied key handles and normalize them, while keeping
    the input format too (as value in a dictionary). The input format is
    used in AEAD filename paths.
    """
    new_handles = {}
    for val in args.key_handles:
        for this in val.split(','):
            n = pyhsm.util.key_handle_to_int(this)
            new_handles[n] = this

    args.key_handles = new_handles


def gen_keys(hsm, args):
    """
    The main key generating loop.
    """

    if args.verbose:
        print "Generating %i keys :\n" % (args.count)
    else:
        print "Generating %i keys" % (args.count)

    for int_id in range(args.start_id, args.start_id + args.count):
        public_id = "%x" % int_id
        modhex_id = pyhsm.yubikey.modhex_encode(public_id)
        padded_id = modhex_id.rjust(args.public_id_chars, 'c')

        if args.verbose:
            print "  %s" % (padded_id)

        num_bytes = len(pyhsm.aead_cmd.YHSM_YubiKeySecret('a' * 16, 'b' * 6).pack())
        hsm.load_random(num_bytes)
        for kh in args.key_handles.keys():
            aead = hsm.generate_aead(public_id, kh)

            filename = output_filename(args.output_dir, args.key_handles[kh], padded_id)

            if args.verbose:
                print "    %4s, %i bytes (%s) -> %s" % \
                    (args.key_handles[kh], len(aead.data), shorten_aead(aead), filename)

            f = open(filename, "w")
            f.write(aead.data)
            f.close()

        if args.verbose:
            print ""

    print "\nDone\n"

def shorten_aead(aead):
    """ Produce pretty-printable version of long AEAD. """
    head = aead.data[:4].encode('hex')
    tail = aead.data[-4:].encode('hex')
    return "%s...%s" % (head, tail)

def output_filename(output_dir, key_handle, public_id):
    """
    Return an output filename for a generated AEAD. Creates a hashed directory structure
    using the last three bytes of the public id to get equal usage.
    """
    parts = [output_dir, key_handle] + pyhsm.util.group(public_id, 2)
    path = os.path.join(*parts)

    if not os.path.isdir(path):
        os.makedirs(path)

    return os.path.join(path, public_id)


def main():
    args = parse_args()

    args_fixup(args)

    print "output dir		: %s" % (args.output_dir)
    print "keys to generate	: %s" % (args.count)
    print "key handles		: %s" % (args.key_handles)
    print "start public_id		: %s (0x%x)" % (args.start_id, args.start_id)
    print "YHSM device		: %s" % (args.device)
    print ""

    hsm = pyhsm.YHSM(device = args.device)

    gen_keys(hsm, args)


if __name__ == '__main__':
    main()
