#!/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: mkpubkeys [manifestfile.yaml [pubkeyfile.yaml]]

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__":
    if len(sys.argv) > 1:
        MANIFEST = sys.argv[1]
    if len(sys.argv) > 2:
        PUBKEYS = sys.argv[2]

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

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

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