#!/usr/bin/python3
#
# Copyright (c) 2019-2020 AT&T Intellectual Property.
# All rights reserved.
#
# SPDX-License-Identifier: GPL-2.0-only
#

""" This is run to get the CGNAT subscriber information from the dataplane
and provide it as YANG RPC information. """


import os
import sys
import getopt
import vplaned
import json
from vyatta.npf.npf_debug import NpfDebug


PROGNAME = os.path.basename(__file__)
DATAPLANE_CMD = 'cgn-op show subscriber'

# class used for printing debugs
dbg = NpfDebug()


def err(msg):
    print(msg, file=sys.stderr)


def process_options():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "d", ['debug'])

    except getopt.GetoptError as r:
        err(r)
        err("usage: {} [-d|--debug] ".format(sys.argv[0]))
        sys.exit(2)

    for opt, arg in opts:
        if opt in ('-d', '--debug'):
            dbg.enable()


def sub_rpc(entry):
    rpc_prot_list = []

    for prot in ['tcp', 'udp', 'other']:
        rpc_prot_entry = {}
        rpc_prot_entry['protocol'] = prot
        rpc_prot_entry['active-mappings'] = entry[prot + '_ports_used']
        rpc_prot_list.append(rpc_prot_entry)

    rpc_entry = {}
    rpc_entry['subscriber'] = entry['address']
    rpc_entry['paired-address'] = entry['paired_addr']
    rpc_entry['duration'] = entry['duration']
    rpc_entry['expired'] = ((entry['flags'] & 0x01) != 0)
    rpc_entry['active-blocks'] = entry['block_count']
    rpc_entry['active-mappings'] = entry['map_active']
    rpc_entry['total-mapping-requests'] = entry['map_reqs']
    rpc_entry['total-mapping-failures'] = entry['map_fails']
    rpc_entry['sessions-created'] = entry['sess_crtd']
    rpc_entry['sessions-destroyed'] = entry['sess_dstrd']
    rpc_entry['session-rate-20s'] = entry['sess_rate_20s']
    rpc_entry['session-rate-1m'] = entry['sess_rate_1m']
    rpc_entry['session-rate-5m'] = entry['sess_rate_5m']
    rpc_entry['session-rate-max'] = entry['sess_rate_max']
    rpc_entry['session-rate-max-time'] = entry['sess_rate_max_tm']
    rpc_entry['session-rate-1m-max'] = entry['sess_rate_1m_max']
    rpc_entry['session-rate-1m-max-time'] = entry['sess_rate_1m_max_tm']
    rpc_entry['sub-sess-created'] = entry['sess2_crtd']
    rpc_entry['sub-sess-destroyed'] = entry['sess2_dstrd']
    rpc_entry['packets-out'] = entry['out_pkts']
    rpc_entry['packets-in'] = entry['in_pkts']
    rpc_entry['bytes-out'] = entry['out_bytes']
    rpc_entry['bytes-in'] = entry['in_bytes']
    rpc_entry['unknown-source'] = entry['unk_pkts_in']
    rpc_entry['assigned-ports'] = entry['port_count']
    rpc_entry['protocols'] = rpc_prot_list

    dbg.pprint("rpc entry: {}".format(rpc_entry))
    return rpc_entry


param_mappings = {
    'ip-address-prefix': 'address',
    'start-index': 'start',
    'max-entries': 'count',
}


def get_cgnat_subscriber_info():
    dbg.pprint("get_cgnat_subscriber_info()")

    try:
        rpc_input = json.load(sys.stdin)
    except ValueError as exc:
        err("Failed to parse input JSON: {}".format(exc))
        return None, 1

    args = DATAPLANE_CMD
    start = None
    count = None

    for param, value in rpc_input.items():
        dp_param = param_mappings.get(param)
        if dp_param is None:
            err("{}: unknown rpc input option: {}".format(PROGNAME, param))
            return None, 2

        if dp_param == 'start':
            start = value
        elif dp_param == 'count':
            count = value
        else:
            args += " {} {}".format(dp_param, value)

    # dataplane requires a start if given a count
    if count is not None and start is None:
        start = "1"

    if start is not None:
        args += " start {}".format(start)

    if count is not None:
        args += " count {}".format(count)

    dbg.pprint("dp command: {}".format(args))

    sub_info_list = []
    with vplaned.Controller() as controller:
        for dp in controller.get_dataplanes():
            with dp:
                dp_dict = dp.json_command(args)
                if dp_dict and dp_dict.get('subscribers'):
                    dbg.pprint("dataplane dict: {}".format(
                        dp_dict['subscribers']))
                    for subscriber in dp_dict['subscribers']:
                        sub_info_list.append(sub_rpc(subscriber))

    if sub_info_list:
        return {'subscribers': sub_info_list}, 0
    else:
        return None, 0


if __name__ == "__main__":
    process_options()
    info, ret = get_cgnat_subscriber_info()
    if info:
        print(json.dumps(info))
    exit(ret)
