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

""" This is run after there has been a change under 'service nat'.
It looks for changes in the CGNAT configuration and sends the changes to
the dataplane. """


import os
import sys
import getopt

from collections import defaultdict
from subprocess import call
from vyatta import configd
from vyatta.npf.npf_debug import NpfDebug
from vyatta.npf.npf_store import store_cfg

FORCE = False
DO_CGNAT = False
DO_EXPORT = False

CFG_DIR = "/etc/td-agent-bit"
CFG_WRAPPER = "cgnat.conf"
CFG_SERVICE = "cgnat-service.conf"
CFG_INPUT = "cgnat-input.conf"
CFG_OUTPUT = "cgnat-output.conf"

FLUENTBIT_STOP_CMD = "systemctl stop td-agent-bit@{}.service"
FLUENTBIT_START_CMD = "systemctl restart td-agent-bit@{}.service"

CONFIG_CANDIDATE = configd.Client.CANDIDATE
CONFIG_RUNNING = configd.Client.RUNNING

BASE_CGNAT_PATH = "service nat cgnat"
BASE_SYSTEM_CGNAT_PATH = "system export"

# class used for printing debugs
dbg = NpfDebug()


# Used for easy assignment of nested dictionaries
def nested_dict():
    return defaultdict(nested_dict)


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


def set_cgnat_config(key, cmd, intf="ALL"):
    fcmd = "cgn-cfg {} ".format(cmd)
    store_cfg(key, fcmd, "SET", dbg, intf)


def delete_cgnat_config(key, cmd, intf="ALL"):
    fcmd = "cgn-cfg {}".format(cmd)
    store_cfg(key, fcmd, "DELETE", dbg, intf)


def create_file(filename):
    """ Open the specified CGNAT config file for writing,
        creating the config dir if required.

        Return the file handle.
    """

    os.makedirs(CFG_DIR, mode=0o755, exist_ok=True)
    return open("{}/{}".format(CFG_DIR, filename), "w")


def remove_file(filename):
    """ Silently remove the specified CGNAT config file.
        We don't care if there are errors,
        eg the file doesn't exist.
    """

    try:
        os.remove("{}/{}".format(CFG_DIR, filename))
    except:
        pass


def process_options():

    """ Process command line options """

    global FORCE, DO_CGNAT, DO_EXPORT

    try:
        opts, args = getopt.getopt(sys.argv[1:], "fd",
                                   ['force', 'debug', 'cgnat', 'export'])

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

    for opt, arg in opts:
        if opt in ('-f', '--force'):
            FORCE = True
        elif opt in ('-d', '--debug'):
            dbg.enable()
        elif opt in '--cgnat':
            DO_CGNAT = True
        elif opt in '--export':
            DO_EXPORT = True


def build_cgnat_event_command(cfg):

    """ Build the command for "cgnat policy select" and "cgnat policy log" """

    event_cmd = ''
    event_sess_all = 'no'
    event_sess_creation = 'no'
    event_sess_deletion = 'no'
    event_sess_periodic = 0
    event_subs = 'no'

    s = None
    if 'select' in cfg:
        if 'event' in cfg['select']:
            if 'subscriber' in cfg['select']['event']:
                event_subs = 'yes'
            if 'session' in cfg['select']['event']:
                s = cfg['select']['event']['session']

    if 'log' in cfg:
        if 'subscriber' in cfg['log']:
            event_subs = 'yes'
        if 'session' in cfg['log']:
            s = cfg['log']['session']

    if s:
        if 'address-group' in s:
            event_cmd += "log-sess-group={} ".format(s['address-group'])
        if 'all-subscribers' in s:
            event_sess_all = 'yes'
        if 'creation' in s:
            event_sess_creation = 'yes'
        if 'deletion' in s:
            event_sess_deletion = 'yes'
        if 'periodic' in s:
            event_sess_periodic = s['periodic']

    event_cmd += ("log-sess-all={} log-sess-creation={} "
                  "log-sess-deletion={} log-sess-periodic={} "
                  "log-subs={} ").format(
                event_sess_all, event_sess_creation, event_sess_deletion,
                event_sess_periodic, event_subs)

    return event_cmd


def build_cgnat_policy_config(commands, cfg, tree):

    """ Process 'cgnat policy' """

    for name in cfg:
        key = "{} policy {}".format(BASE_CGNAT_PATH, name)
        cmd = ""
        dbg.pprint("name: {}".format(name))
        p = cfg[name]

        if 'priority' in p:
            cmd += "priority={} ".format(p['priority'])

        # Policy match parameters
        if 'match' in p:
            if 'source' in p['match']:
                if 'address-group' in p['match']['source']:
                    cmd += "match-ag={} ".format(
                        p['match']['source']['address-group'])

        # Policy translation parameters
        if 'translation' in p:
            if 'pool' in p['translation']:
                cmd += "pool={} ".format(p['translation']['pool'])

        cmd += build_cgnat_event_command(p)

        commands[name][tree] = (key, cmd)
        dbg.pprint("POLICY - KEY: {}".format(key))
        dbg.pprint("         CMD: {}".format(cmd))


#
# Parse port number.
#
# This is currently just a number.  A future change will allow selection from
# a limited list of strings ("dns", "http" etc.), which will then be converted
# to a port number in this function.
#
def cgnat_port_str(number_or_string):
    # Convert to string so that we can use isdigit method
    port = str(number_or_string)

    # Is port already a number?
    if port.isdigit():
        return port

    raise NameError("{} is not a recognised port".format(port))


#
# Build the per-state session timeouts command
#
# The state based session timeouts (without ports) is stored as one long string,
# which configures the values for all protocols and states, for example:
#
# cgn-cfg session-timeouts other-estab 100 other-opening 70 tcp-estab 8000
#         tcp-closing 300 tcp-opening 240 udp-estab 300 udp-opening 30

# However, the per-port timeout config is stored as a string for each
# protocol and port, for example:
#
# cgn-cfg session-timeouts tcp-estab port 2001 timeout 60
#
# The command array used to hold the lines to store needs to
# index on protocol and port, as there can be multiple entries.
#
# The two types of commands are built up under different keys ('global'
# and 'per-port') so they can be processed separately in the
# send_cgnat_sess_timeout_config() function for storing in the cstore (which
# results in the commands going to the dataplane).
#
def build_cgnat_sess_timeout_command(commands, cfg, tree):

    """ Process 'cgnat session-timeout' """

    state_to_cmd = {
        'established': 'estab',
        'partially-open': 'opening',
        'partially-closed': 'closing'
    }

    key = "{} session-timeout".format(BASE_CGNAT_PATH)
    cmd = "session-timeouts "

    for prot in cfg:
        for param2 in cfg[prot]:
            if param2 == 'port':
                # process as per-port configuration
                base_cmd = "session-timeouts {}-estab".format(prot)
                for port in cfg[prot][param2]:
                    timeout = cfg[prot][param2][port]['established']

                    port_num = cgnat_port_str(port)
                    if not port_num:
                        continue

                    pkey = "{} {} port {}".format(key, prot, port_num)
                    pcmd = "{} port {} timeout {}".format(base_cmd, port_num,
                                                          timeout)

                    commands['per-port'][prot][port_num][tree] = (pkey, pcmd)
                    dbg.pprint("PORT-SESS-TIMEOUT KEY: {}".format(pkey))
                    dbg.pprint("                  CMD: {}".format(pcmd))

            else:
                # If param2 is not 'port, then it is a state, so process to
                # create a single line for the variables protocols and states
                if param2 not in state_to_cmd:
                    err("Unexpected state of {} for {}".format(param2, prot))
                    continue
                cmd += "{}-{} {} ".format(
                    prot, state_to_cmd[param2], cfg[prot][param2])

    commands['global'][tree] = (key, cmd)
    dbg.pprint("SESS-TIMEOUT KEY: {}".format(key))
    dbg.pprint("             CMD: {}".format(cmd))


def build_cgnat_interface_config(commands, cfg, tree):

    """ Process 'cgnat interface' """

    for ifname in cfg:
        key = "{} interface {}".format(BASE_CGNAT_PATH, ifname)
        dbg.pprint("ifname: {}".format(ifname))

        if 'policy' in cfg[ifname]:
            for policy in cfg[ifname]['policy']:
                pkey = "{} {}".format(key, policy)
                cmd = "intf={} name={}".format(ifname, policy)
                commands[ifname][policy][tree] = (pkey, cmd)
                dbg.pprint("INTERFACE - KEY: {}".format(pkey))
                dbg.pprint("            CMD: {}".format(cmd))


def build_cgnat_warning_config(commands, cfg, tree):

    """ Process 'cgnat select warning event resource constraint' """

    for name in cfg:
        key = "{} select warning {}".format(BASE_CGNAT_PATH, name)
        cmd = ""
        p = cfg[name]

        if 'threshold' in p:
            cmd += "threshold {}".format(p['threshold'])
            if 'interval' in p:
                cmd += " interval {}".format(p['interval'])

        commands[name][tree] = (key, cmd)
        dbg.pprint("POLICY - KEY: {}".format(key))
        dbg.pprint("         CMD: {}".format(cmd))


def build_cgnat_events_config(commands, cfg, tree, kind):

    """ Process 'cgnat log' and 'cgnat export' """

    # kind = 'rte_log' | 'protobuf'

    for event in cfg:   # event = session, subscriber, resource-constraint, ...
        key = "{} events {} {}".format(BASE_CGNAT_PATH, kind, event)
        dbg.pprint("event: {}".format(event))
        cmd = "{} enable".format(event)
        commands[event]['enable'][tree] = (key, cmd)

        dbg.pprint("EVENTS - KEY: {}".format(key))
        dbg.pprint("         CMD: {}".format(cmd))

        # Save any kafka config
        if 'using' in cfg[event]:
            if 'kafka' in cfg[event]['using']:
                if 'cluster' in cfg[event]['using']['kafka']:
                    commands[event]['cluster'][tree] = \
                        cfg[event]['using']['kafka']['cluster']

                if 'with' in cfg[event]['using']['kafka']:
                    commands[event]['with'][tree]['topic'] = \
                        cfg[event]['using']['kafka']['with'] \
                        .get('topic')
                    commands[event]['with'][tree]['key-field'] = \
                        cfg[event]['using']['kafka']['with'] \
                        .get('key-field')
                    commands[event]['with'][tree]['field-delimiter'] = \
                        cfg[event]['using']['kafka']['with'] \
                        .get('field-delimiter')
                    commands[event]['with'][tree]['priority'] = \
                        cfg[event]['using']['kafka']['with'] \
                        .get('priority')
                    commands[event]['with'][tree]['storage-limit'] = \
                        cfg[event]['using']['kafka']['with'] \
                        .get('storage-limit')


def build_cgnat_config(commands, cfg, tree):

    """ Process keywords under the top-level 'cgnat' """

    if 'policy' in cfg:
        build_cgnat_policy_config(commands['policy'], cfg['policy'], tree)

    if 'disable-hairpinning' in cfg:
        key = "{} disable-hairpinning".format(BASE_CGNAT_PATH)
        commands['disable-hairpinning'][tree] = key

    if 'snat-alg-bypass' in cfg:
        key = "{} snat-alg-bypass".format(BASE_CGNAT_PATH)
        commands['snat-alg-bypass'][tree] = key

    if 'max-sessions' in cfg:
        key = "{} max-sessions".format(BASE_CGNAT_PATH)
        commands['max-sessions'][tree] = (key, cfg['max-sessions'])

    if 'max-dest-per-session' in cfg:
        key = "{} max-dest-per-session".format(BASE_CGNAT_PATH)
        commands['max-dest-per-session'][tree] = (
            key, cfg['max-dest-per-session'])

    if 'session-timeout' in cfg:
        build_cgnat_sess_timeout_command(commands['session-timeout'],
                                         cfg['session-timeout'], tree)

    if 'interface' in cfg:
        build_cgnat_interface_config(commands['interface'], cfg['interface'],
                                     tree)

    if 'select' in cfg:
        if 'warning' in cfg['select']:
            if 'event' in cfg['select']['warning']:
                if 'resource-constraint' in cfg['select']['warning']['event']:
                    key = "{} resource-constraint".format(BASE_CGNAT_PATH)
                    build_cgnat_warning_config(
                        commands['resource-constraint'],
                        cfg['select']['warning']['event']
                           ['resource-constraint'],
                        tree)

    if 'log' in cfg:
        if 'event' in cfg['log']:
            build_cgnat_events_config(commands['events']['rte_log'],
                                      cfg['log']['event'],
                                      tree, "rte_log")
    if 'export' in cfg:
        if 'event' in cfg['export']:
            build_cgnat_events_config(commands['events']['protobuf'],
                                      cfg['export']['event'],
                                      tree, "protobuf")
    if 'cpu-affinity' in cfg:
        if 'event' in cfg['cpu-affinity']:
            if 'session' in cfg['cpu-affinity']['event']:
                key = "{} cpu-affinity event session".format(BASE_CGNAT_PATH)
                commands['session-cpu-core'][tree] = (
                    key, cfg['cpu-affinity']['event']['session'])


def build_export_kafka_cluster_config(commands, cfg):

    """ Process 'system export kafka cluster' """

    if 'bootstrap' not in cfg.keys():
        return

    commands['ipv4-address'] = cfg['bootstrap'].get('ipv4-address')
    commands['ipv6-address'] = cfg['bootstrap'].get('ipv6-address')
    commands['routing-instance'] = cfg['bootstrap'].get('routing-instance')


def build_export_kafka_config(commands, cfg, tree):

    """ Process 'system export kafka' """

    if 'cluster' in cfg:
        for cluster in cfg['cluster']:
            build_export_kafka_cluster_config(
                    commands[cluster][tree],
                    cfg['cluster'][cluster])


def build_export_config(commands, cfg, tree):

    """ Process 'system export' """

    if 'kafka' in cfg:
        build_export_kafka_config(commands['export'], cfg['kafka'], tree)


def send_cgnat_policy_config(commands):

    """ Send config for 'cgnat policy' """

    for name in commands:

        rc = commands[name].get('running')
        cc = commands[name].get('cand')

        if rc is None:
            if cc is not None:
                # this is new configuration
                set_cgnat_config(cc[0], "policy add {} {}".format(name, cc[1]))
        else:
            if cc is None:
                # configuration is being deleted
                delete_cgnat_config(rc[0], "policy delete {}".format(name))
            else:
                # configuration is being updated
                if cc[1] != rc[1]:
                    # only send if changed
                    set_cgnat_config(cc[0], "policy add {} {}".format(name,
                                     cc[1]))
                else:
                    dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".
                               format(cc[0], cc[1]))


def send_cgnat_disable_hairpinning_config(commands):

    """ Send config for 'cgnat disable-hairpinning' """

    rc = commands.get('running')
    cc = commands.get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc, "hairpinning off")
    else:
        if cc is None:
            # configuration is being deleted
            delete_cgnat_config(rc, "hairpinning on")


def send_cgnat_snat_alg_bypass_config(commands):

    """ Send config for 'cgnat snat-alg-bypass' """

    rc = commands.get('running')
    cc = commands.get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc, "snat-alg-bypass off")
    else:
        if cc is None:
            # configuration is being deleted
            delete_cgnat_config(rc, "snat-alg-bypass on")


def send_cgnat_max_sessions_config(commands):

    """ Send config for 'cgnat max-sessions' """

    rc = commands.get('running')
    cc = commands.get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc[0], "max-sessions {}".format(cc[1]))
    else:
        if cc is None:
            # configuration is being deleted - 0 indicates use default
            delete_cgnat_config(rc[0], "max-sessions 0")
        else:
            # configuration is being updated
            if cc[1] != rc[1]:
                # only send if changed
                set_cgnat_config(cc[0], "max-sessions {}".format(cc[1]))
            else:
                dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".
                           format(cc[0], cc[1]))


def send_cgnat_max_dest_per_session_config(commands):

    """ Send config for 'cgnat max-dest-per-session' """

    rc = commands.get('running')
    cc = commands.get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc[0], "max-dest-per-session {}".format(cc[1]))
    else:
        if cc is None:
            # configuration is being deleted - 0 indicates use default
            delete_cgnat_config(rc[0], "max-dest-per-session 0")
        else:
            # configuration is being updated
            if cc[1] != rc[1]:
                # only send if changed
                set_cgnat_config(cc[0], "max-dest-per-session {}".format(
                    cc[1]))
            else:
                dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".format(
                    cc[0], cc[1]))


def send_cgnat_sess_timeout_config(commands):

    """ Send config for 'cgnat session-timeout' """

    # First handle the protocol and state config, which is stored as
    # a single line under the 'global' entry.

    rc = commands.get('global', {}).get('running')
    cc = commands.get('global', {}).get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc[0], cc[1])
    else:
        if cc is None:
            # configuration is being deleted - put back to defaults
            delete_cgnat_config(rc[0], "session-timeouts tcp-opening 240 "
                                "tcp-estab 7440 tcp-closing 240 "
                                "udp-opening 30 udp-estab 300 "
                                "other-opening 30 other-estab 240")
        else:
            # configuration is being updated
            if cc[1] != rc[1]:
                # only send if changed
                set_cgnat_config(cc[0], cc[1])
            else:
                dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".format(
                    cc[0], cc[1]))

    # Next handle the per-port settings, which has entries per-protocol
    # and per-port and is under the 'per-port' entry.

    pcommands = commands.get('per-port', {})

    for prot in pcommands:
        for port in pcommands[prot]:

            rc = pcommands[prot][port].get('running')
            cc = pcommands[prot][port].get('cand')

            if rc is None:
                if cc is not None:
                    # this is new configuration
                    set_cgnat_config(cc[0], cc[1])
            else:
                if cc is None:
                    # configuration is being deleted, so put back to the default
                    # - this is done by setting the timeout to 0

                    # remove last word (i.e. old timeout) and add timeout 0
                    cmd = rc[1].rsplit(' ', 1)[0] + ' 0'
                    delete_cgnat_config(rc[0], cmd)
                else:
                    # configuration is being updated
                    if cc[1] != rc[1]:
                        # only send if changed
                        set_cgnat_config(cc[0], cc[1])
                    else:
                        dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".
                                   format(cc[0], cc[1]))


def send_cgnat_interface_config(commands):

    """ Send config for 'cgnat interface' """

    for ifname in commands:
        for policy in commands[ifname]:

            rc = commands[ifname][policy].get('running')
            cc = commands[ifname][policy].get('cand')

            if rc is None:
                if cc is not None:
                    # this is new configuration
                    set_cgnat_config(cc[0], "policy attach {}".format(cc[1]),
                                     ifname)
            else:
                if cc is None:
                    # configuration is being deleted
                    delete_cgnat_config(rc[0], "policy detach {}".format(
                        rc[1]), ifname)
                else:
                    # configuration is being updated
                    if cc[1] != rc[1]:
                        # only send if changed
                        set_cgnat_config(cc[0], "policy attach {}".format(
                            cc[1]), ifname)
                    else:
                        dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".
                                   format(cc[0], cc[1]))


def send_cgnat_resource_constraint_config(commands):

    """ Send config for 'cgnat resource-constraint' """

    for name in commands:

        rc = commands[name].get('running')
        cc = commands[name].get('cand')

        if rc is None:
            if cc is not None:
                # this is new configuration
                set_cgnat_config(cc[0], "warning add {} {}".format(name,
                                                                   cc[1]))
        else:
            if cc is None:
                # configuration is being deleted
                delete_cgnat_config(rc[0], "warning del {}".format(name))
            else:
                # configuration is being updated
                if cc[1] != rc[1]:
                    # only send if changed
                    set_cgnat_config(cc[0], "warning add {} {}".format(name,
                                     cc[1]))
                else:
                    dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".
                               format(cc[0], cc[1]))


def send_cgnat_fluentbit_service_config(create):

    """ Write the fluentbit service file

        "create" indicates whether to create (True) or remove (False) the file.
    """

    # TODO: this could be a static file

    if create:
        f = create_file(CFG_SERVICE)

        f.write("[SERVICE]\n")
        f.write("    Flush             1\n")
        f.write("    Daemon            Off\n")

        if dbg.is_enabled():
            f.write("    Log_Level         debug\n")
        else:
            f.write("    Log_Level         info\n")

        f.write("    HTTP_Server       On\n")
        f.write("    HTTP_Listen       127.0.0.1\n")
        f.write("    HTTP_Port         2020\n")
        f.write("    Grace             0\n")

        f.write("    storage.path      /opt/vyatta/tmp/td-agent-bit/storage\n")
        f.write("    storage.sync      full\n")
        f.write("    storage.checksum  on\n")

        f.write("\n")

        f.close()
    else:
        remove_file(CFG_SERVICE)


def send_cgnat_fluentbit_input_config(commands):

    """ Write the fluentbit input config file """

    # Define High Water Mark
    zmq_hwm_critical = 100000      # ie very large but not infinite
    zmq_hwm_noncritical = 100000   # ie very large but not infinite

    # First, tell the dataplane to stop existing kafka exports.

    for event in commands:    # 'session', 'subscriber', ...
        if 'enable' in commands[event]:
            if 'running' in commands[event]['enable']:
                rc = commands[event]['enable']['running']
                delete_cgnat_config(rc[0],
                                    "events protobuf {} disable"
                                    .format(event))

    # Now rewrite the input config file

    f = create_file(CFG_INPUT)

    wrote_output = False

    for event in commands:    # 'session', 'subscriber', ...

        # 'cluster' and 'with' are mandatory
        if 'cluster' not in commands[event] or 'with' not in commands[event]:
            continue

        cluster = commands[event]['cluster'].get('cand')
        withcfg = commands[event]['with'].get('cand')

        # These are mandatory in the yang,
        # but check anyway.
        if not cluster or not withcfg:
            continue

        f.write("[INPUT]\n")
        f.write("    Name            zmq\n")
        f.write("    Endpoint        "
                "ipc:///var/run/vyatta/cgnat-event-{}\n".format(event))

        # TODO: del_cgnat_config("cgnat events protobuf {} hwm".format(event))
        if withcfg.get('priority') == 'critical':
            # priority critical
            f.write("    Storage.type    filesystem\n")
            f.write("    Hwm             {}\n".format(zmq_hwm_critical))
            # TODO: compare this key with the other cc[0] keys
            set_cgnat_config("cgnat events protobuf {} hwm"
                             .format(event),
                             "events protobuf {} hwm {}"
                             .format(event, zmq_hwm_critical))
        else:
            # priority not critical
            f.write("    Storage.type    memory\n")
            f.write("    Hwm             {}\n".format(zmq_hwm_noncritical))
            set_cgnat_config("cgnat events protobuf {} hwm"
                             .format(event),
                             "events protobuf {} hwm {}"
                             .format(event, zmq_hwm_noncritical))

        if withcfg.get('storage-limit') is not None:
            f.write("    Mem_Buf_Limit   {}M\n"
                    .format(withcfg['storage-limit']))
        else:
            # default storage-limit
            f.write("    Mem_Buf_Limit   5M\n")

        f.write("    Log_Type        {}\n".format(event))
        f.write("    Tag             cgnat-{}\n".format(event))
        f.write("    Topic           {}\n".format(withcfg['topic']))

        if withcfg.get('key-field') is not None:
            f.write("    Key_field      ")
            for keyfield in withcfg['key-field']:
                f.write(" {}".format(keyfield))
            f.write("\n")

        if withcfg.get('field-delimiter') is not None:
            f.write("    Field_delimiter {}\n".
                    format(withcfg['field-delimiter']))
        else:
            # default field_delimiter
            f.write("    Field_delimiter _\n")

        f.write("\n")
        wrote_output = True

    f.close()

    if not wrote_output:
        remove_file(CFG_INPUT)

    # Finally, tell the dataplane to re-enable existing kafka exports.

    for event in commands:    # 'session', 'subscriber', ...
        if 'enable' in commands[event]:
            if 'running' in commands[event]['enable']:
                rc = commands[event]['enable']['running']
                set_cgnat_config(rc[0],
                                 "events protobuf {} enable"
                                 .format(event))


def send_cgnat_events_cfg_dataplane(commands):

    """ Configure the dataplane rte_log and protobuf exports. """

    for kind in commands:               # 'rte_log' or 'protobuf'
        for event in commands[kind]:    # 'session', 'subscriber', ...

            if 'enable' in commands[kind][event]:
                rc = commands[kind][event]['enable'].get('running')
                cc = commands[kind][event]['enable'].get('cand')

                if rc is None:
                    if cc is not None:
                        # this is new configuration
                        set_cgnat_config(cc[0],
                                         "events {} {} enable"
                                         .format(kind, event))
                else:
                    if cc is None:
                        # configuration is being deleted
                        delete_cgnat_config(rc[0],
                                            "events {} {} disable"
                                            .format(kind, event))
                    else:
                        # configuration is being updated
                        if cc[1] != rc[1]:
                            # only send if changed
                            set_cgnat_config(cc[0],
                                             "events {} {}"
                                             .format(kind, event))

                        else:
                            dbg.pprint("No change so not sent: "
                                       "KEY: {}; PARAM: {}"
                                       .format(cc[0], cc[1]))


def restart_fluentbit(r):

    """ Start or stop fluentbit as required """

    if r:
        # First stop any existing FB instances
        restart_fluentbit(False)

        # Try to start fluentbit
        try:
            cfg = (client.tree_get_dict("system export kafka",
                                        CONFIG_CANDIDATE, 'internal')
                   ['kafka']['cluster'])
            for cluster in cfg.keys():
                vrf = cfg[cluster]['bootstrap']['routing-instance']
                call(FLUENTBIT_START_CMD.format(vrf).split())
        except:
            # "system export kafka" isn't configured.
            pass

    else:
        # Try to stop fluentbit
        try:
            cfg = (client.tree_get_dict("system export kafka",
                                        CONFIG_RUNNING, 'internal')
                   ['kafka']['cluster'])
            for cluster in cfg.keys():
                vrf = cfg[cluster]['bootstrap']['routing-instance']
                call(FLUENTBIT_STOP_CMD.format(vrf).split())
        except:
            # "system export kafka" isn't configured.
            pass


def send_cgnat_events_config(commands):

    """ Send config for 'cgnat log event' and 'cgnat export event' """

    global client

    if 'protobuf' in commands:
        send_cgnat_fluentbit_input_config(commands['protobuf'])

    # Finally, tell the dataplane which events to export
    send_cgnat_events_cfg_dataplane(commands)


def send_cgnat_session_cpu_core_config(commands):

    """ Send config for 'cgnat cpu-affinity event session' """

    rc = commands.get('running')
    cc = commands.get('cand')

    if rc is None:
        if cc is not None:
            # this is new configuration
            set_cgnat_config(cc[0], "events core {}".format(cc[1]))
    else:
        if cc is None:
            # configuration is being deleted
            delete_cgnat_config(rc[0], "events core")
        else:
            # configuration is being updated
            if cc[1] != rc[1]:
                # only send if changed
                set_cgnat_config(cc[0], "events core {}".format(cc[1]))
            else:
                dbg.pprint("No change so not sent: KEY: {}; PARAM: {}".format(
                    cc[0], cc[1]))


def send_cgnat_config(commands):

    """ Send config for top-level 'cgnat' """

    dbg.pprint("send_cgnat_config()")

    if 'policy' in commands:
        send_cgnat_policy_config(commands['policy'])

    if 'disable-hairpinning' in commands:
        send_cgnat_disable_hairpinning_config(commands['disable-hairpinning'])

    if 'snat-alg-bypass' in commands:
        send_cgnat_snat_alg_bypass_config(commands['snat-alg-bypass'])

    if 'max-sessions' in commands:
        send_cgnat_max_sessions_config(commands['max-sessions'])

    if 'max-dest-per-session' in commands:
        send_cgnat_max_dest_per_session_config(
            commands['max-dest-per-session'])

    if 'session-timeout' in commands:
        send_cgnat_sess_timeout_config(commands['session-timeout'])

    if 'interface' in commands:
        send_cgnat_interface_config(commands['interface'])

    if 'resource-constraint' in commands:
        send_cgnat_resource_constraint_config(commands['resource-constraint'])

    if 'events' in commands:
        send_cgnat_events_config(commands['events'])

    if 'session-cpu-core' in commands:
        send_cgnat_session_cpu_core_config(commands['session-cpu-core'])


def program_cgnat_config():

    """ Build and send the top-level 'cgnat' config """

    global FORCE, client
    dbg.pprint("program_cgnat_config()")

    commands = nested_dict()

    try:
        status = client.node_get_status(CONFIG_CANDIDATE, BASE_CGNAT_PATH)

        if status == client.UNCHANGED and not FORCE:
            dbg.pprint("unchanged: {} so no work to do".
                       format(BASE_CGNAT_PATH))
            return 0

        try:
            cand_cfg = (client.tree_get_dict(BASE_CGNAT_PATH,
                                             CONFIG_CANDIDATE,
                                             'internal')
                        ['cgnat'])
            dbg.pprint("BUILD CANDIDATE")
            build_cgnat_config(commands, cand_cfg, 'cand')
        except configd.Exception:
            dbg.pprint("failed getting candidate tree for {}".
                       format(BASE_CGNAT_PATH))

    except configd.Exception:
        dbg.pprint("there is no configuration under {}".format(
                   BASE_CGNAT_PATH))

    try:
        running_cfg = (client.tree_get_dict(BASE_CGNAT_PATH,
                                            CONFIG_RUNNING,
                                            'internal')
                       ['cgnat'])
        dbg.pprint("BUILD RUNNING")
        build_cgnat_config(commands, running_cfg, 'running')
    except configd.Exception:
        dbg.pprint("failed getting running tree for {}".format(
                   BASE_CGNAT_PATH))

    # Send commands to the dataplane using cstore, which will change
    # the running configuration into the candidate configuration
    send_cgnat_config(commands)
    return 0


def write_fluentbit_wrapper_config(create):

    """ Write the fluentbit wrapper file
        which pulls together the system, input, and output files.

        "create" indicates whether to create (True) or remove (False) the file.
    """

    # TODO: this could be a static file

    if create:
        f = create_file(CFG_WRAPPER)
        f.write("@INCLUDE cgnat-*.conf\n")
        f.close()
    else:
        remove_file(CFG_WRAPPER)


def write_fluentbit_output_config(events_cfg, commands):

    """ Write the fluentbit output file

        If there are no events configured then delete the file.
    """

    # Track whether any OUTPUT sections were written to the file
    wrote_output = False

    # Don't write any output if there are no events configured
    if events_cfg:
        f = create_file(CFG_OUTPUT)

        # Iterate through "system export kafka cluster N"
        for cluster in commands.keys():

            f.write("[OUTPUT]\n")
            f.write("    Name                  kafka\n")
            f.write("    Format                raw\n")
            f.write("    Topic_Key             topic\n")
            f.write("    Message_Key           key\n")
            f.write("    Message_Key_Is_Name   true\n")

            # Brokers comes from commands['ipv4-address'] and ['ipv6-address']
            # Either list may be "None" if empty.
            ipv4_addrs = commands[cluster]['cand'].get('ipv4-address') or set()
            ipv6_addrs = commands[cluster]['cand'].get('ipv6-address') or set()

            if ipv4_addrs or ipv6_addrs:
                ipv4_addrs.update(ipv6_addrs)
                f.write("    Brokers               {}\n".
                        format("[" + '],['.join(ipv4_addrs) + "]"))

            f.write("    Routing_Instance      {}\n".
                    format(commands[cluster]['cand']
                           ['routing-instance']))

            # Build a list of topics
            # for the configured events ('session', 'subscriber', ...)
            # which are using this cluster.
            events = set()
            for event in events_cfg:
                if cluster in \
                        events_cfg[event]['using']['kafka']['cluster']:
                    events.add(events_cfg[event]['using']['kafka']
                                         ['with']['topic'])
            f.write("    Topics                {}\n".format(",".join(events)))

            f.write("    Match                 *\n")
            f.write("    Retry_Limit           false\n")

            wrote_output = True

        f.close()

    else:
        remove_file(CFG_OUTPUT)

    return wrote_output


def send_system_export_config(commands):

    """ Send config for 'system_export' """

    dbg.pprint("send_system_export_config()")

    if 'export' in commands:

        # Get the "cgant export event" config
        # so we can extract the names of the configured events.

        events_cfg = None

        try:
            events_cfg = (client.tree_get_dict(BASE_CGNAT_PATH,
                                               CONFIG_CANDIDATE,
                                               'internal')
                          ['cgnat'])
            if 'export' in events_cfg and 'event' in events_cfg['export']:
                events_cfg = events_cfg['export']['event']
            else:
                events_cfg = None

            # Ensure that the fluentbit config wrapper and service file exist
            write_fluentbit_wrapper_config(True)
            send_cgnat_fluentbit_service_config(True)

            # Write the "output" config file.
            # The return value tells us whether any output sections
            # were written,
            # which tells us whether to stop or to restart fluentbit.
            r = write_fluentbit_output_config(events_cfg, commands['export'])
            restart_fluentbit(r)

        except configd.Exception:
            # There are no events configured, and that's OK.
            dbg.pprint("failed getting candidate tree for {}".
                       format(BASE_CGNAT_PATH))

            # Delete all the fluentbit cgnat files and stop fluentbit.
            write_fluentbit_wrapper_config(False)
            send_cgnat_fluentbit_service_config(False)
            write_fluentbit_output_config(None, None)
            restart_fluentbit(False)


def program_system_export_config(FORCE):

    """ Build and send top-level 'system export' config """

    global client

    commands = nested_dict()

    try:
        status = client.node_get_status(CONFIG_CANDIDATE,
                                        BASE_SYSTEM_CGNAT_PATH)

        if status == client.UNCHANGED and not FORCE:
            dbg.pprint("unchanged: {} so no work to do".
                       format(BASE_SYSTEM_CGNAT_PATH))
            return 0

        try:
            cand_cfg = (client.tree_get_dict(BASE_SYSTEM_CGNAT_PATH,
                                             CONFIG_CANDIDATE,
                                             'internal')
                        ['export'])
            dbg.pprint("BUILD CANDIDATE")
            build_export_config(commands, cand_cfg, 'cand')
        except configd.Exception:
            dbg.pprint("failed getting candidate tree for {}".
                       format(BASE_SYSTEM_CGNAT_PATH))

    except configd.Exception:
        dbg.pprint("there is no configuration under {}".format(
                   BASE_SYSTEM_CGNAT_PATH))

    try:
        running_cfg = (client.tree_get_dict(BASE_SYSTEM_CGNAT_PATH,
                                            CONFIG_RUNNING,
                                            'internal')
                       ['export'])
        dbg.pprint("BUILD RUNNING")
        build_export_config(commands, running_cfg, 'running')
    except configd.Exception:
        dbg.pprint("failed getting running tree for {}".format(
                   BASE_SYSTEM_CGNAT_PATH))

    # Send commands to the dataplane using cstore, which will change
    # the running configuration into the candidate configuration
    send_system_export_config(commands)
    return 0


if __name__ == "__main__":
    process_options()

    try:
        client = configd.Client()
    except Exception as exc:
        err("Cannot establish client session: '{}'".format(str(exc).strip()))
        exit(1)

    FORCE = True

    # This script is called from both the "service nat cgnat"
    # and "system export" config trees.
    #
    # We need to do some work if either tree is modified,
    # but not duplicate the work and do things twice.
    #
    # If we're called for the "service nat cgnat" tree (DO_CGNAT)
    # and the "system export ..." config is unchanged,
    # then we won't be called again,
    # so we need to do the work here.
    #
    # However if the "system export ..." config has changed then we defer,
    # because we'll be called again momentarily for the "system export" tree.
    #
    # If we're called for the "system export" tree (DO_EXPORT),
    # then we always do the necessary work.

    if (DO_EXPORT or
        (DO_CGNAT and
            ((not client.node_exists(client.CANDIDATE,
                                     BASE_SYSTEM_CGNAT_PATH)) or
             (client.node_get_status(CONFIG_CANDIDATE,
                                     BASE_SYSTEM_CGNAT_PATH)
              == client.UNCHANGED)))):

        program_cgnat_config()
        program_system_export_config(FORCE)
