Add prefix to NIPAP

def add_prefix(arg, opts, shell_opts):
    """ Add prefix to NIPAP
    """

    # sanity checks
    if 'from-pool' not in opts and 'from-prefix' not in opts and 'prefix' not in opts:
        print("ERROR: 'prefix', 'from-pool' or 'from-prefix' must be specified.", file=sys.stderr)
        sys.exit(1)

    if len([opt for opt in opts if opt in ['from-pool', 'from-prefix', 'prefix']]) > 1:
        print("ERROR: Use either assignment 'from-pool', 'from-prefix' or manual mode (using 'prefix')", file=sys.stderr)
        sys.exit(1)

    if 'from-pool' in opts:
        return add_prefix_from_pool(arg, opts)

    args = {}
    p = _prefix_from_opts(opts)
    p.vrf = get_vrf(opts.get('vrf_rt'), abort=True)

    for avp in opts.get('extra-attribute', []):
        try:
            key, value = avp.split('=', 1)
        except ValueError:
            print("ERROR: Incorrect extra-attribute: %s. Accepted form: 'key=value'\n" % avp, file=sys.stderr)
            return
        p.avps[key] = value

    if 'from-prefix' in opts:
        args['from-prefix'] = [ opts['from-prefix'], ]

    if 'prefix_length' in opts:
        args['prefix_length'] = int(opts['prefix_length'])

    if 'family' in opts:
        if opts['family'] == 'ipv4':
            family = 4
        elif opts['family'] == 'ipv6':
            family = 6
        elif opts['family'] == 'dual-stack':
            print("ERROR: dual-stack mode only valid for from-pool assignments", file=sys.stderr)
            sys.exit(1)

        args['family'] = family

    # try to automatically figure out type for new prefix when not
    # allocating from a pool

    # get a list of prefixes that contain this prefix
    vrf_id = 0
    if p.vrf:
        vrf_id = p.vrf.id

    if 'from-prefix' in args:
        parent_prefix = args['from-prefix'][0]
        parent_op = 'equals'
    else:
        # If no prefix length is specified it is assumed to be a host and we do
        # a search for prefixes that contains the specified prefix. The last
        # entry will be the parent of the new prefix and we can look at it to
        # determine type.
        # If prefix length is specified (i.e. CIDR format) we check if prefix
        # length equals max length in which case we assume a host prefix,
        # otherwise we search for the network using an equal match and by
        # zeroing out bits in the host part.
        if len(opts.get('prefix').split("/")) == 2:
            ip = IPy.IP(opts.get('prefix').split("/")[0])
            plen = int(opts.get('prefix').split("/")[1])
            if ip.version() == 4 and plen == 32 or ip.version() == 6 and plen == 128:
                parent_prefix = str(ip)
                parent_op = 'contains'
            else:
                parent_prefix = str(IPy.IP(opts.get('prefix'), make_net=True))
                parent_op = 'equals'
        else:
            parent_prefix = opts.get('prefix')
            parent_op = 'contains'

    auto_type_query = {
            'val1': {
                'val1'      : 'prefix',
                'operator'  : parent_op,
                'val2'      : parent_prefix
                },
            'operator': 'and',
            'val2': {
                'val1'      : 'vrf_id',
                'operator'  : 'equals',
                'val2'      : vrf_id
                }
        }
    res = Prefix.search(auto_type_query, { })

    # no results, ie the requested prefix is a top level prefix
    if len(res['result']) == 0:
        if p.type is None:
            print("ERROR: Type of prefix must be specified ('assignment' or 'reservation').", file=sys.stderr)
            sys.exit(1)
    else:
        # last prefix in list will be the parent of the new prefix
        parent = res['result'][-1]

        # if the parent is an assignment, we can assume the new prefix to be
        # a host and act accordingly
        if parent.type == 'assignment':
            # automatically set type
            if p.type is None:
                print("WARNING: Parent prefix is of type 'assignment'. Automatically setting type 'host' for new prefix.", file=sys.stderr)
            elif p.type == 'host':
                pass
            else:
                print("WARNING: Parent prefix is of type 'assignment'. Automatically overriding specified type '%s' with type 'host' for new prefix." % p.type, file=sys.stderr)
            p.type = 'host'

            # if it's a manually specified prefix
            if 'prefix' in opts:
                # fiddle prefix length to all bits set
                if parent.family == 4:
                    p.prefix = p.prefix.split('/')[0] + '/32'
                else:
                    p.prefix = p.prefix.split('/')[0] + '/128'

            # for from-prefix, we set prefix_length to host length
            elif 'from-prefix' in opts:
                if parent.family == 4:
                    args['prefix_length'] = 32
                else:
                    args['prefix_length'] = 128

    try:
        p.save(args)
    except NipapError as exc:
        print("Could not add prefix to NIPAP: %s" % str(exc), file=sys.stderr)
        sys.exit(1)

    if p.type == 'host':
        print("Host %s added to %s: %s" % (p.display_prefix,
                vrf_format(p.vrf), p.node or p.description))
    else:
        print("Network %s added to %s: %s" % (p.display_prefix,
                vrf_format(p.vrf), p.description))

    if opts.get('add-hosts') is not None:
        if p.type != 'assignment':
            print("ERROR: Not possible to add hosts to non-assignment", file=sys.stderr)
            sys.exit(1)

        for host in opts.get('add-hosts').split(','):
            h_opts = {
                    'from-prefix': p.prefix,
                    'vrf_rt': p.vrf.rt,
                    'type': 'host',
                    'node': host
                    }
            add_prefix({}, h_opts, {})