#!/usr/bin/env python3

# Read a manifest.yaml file, and for each listed node, execute it with
# the --gen option to generate (or load if they already exist) a
# private/public key pair.  It will output the public key to stdout,
# which we capture.  At the end, output all of the public keys to a
# pubkeys.yaml file.

# Usage: getpubkeys [-m manifest.yaml] [-p pubkeys.yaml]

import argparse
import re
import shlex
import subprocess
import sys
import yaml

# The default input file
MANIFEST = "manifest.yaml"

# The default output file
PUBKEYS = "pubkeys.yaml"

# The TEEMS binary
TEEMS = "./teems"

def getkey(node, manifestdata):
    """Get the public key for a particular node (as a 128-character hex
    string).  The passed manifestdata is the dictionary corresponding to
    that node in the manifest."""
    cmdline = ''
    if 'launchprefix' in manifestdata:
        cmdline = manifestdata['launchprefix'] + ' '
    cmdline += TEEMS + " --gen %(sprvfile)s %(pubfile)s" % manifestdata
    pubkey = None
    try:
        print("Fetching pubkey for", node)
        out = subprocess.run(shlex.split(cmdline), stdout=subprocess.PIPE) \
            .stdout.decode('utf-8')
        res = re.search('^Pubkey: ([0-9a-f]{128})', out, re.MULTILINE)
        if res:
            pubkey = res.group(1)
        else:
            print(res)
    except:
        # Couldn't find the output
        pass
    return pubkey

if __name__ == "__main__":
    aparse = argparse.ArgumentParser(
        description='Create a TEEMS pubkeys.yaml file from a manifest.yaml file'
    )
    aparse.add_argument('-m', default=MANIFEST,
        help='manifest.yaml input file')
    aparse.add_argument('-p', default=PUBKEYS,
        help='pubkeys.yaml output file')
    args = aparse.parse_args()

    # A dictionary to store the output pubkeys
    pubkeys = {}

    # Read the manifest
    with open(args.m) as mf:
        manifest = yaml.safe_load(mf)
        for (node, data) in manifest.items():
            if node != "params":
                key = getkey(node, data)
                if key is not None:
                    pubkeys[node] = key

    print('')
    print(yaml.dump(pubkeys))
    with open(args.p, 'w') as pf:
        yaml.dump(pubkeys, pf)