tools/test/bin/onos-gen-config

#!/usr/bin/env python
"""
usage: onos-gen-config [-h] [-s PARTITION_SIZE] [-n NUM_PARTITIONS]
                            [filename] [node_ip [node_ip ...]]

Generate the partitions json file given a list of IPs or from the $OCC*
environment variables.

positional arguments:
  filename              File to write output to. If none is provided, output
                        is written to stdout.
  node_ip               IP Address(es) of the node(s) in the cluster. If no
                        IPs are given, will use the $OC* environment
                        variables. NOTE: these arguments are only processed
                        after the filename argument.

optional arguments:
  -h, --help            show this help message and exit
  -s PARTITION_SIZE, --partition-size PARTITION_SIZE
                        Number of nodes per partition. Note that partition
                        sizes smaller than 3 are not fault tolerant. Defaults
                        to 3.
  -n NUM_PARTITIONS, --num-partitions NUM_PARTITIONS
                        Number of partitions. Defaults to the number of nodes
                        in the cluster.
"""

from os import environ
import argparse
import re
import json

convert = lambda text: int(text) if text.isdigit() else text.lower()
alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]


def get_vars(type):
    types = {'controller': 'OC', 'storage': 'OCC'}
    vars = []
    for var in environ:
        if re.match(r"{}[0-9]+".format(types[type]), var):
            vars.append(var)
    return sorted(vars, key=alphanum_key)


def get_nodes(type, ips=None, default_port=5679):
    node = lambda id, ip, port: {'id': id, 'ip': ip, 'port': port}
    prefixes = {'controller': 'onos', 'storage': 'atomix'}
    result = []
    if not ips:
        ips = [environ[v] for v in get_vars(type)]
    i = 1
    for ip_string in ips:
        address_tuple = ip_string.split(":")
        if len(address_tuple) == 3:
            id = address_tuple[0]
            ip = address_tuple[1]
            port = int(address_tuple[2])
        else:
            id = '{}-{}'.format(prefixes[type], i)
            i += 1
            ip = ip_string
            port = default_port
        result.append(node(id, ip, port))
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Generate the partitions json file given a list of IPs or from environment variables.")
    parser.add_argument(
        'filename', metavar='filename', type=str, nargs='?',
        help='File to write output to. If none is provided, output is written to stdout.')
    parser.add_argument(
        '--controller-nodes', '-c', metavar='node_ip', type=str, nargs='+',
        help='IP Address(es) of the controller nodes. If no IPs are given, ' +
             'will use the $OC* environment variables. NOTE: these arguemnts' +
             ' are only processed after the filename argument.')
    parser.add_argument(
        '--storage-nodes', '-s', metavar='node_ip', type=str, nargs='+',
        help='IP Address(es) of the storage nodes. If no IPs are given, ' +
             'will use the $OCC* environment variables. NOTE: these arguemnts' +
             ' are only processed after the filename argument.')

    args = parser.parse_args()
    filename = args.filename
    controller = get_nodes('controller', args.controller_nodes)
    storage = get_nodes('storage', args.storage_nodes)

    if len(storage) > 0:
        data = {
            'name': 'onos',
            'storage': storage
        }
    else:
        data = {
            'name': 'onos',
            'controller': controller
        }
    output = json.dumps(data, indent=4)

    if filename:
        with open(filename, 'w') as f:
            f.write(output)
    else:
        print output