helpers.py

"""
Helper functions used to group interfaces together and generate the
necessary data to generate the dashboards from templates.
"""
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor
import logging
import json
from itertools import product
from functools import partial, reduce
from string import ascii_uppercase
from brian_dashboard_manager.templating.render import create_panel, \
    create_panel_target, create_dropdown_panel

NUM_PROCESSES = 4

PANEL_HEIGHT = 12
PANEL_WIDTH = 24

logger = logging.getLogger(__file__)


def num_generator(start=30):
    """
    Generator for numbers starting from the value of `start`

    :param start: number to start at
    :return: generator of numbers
    """

    num = start
    while True:
        yield num
        num += 1


def gridPos_generator(id_generator, start=1, agg=False):
    """
    Generator of gridPos objects used in Grafana dashboards to position panels.

    :param id_generator: generator of panel ids
    :param start: panel number to start from
    :param agg: whether to generate a panel for the aggregate dashboards,
    which has two panels per row

    :return: generator of gridPos objects
    """

    num = start
    while True:
        yield {
            "height": PANEL_HEIGHT,
            "width": PANEL_WIDTH if not agg else PANEL_WIDTH // 2,
            "x": 0,
            "y": num * PANEL_HEIGHT,
            "id": next(id_generator)
        }
        if agg:
            yield {
                "height": PANEL_HEIGHT,
                "width": PANEL_WIDTH // 2,
                "x": PANEL_WIDTH // 2,
                "y": num * PANEL_HEIGHT,
                "id": next(id_generator)
            }
        num += 1


def letter_generator():
    """
    Generator for letters used to generate refIds for panel targets.

    :return: generator of strings
    """
    i = 0
    j = 0
    num_letters = len(ascii_uppercase)
    while True:
        result = ascii_uppercase[i % num_letters]

        # tack on an extra letter if we are out of them
        if (i >= num_letters):
            result += ascii_uppercase[j % num_letters]
            j += 1
            if (j != 0 and j % num_letters == 0):
                i += 1
        else:
            i += 1

        yield result


def get_nren_interface_data_old(interfaces):
    """
    Helper for grouping interfaces into groups of NRENs
    Extracts information from interfaces to be used in panels.
    NREN dashboards have aggregate panels at the top and
    dropdowns for services / physical interfaces.
    """
    result = {}

    for interface in interfaces:

        description = interface['description'].strip()
        interface_name = interface['name']
        host = interface['router']

        router = host.replace('.geant.net', '')
        location = host.split('.')[1].upper()
        panel_title = f"{router} - {{}} - {interface_name} - {description}"

        dashboards_info = interface['dashboards_info']

        for info in dashboards_info:
            dashboard_name = info['name']

            dashboard = result.get(dashboard_name, {
                'AGGREGATES': [],
                'SERVICES': [],
                'PHYSICAL': []
            })

            if info['interface_type'] == 'AGGREGATE':
                dashboard['AGGREGATES'].append({
                    'interface': interface_name,
                    'hostname': host,
                    'alias':
                        f"{location} - {dashboard_name} ({interface_name})"
                })

                # link aggregates are also shown
                # under the physical dropdown
                dashboard['PHYSICAL'].append({
                    'title': panel_title,
                    'hostname': host,
                    'interface': interface_name
                })

            elif info['interface_type'] == 'LOGICAL':
                dashboard['SERVICES'].append({
                    'title': panel_title,
                    'hostname': host,
                    'interface': interface_name
                })
            elif info['interface_type'] == 'PHYSICAL':
                dashboard['PHYSICAL'].append({
                    'title': panel_title,
                    'hostname': host,
                    'interface': interface_name
                })

            result[dashboard_name] = dashboard
    return result


def get_nren_interface_data(services, interfaces, excluded_dashboards):
    """
    Helper for grouping interface data to be used for generating
    dashboards for NRENs.

    Extracts information from interfaces to be used in panels.

    :param services: list of services
    :param interfaces: list of interfaces
    :param excluded_dashboards: list of dashboards to exclude for
    the organization we are generating dashboards for

    :return: dictionary of dashboards and their service/interface data
    """

    result = {}

    customers = defaultdict(list)
    aggregate_interfaces = dict()

    for service in services:
        _customers = service.get('customers')
        for cust in _customers:
            if cust.lower() in excluded_dashboards:
                continue
            customers[cust].append(service)

    for customer, services in customers.items():
        if not any([s['service_type'] == 'GEANT IP' for s in services]):
            # NREN access customers must have at least one GEANT IP service
            continue

        dashboard = result.setdefault(customer, {
            'AGGREGATES': [],
            'SERVICES': [],
            'PHYSICAL': []
        })

        for service in services:
            _interfaces = service.get('endpoints')
            name = service.get('name')
            sid = service.get('sid')
            scid = service.get('scid')
            service_type = service.get('service_type')

            measurement = 'scid_rates'

            lag_service = 'GA-' in sid and service_type == 'ETHERNET'

            if len(_interfaces) == 0:
                continue

            if 'interface' in _interfaces[0]:
                if_name = _interfaces[0].get('interface')
                router = _interfaces[0].get('hostname')
            else:
                if_name = _interfaces[0].get('port')
                router = _interfaces[0].get('equipment')
            router = router.replace('.geant.net', '')
            title = f'{router} - {{}} - {if_name} - {name} ({sid})'

            if lag_service:
                if len(_interfaces) > 1:
                    logger.info(
                        f'{sid} {name} aggregate service has > 1 interface')
                    continue

                aggregate_interfaces[f'{router}:::{if_name}'] = True
                dashboard['AGGREGATES'].append({
                    'measurement': measurement,
                    'alias': title,
                    'scid': scid
                })

            if 'MDVPN' in service['service_type']:
                # MDVPN type services don't have data in BRIAN
                continue

            has_v6_interface = False
            for interface in _interfaces:
                if 'addresses' in interface:
                    for address in interface['addresses']:
                        if address.find(':') > 0:
                            has_v6_interface = True
                            break

            dashboard['SERVICES'].append({
                'measurement': measurement,
                'title': title,
                'scid': scid,
                'sort': (sid[:2], name),
                'has_v6': has_v6_interface
            })

    def _check_in_aggregate(router, interface):
        return f'{router}:::{interface}' in aggregate_interfaces

    for interface in interfaces:

        description = interface['description'].strip()
        interface_name = interface['name']
        host = interface['router']
        port_type = interface.get('port_type', 'unknown').lower()

        router = host.replace('.geant.net', '')
        panel_title = f"{router} - {{}} - {interface_name} - {description}"

        dashboards_info = interface['dashboards_info']

        for info in dashboards_info:
            dashboard_name = info['name']

            dashboard = result.get(dashboard_name, {
                'AGGREGATES': [],
                'SERVICES': [],
                'PHYSICAL': []
            })

            _covered_by_service = _check_in_aggregate(router, interface_name)
            if port_type == 'access' and not _covered_by_service:
                dashboard['AGGREGATES'].append({
                    'interface': interface_name,
                    'hostname': host,
                    'alias':
                    f"{router} - {interface_name} - {dashboard_name} "
                })

            if info['interface_type'] == 'AGGREGATE':
                # link aggregates are also shown
                # under the physical dropdown
                dashboard['PHYSICAL'].append({
                    'title': panel_title,
                    'hostname': host,
                    'interface': interface_name
                })
            elif info['interface_type'] == 'PHYSICAL':
                dashboard['PHYSICAL'].append({
                    'title': panel_title,
                    'hostname': host,
                    'interface': interface_name
                })

            result[dashboard_name] = dashboard

    for customer in list(result.keys()):
        lengths = [len(val) for val in result[customer].values()]
        if sum(lengths) == 0:
            # no services/interfaces, so remove it
            del result[customer]
    return result


def get_interface_data(interfaces):
    """
    Helper for grouping interface data to be used for generating
    various dashboards

    Extracts information from interfaces to be used in panels.

    :param interfaces: list of interfaces

    :return: dictionary of dashboards and their interface data
    """

    result = {}

    for interface in interfaces:

        description = interface['description'].strip()
        interface_name = interface['name']
        host = interface['router']

        router = host.replace('.geant.net', '')
        panel_title = f"{router} - {{}} - {interface_name} - {description}"

        dashboards_info = interface['dashboards_info']

        for info in dashboards_info:
            dashboard_name = info['name']

            dashboard = result.get(dashboard_name, [])
            dashboard.append({
                'title': panel_title,
                'interface': interface_name,
                'hostname': host,
                'has_v6': len(interface.get('ipv6', [])) > 0
            })

            result[dashboard_name] = dashboard
    return result


def get_aggregate_interface_data(interfaces, agg_name, group_field):
    """
    Helper for grouping interface data to be used for generating
    aggregate dashboards.
    Aggregate dashboards have panels with multiple targets (timeseries)
    that are grouped by a field (`group_field`). This function
    groups the interfaces by the `group_field` and returns a dictionary
    of aggregate dashboards and their interface data.

    One of the panels is a special panel that has all the targets
    in a single panel, as an aggregate of all data for that dashboard.

    :param interfaces: list of interfaces
    :param agg_name: name of the aggregate dashboard
    :param group_field: field to group the interfaces by

    :return: dictionary of aggregate dashboards and their interface data
    """

    result = []

    def get_reduce_func_for_field(field):
        def reduce_func(prev, curr):
            groups = prev.get(curr[field], [])
            groups.append(curr)
            all_agg = prev.get('EVERYSINGLETARGET', [])
            all_agg.append(curr)
            prev[curr[field]] = groups
            prev['EVERYSINGLETARGET'] = all_agg
            return prev
        return reduce_func

    for interface in interfaces:

        interface_name = interface.get('name')
        host = interface.get('router', '')
        router = host.replace('.geant.net', '')

        for info in interface['dashboards_info']:
            remote = info['name']
            location = host.split('.')[1].upper()
            result.append({
                'type': agg_name,
                'interface': interface_name,
                'hostname': host,
                'remote': remote,
                'location': location,
                'alias': f"{router} - {remote} - {interface_name}",
            })
    return reduce(get_reduce_func_for_field(group_field), result, {})


def get_aggregate_targets(targets):
    """
    Helper for generating targets for aggregate panels.

    Aggregate panels have multiple targets (timeseries) that are
    grouped by a field (`group_field`).

    This function generates the targets for the aggregate panel.

    :param targets: list of targets

    :return: tuple of ingress and egress targets for the ingress and egress
    aggregate panels respectively
    """

    ingress = []
    egress = []

    # used to generate refIds
    letters = letter_generator()

    for target in targets:
        ref_id = next(letters)
        in_data = {
            **target,
            'alias': f"{target['alias']} - Ingress Traffic",
            'refId': ref_id,
            'select_field': 'ingress'
        }
        out_data = {
            **target,
            'alias': f"{target['alias']} - Egress Traffic",
            'refId': ref_id,
            'select_field': 'egress'
        }
        ingress_target = create_panel_target(in_data)
        egress_target = create_panel_target(out_data)
        ingress.append(ingress_target)
        egress.append(egress_target)

    return ingress, egress


def get_panel_fields(panel, panel_type, datasource):
    """
    Helper for generating panels.

    Generates the fields for the panel based on the panel type.

    :param panel: panel data
    :param panel_type: type of panel (traffic, errors, etc.)
    :param datasource: datasource to use for the panel

    :return: generated panel definition from the panel data and panel type
    """

    letters = letter_generator()

    def get_target_data(alias, field):
        return {
            # panel includes identifying information
            # such as hostname, interface, etc.
            **panel,
            'alias': alias,
            'refId': next(letters),
            'select_field': field,
            'percentile': 'percentile' in alias.lower(),
            'errors': panel_type == 'errors'  # used to remove *8 on value
        }

    error_fields = [('Ingress Errors', 'errorsIn'),
                    ('Egress Errors', 'errorsOut'),
                    ('Ingress Discards', 'discardsIn'),
                    ('Egress Discards', 'discardsOut')]

    ingress = ['Ingress Traffic', 'Ingress 95th Percentile']
    egress = ['Egress Traffic', 'Egress 95th Percentile']

    is_v6 = panel_type == 'IPv6'
    is_error = panel_type == 'errors'
    in_field = 'ingressv6' if is_v6 else 'ingress'

    out_field = 'egressv6' if is_v6 else 'egress'

    fields = [*product(ingress, [in_field]), *product(egress, [out_field])]

    targets = error_fields if is_error else fields

    return create_panel({
        **panel,
        'datasource': datasource,
        'linewidth': 1,
        'title': panel['title'].format(panel_type),
        'panel_targets': [get_target_data(*target) for target in targets],
        'y_axis_type': 'errors' if is_error else 'bits',
    })


def default_interface_panel_generator(gridPos, use_all_traffic=True, use_ipv6=True):
    """
    Helper for generating panel definitions for dashboards.

    Generates the panel definitions for the dashboard based on the
    panel data and panel type.

    :param gridPos: generator for grid positions
    :param use_all_traffic: include ipv4 + other traffic interfaces
    :param use_ipv6: include ipv6 interfaces

    :return: function that generates panel definitions
    """

    def get_panel_definitions(panels, datasource, errors=False):
        """
        Generates the panel definitions for the dashboard based on the
        panel data for the panel types (traffic, errors, IPv6).

        IPv6 and errors are optional / determined by the presence of the
        `has_v6` field in the panel data, and the `errors` parameter.

        :param panels: panel data
        :param datasource: datasource to use for the panel
        :param errors: whether to include an error panel

        :return: list of panel definitions
        """
        result = []

        for panel in panels:
            if use_all_traffic:
                result.append(get_panel_fields({
                    **panel,
                    **next(gridPos)
                }, 'traffic', datasource))
            if use_ipv6:
                if panel.get('has_v6', False):
                    result.append(get_panel_fields({
                        **panel,
                        **next(gridPos)
                    }, 'IPv6', datasource))
            if errors:
                result.append(get_panel_fields({
                    **panel,
                    **next(gridPos)
                }, 'errors', datasource))

        return result

    return get_panel_definitions


def get_nren_dashboard_data_single(data, datasource, tag):
    """
    Helper for generating dashboard definitions for a single NREN.

    NREN dashboards have two aggregate panels (ingress and egress),
    and two dropdown panels for services and interfaces.

    :param data: data for the dashboard, including the NREN name and
    the panel data
    :param datasource: datasource to use for the panels
    :param tag: tag to use for the dashboard, used for dashboard dropdowns on
    the home dashboard.

    :return: dashboard definition for the NREN dashboard
    """

    nren, dash = data
    id_gen = num_generator()

    if len(dash['AGGREGATES']) > 0:
        agg_panels = create_aggregate_panel(
            f'Aggregate - {nren}',
            gridPos_generator(id_gen, agg=True),
            dash['AGGREGATES'], datasource)
        gridPos = gridPos_generator(id_gen, start=2)
    else:
        gridPos = gridPos_generator(id_gen)
        agg_panels = []

    panel_gen = default_interface_panel_generator(gridPos, use_all_traffic=True, use_ipv6=False)
    panel_ipv6_gen = default_interface_panel_generator(gridPos, use_all_traffic=False, use_ipv6=True)

    services_dropdown = create_dropdown_panel('Services', **next(gridPos))

    def sort_key(panel):
        sort = panel.get('sort')
        if not sort:
            return 'ZZZ'+panel.get('hostname')  # sort to end
        return sort

    service_panels = panel_gen(
        sorted(dash['SERVICES'], key=sort_key), datasource)

    services_ipv6_dropdown = create_dropdown_panel('Services - IPv6 Only', **next(gridPos))
    service_ipv6_panels = panel_ipv6_gen(
        sorted(dash['SERVICES'], key=sort_key), datasource
    )

    iface_dropdown = create_dropdown_panel('Interfaces', **next(gridPos))
    phys_panels = panel_gen(dash['PHYSICAL'], datasource, True)

    dropdown_groups = [{
        'dropdown': services_dropdown,
        'panels': service_panels,
    }]
    if len(service_ipv6_panels) > 0:
        dropdown_groups.append({
            'dropdown': services_ipv6_dropdown,
            'panels': service_ipv6_panels
        })
    dropdown_groups.append({
        'dropdown': iface_dropdown,
        'panels': phys_panels,
    })

    result = {
        'nren_name': nren,
        'datasource': datasource,
        'aggregate_panels': agg_panels,
        'dropdown_groups': dropdown_groups
    }
    if isinstance(tag, list):
        result['tags'] = tag
    else:
        result['tag'] = tag

    return result


def get_nren_dashboard_data(data, datasource, tag):
    """
    Helper for generating dashboard definitions for all NRENs.
    Uses multiprocessing to speed up generation.

    :param data: the NREN names and the panel data for each NREN
    :param datasource: datasource to use for the panels
    :param tag: tag to use for the dashboard, used for dashboard dropdowns on
    the home dashboard.

    :return: generator for dashboard definitions for each NREN
    """

    with ProcessPoolExecutor(max_workers=NUM_PROCESSES) as executor:
        for dash in executor.map(
            partial(
                get_nren_dashboard_data_single,
                datasource=datasource,
                tag=tag),
            data.items()
        ):

            yield dash


def get_dashboard_data_single(
        data, datasource, tag,
        panel_generator=default_interface_panel_generator,
        errors=False):
    """
    Helper for generating dashboard definitions for non-NREN dashboards.

    :param data: data for the dashboard, including the dashboard name and
    the panel data
    :param datasource: datasource to use for the panels
    :param tag: tag to use for the dashboard, used for dashboard dropdowns on
    the home dashboard.
    :param panel_generator: function for generating panel definitions
    :param errors: whether or not to include an error panel for each interface

    :return: dashboard definition for the NREN dashboard
    """

    id_gen = num_generator()
    gridPos = gridPos_generator(id_gen)
    panel_gen = panel_generator(gridPos)

    name, panels = data
    result = {
        'title': name,
        'datasource': datasource,
        'panels': list(panel_gen(panels, datasource, errors)),
    }

    if isinstance(tag, list):
        result['tags'] = tag
    else:
        result['tag'] = tag

    return result


def get_dashboard_data(
        data, datasource, tag,
        panel_generator=default_interface_panel_generator,
        errors=False):
    """
    Helper for generating dashboard definitions for all non-NREN dashboards.
    Uses multiprocessing to speed up generation.

    :param data: the dashboard names and the panel data for each dashboard
    :param datasource: datasource to use for the panels
    :param tag: tag to use for the dashboard, used for dashboard dropdowns on
    the home dashboard.
    :param panel_generator: function for generating panel definitions
    :param errors: whether or not to include an error panel for each interface

    :return: generator for dashboard definitions for each dashboard
    """

    with ProcessPoolExecutor(max_workers=NUM_PROCESSES) as executor:
        for dash in executor.map(
            partial(
                get_dashboard_data_single,
                datasource=datasource,
                tag=tag,
                panel_generator=panel_generator,
                errors=errors),
            data.items()
        ):

            yield dash


def create_aggregate_panel(title, gridpos, targets, datasource):
    """
    Helper for generating aggregate panels. Creates two panels, one for
    ingress and one for egress.

    :param title: title for the panel
    :param gridpos: generator for grid position
    :param targets: list of targets for the panels, used to build separate
    targets for both ingress and egress.
    :param datasource: datasource to use for the panels

    :return: tuple of aggregate panels, one for ingress and one for egress
    """

    ingress_targets, egress_targets = get_aggregate_targets(targets)

    ingress_pos = next(gridpos)
    egress_pos = next(gridpos)

    is_total = 'totals' in title.lower()

    def reduce_alias(prev, curr):
        d = json.loads(curr)
        alias = d['alias']
        if 'egress' in alias.lower():
            prev[alias] = '#0000FF'
        else:
            prev[alias] = '#00FF00'
        return prev

    ingress_colors = reduce(reduce_alias, ingress_targets, {})
    egress_colors = reduce(reduce_alias, egress_targets, {})

    ingress = create_panel({
        **ingress_pos,
        'stack': True,
        'linewidth': 0 if is_total else 1,
        'datasource': datasource,
        'title': title + ' - ingress',
        'targets': ingress_targets,
        'y_axis_type': 'bits',
        'alias_colors': json.dumps(ingress_colors) if is_total else {}
    })

    egress = create_panel({
        **egress_pos,
        'stack': True,
        'linewidth': 0 if is_total else 1,
        'datasource': datasource,
        'title': title + ' - egress',
        'targets': egress_targets,
        'y_axis_type': 'bits',
        'alias_colors': json.dumps(egress_colors) if is_total else {}
    })

    return ingress, egress


def get_aggregate_dashboard_data(title, remotes, datasource, tag):
    """
    Helper for generating aggregate dashboard definitions.
    Aggregate dashboards consist only of aggregate panels that are
    panels with data for multiple interfaces.

    At the top of the dashboard are two aggregate panels showing
    total ingress and egress data for all interfaces.

    Below that are two aggregate panels for each target, one for
    ingress and one for egress.

    :param title: title for the dashboard
    :param remotes: dictionary of targets for the panels, the key is the
    remote (usually a customer) and the value is a list of targets
    for that remote. A single target represents how to fetch
    data for one interface.
    :param datasource: datasource to use for the panels
    :param tag: tag to use for the dashboard, used for dashboard dropdowns on
    the home dashboard.

    :return: dashboard definition for the aggregate dashboard
    """

    id_gen = num_generator()
    gridPos = gridPos_generator(id_gen, agg=True)

    panels = []
    all_targets = remotes.get('EVERYSINGLETARGET', [])

    ingress, egress = create_aggregate_panel(
        title, gridPos, all_targets, datasource)
    panels.extend([ingress, egress])

    totals_title = title + ' - Totals'
    t_in, t_eg = create_aggregate_panel(
        totals_title, gridPos, all_targets, datasource)
    panels.extend([t_in, t_eg])

    if 'EVERYSINGLETARGET' in remotes:
        del remotes['EVERYSINGLETARGET']

    for remote in remotes:
        _in, _out = create_aggregate_panel(
            title + f' - {remote}', gridPos, remotes[remote], datasource)
        panels.extend([_in, _out])

    result = {
        'title': title,
        'datasource': datasource,
        'panels': panels,

    }
    if isinstance(tag, list):
        result['tags'] = tag
    else:
        result['tag'] = tag

    return result