teems/App/mkconfig.py

#!/usr/bin/env python3

# Given a manifest.yaml file, a pubkeys.yaml file, and (optionally) a
# list of nodes to include, generate the config.json file to give to
# each of the nodes.  If the list of nodes is missing, include all nodes
# in the manifest.yaml file.  It is an error to include a node that is
# missing from either the manifest.yaml or the pubkeys.yaml files.

# Usage: mkconfig [-m manifest.yaml] [-p pubkeys.yaml] [-n node1 node2...]

import argparse
import json
import yaml

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

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

def create_json(manifestfile, pubkeysfile, nodelist, params_override):
    """Given a manifest.yaml file, a pubkeys.yaml file, and (optionally) a
    list of nodes to include and a dictionary of parameter settings to
    override, generate the config.json file to give to each of the
    nodes.  If the list of nodes is missing, include all nodes in the
    manifest.yaml file.  It is an error to include a node that is
    missing from either the manifest.yaml or the pubkeys.yaml files."""
    with open(manifestfile) as mf:
        manifest = yaml.safe_load(mf)
    with open(pubkeysfile) as pf:
        pubkeys = yaml.safe_load(pf)
    if nodelist is None or len(nodelist) == 0:
        nodelist = manifest.keys()
    config = {}
    if "params" in manifest:
        config['params'] = manifest['params']
        for ov in params_override:
            if params_override[ov] is not None:
                config['params'][ov] = params_override[ov]
    config['nodes'] = []
    for node in nodelist:
        if node == "params":
            continue
        nodeconf = {}
        m = manifest[node]
        nodeconf['name'] = node
        nodeconf['pubkey'] = pubkeys[node]
        nodeconf['listen'] = m['listen']
        # Optional fields
        for f in ['clisten', 'slisten', 'weight', 'roles']:
            if f in m:
                nodeconf[f] = m[f]
        config['nodes'].append(nodeconf)
    return json.dumps(config)

if __name__ == "__main__":
    aparse = argparse.ArgumentParser(
        description='Create a TEEMS config.json file from a manifest and a pubkeys file'
    )
    aparse.add_argument('-m', default=MANIFEST,
        help='manifest.yaml file')
    aparse.add_argument('-p', default=PUBKEYS,
        help='pubkeys.yaml file')
    aparse.add_argument('-z', default=None,
        help='override message size')
    aparse.add_argument('-u', default=None,
        help='override max number of users')
    aparse.add_argument('-B', default=None,
        help='override max number of outgoing private messages per user per epoch')
    aparse.add_argument('-b', default=None,
        help='override max number of incoming private messages per user per epoch')
    aparse.add_argument('-C', default=None,
        help='override max number of outgoing public messages per user per epoch')
    aparse.add_argument('-c', default=None,
        help='override max number of incoming public messages per user per epoch')
    aparse.add_argument('-n', nargs='*', help='nodes to include')
    args = aparse.parse_args()

    params_overrides = {
        'msg_size': args.z,
        'user_count': args.u,
        'priv_out': args.B,
        'priv_in': args.b,
        'pub_out': args.C,
        'pub_in': args.c,
    }
    json = create_json(args.m, args.p, args.n, params_overrides)

    if json is not None:
        print(json)