#!/usr/bin/env python3
'''
Python framework for operational status test on Arista network.
'''
import os
import re
import sys
import time
import json
import argparse
from jsondiff import diff
from jmespath import search
from pyeapi import load_config
from pyeapi import connect_to
from pyateos.plugins.acl import acl
from pyateos.plugins.as_path import as_path
from pyateos.plugins.bgp_evpn import bgp_evpn
from pyateos.plugins.bgp_ipv4 import bgp_ipv4
from pyateos.plugins.interface import interface
from pyateos.plugins.ip_route import ip_route
from pyateos.plugins.mlag import mlag
from pyateos.plugins.ntp import ntp
from pyateos.plugins.prefix_list import prefix_list
from pyateos.plugins.route_map import route_map
from pyateos.plugins.snmp import snmp
from pyateos.plugins.stp import stp
from pyateos.plugins.vlan import vlan
from pyateos.plugins.vrf import vrf
from pyateos.plugins.vxlan import vxlan

def arguments():
    '''
    usage: pyATEOS [-h] (-B | -A | -C) -t TEST [TEST ...] [-g GROUP [GROUP ...]]
               [-i INVENTORY] -n NODE [NODE ...] [-F FILE [FILE ...]] [-f]

    pyATEOS - A simple python application for operational status test on Arista
    device. Based on pyATS idea and pyeapi library for API calls.

    optional arguments:
    -h, --help            show this help message and exit
    -B, --before          write json file containing the test result BEFORE. To
                            be run BEFORE the config change. File path example:
                            $PWD/before/ntp/router1_ntp.json
    -A, --after           write json file containing the test result AFTER. To
                            be run AFTER the config change. File path example:
                            $PWD/after/ip_route/router1_ip_route.json
    -C, --compare         diff between before and after test files. File path
                            example: $PWD/diff/snmp/router1_snmp.json
    -t TEST [TEST ...], --test TEST [TEST ...]
                            run one or more specific test. Multiple values are
                            accepted separated by space
    -g GROUP [GROUP ...], --group GROUP [GROUP ...]
                            run a subset of test. Options available: mgmt,
                            routing, layer2, ctrl, all Multiple values are
                            accepted separated by space. Works also with -t --test
    -i INVENTORY, --inventory INVENTORY
                            specify pyeapi inventory file path
    -n NODE [NODE ...], --node NODE [NODE ...]
                            specify inventory node. Multiple values are accepted
                            separated by space
    -F FILE [FILE ...], --file_name FILE [FILE ...]
                            provide the 2 filename IDs to compare, separated by
                            space. BEFORE first, AFTER second. i.e [..] -C -f
                            1582386835 1582387929
    -f, --filter          filter counters where present
    '''

    parser = argparse.ArgumentParser(
        prog='pyATEOS',
        description='''pyATEOS - A simple python application for operational status test on
        Arista device. Based on pyATS idea and pyeapi library for API calls.'''
        )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        '-B',
        '--before',
        dest='before',
        action='store_true',
        help='''write json file containing the test result BEFORE.
        To be run BEFORE the config change. 
        File path example: $PWD/before/ntp/router1_ntp.json'''
        )
    group.add_argument(
        '-A',
        '--after',
        dest='after',
        action='store_true',
        help='''write json file containing the test result AFTER.
        To be run AFTER the config change.
        File path example: $PWD/after/ip_route/router1_ip_route.json'''
        )
    group.add_argument(
        '-C',
        '--compare',
        dest='compare',
        action='store_true',
        help='''diff between before and after test files.
        File path example: $PWD/diff/snmp/router1_snmp.json'''
        )
    parser.add_argument(
        '-t',
        '--test',
        dest='test',
        help='run one or more specific test. Multiple values are accepted separated by space',
        nargs='+',
        )
    parser.add_argument(
        '-g',
        '--group',
        dest='group',
        help='''run a subset of test. Options available: mgmt, routing, layer2, ctrl, all
        Multiple values are accepted separated by space. Works also with -t --test''',
        nargs='+',
        )
    parser.add_argument(
        '-i',
        '--inventory',
        dest='inventory',
        help='specify pyeapi inventory file path'
        )
    parser.add_argument(
        '-n',
        '--node',
        dest='node',
        help='specify inventory node. Multiple values are accepted separated by space',
        nargs='+',
        required=True
        )
    parser.add_argument(
        '-F',
        '--file_name',
        dest='file',
        help='''provide the 2 filename IDs to compare, separated by space.
        BEFORE first, AFTER second. i.e [..] -C -f 1582386835 1582387929''',
        nargs='+',
        type=int,
        )
    parser.add_argument(
        '-f',
        '--filter',
        dest='diff_filter',
        help='''filter counters where present''',
        action='store_true',
        )

    if parser.parse_args().compare and parser.parse_args().file is None:
        parser.error("-C/--compare requires -F/--file_name.")

    if not (parser.parse_args().test or parser.parse_args().group):
        parser.error("Either -t/--test or -g/--group is required.")

    return parser.parse_args()


def flags(args):
    '''
    arg.parser() assertion and unpacking
    '''
    returned_dict = dict()

    if args.compare:
        if args.file:
            assert len(args.file) == 2 and (args.file[0] - args.file[1]) < 0, '''
            provide the 2 filename IDs to compare, separated by space. 
            BEFORE first, AFTER second. i.e [..] -C -f 1582386835 1582387929'''
            file_name = args.file
    else:
        file_name = None

    returned_dict.update(
        inventory=args.inventory,
        node=args.node,
        test=args.test,
        before=args.before,
        after=args.after,
        compare=args.compare,
        file_name=file_name,
        diff_filter=args.diff_filter,
        group=args.group,
        )

    return returned_dict


class WriteFile():
    '''
    entry point for test file write and filesystem build
    '''

    def __init__(self, test, node, diff_filter=None, file_name=None):
        '''
        class' vars initialization
        '''
        self.test = test
        self.node = node
        self.file_name = file_name
        self.diff_filter = diff_filter
        self.pwd_before = os.getcwd() + '/tests/before/{0}/{1}/'.format(self.test, self.node)
        self.pwd_after = os.getcwd() + '/tests/after/{0}/{1}/'.format(self.test, self.node)
        self.pwd_diff = os.getcwd() + '/tests/diff/{0}/{1}/'.format(self.test, self.node)
        self.time_file = round(time.time())

    def write_before(self, result):
        '''
        build filesystem and write file for BEFORE tests.
        '''
        if not os.path.exists(self.pwd_before):
            os.makedirs(self.pwd_before)

        with open('{0}/{1}.json'.format(
                self.pwd_before,
                self.time_file,
            ), 'w', encoding='utf-8') as file:
            json.dump(result, file, ensure_ascii=False, indent=4)
            print('BEFORE file ID for {} test: {}'.format(self.test.upper(), self.time_file))


    def write_after(self, result):
        '''
        build filesystem and write file for AFTER tests.
        '''
        if not os.path.exists(self.pwd_after):
            os.makedirs(self.pwd_after)

        with open('{0}/{1}.json'.format(
                self.pwd_after,
                self.time_file,

            ), 'w', encoding='utf-8') as file:
            json.dump(result, file, ensure_ascii=False, indent=4)
            print('AFTER file ID for {} test: {}'.format(self.test.upper(), self.time_file))


    def write_diff(self):
        '''
        build filesystem, write and legalize file for COMPARE tests.
        '''
        class CustomFilter():
            '''
            class for filtering json outputs
            '''
            def filter_jmespath(self, test, legal_json_diff):
                '''jmespath filter to reduce output verbosity.
                Valid only for dict(dict()) type, not dict(list())'''

                plugins_filter = {
                    'ntp': 'peers.{delete: delete, insert: insert}',
                    'vlan': 'vlans.{delete: delete, insert: insert}',
                    'as_path': 'activeIpAsPathLists.{delete: delete, insert: insert}',
                    'lldp': 'lldpNeighbors.{delete: delete, insert: insert}'
                }

                if plugins_filter.get(test):
                    final_diff = search(plugins_filter.get(test), legal_json_diff)
                elif test == 'interface':
                    final_diff = CustomFilter().filter_iface_counters(legal_json_diff)
                elif test == 'acl':
                    final_diff = CustomFilter().filter_acls_counters(legal_json_diff)
                else:
                    final_diff = legal_json_diff

                return final_diff


            def filter_iface_counters(self, legal_json_diff):
                '''
                "Ethernet47": {
                    "interfaceStatus": [
                        "connected",
                        "disabled"
                    ],
                ----------------------------------------
                "Recirc-Channel1500": {
                    "memberInterfaces": {
                        "insert": {
                            "Ethernet47": {
                                "duplex": "duplexFull",
                                "bandwidth": 25000000000
                            }
                        }
                    },
                ----------------------------------------
                "insert": {
                    "Loopback999": {
                ----------------------------------------
                "delete": {
                    "Loopback270": { 
                '''
                return_dict = {'interfaces': {}}

                for ifaces in legal_json_diff.values():
                    for iface_name, iface_values in ifaces.items():

                        if iface_values.get('interfaceStatus'):
                            return_dict['interfaces'][iface_name] = legal_json_diff['interfaces'][iface_name]['interfaceStatus']

                        if iface_values.get('memberInterfaces'):
                            if iface_values['memberInterfaces'].get('delete') or iface_values['memberInterfaces'].get('insert'):
                                return_dict['interfaces'][iface_name] = legal_json_diff['interfaces'][iface_name]['memberInterfaces']

                        if iface_name == 'insert' or iface_name == 'delete':
                            return_dict['interfaces'][iface_name] = legal_json_diff['interfaces'][iface_name]

                return return_dict


            def filter_acls_counters(self, legal_json_diff):
                '''
                {
                "aclList": {
                    "0:": {
                        "sequence": {
                            "0:": {
                                "ruleFilter": {
                                    "source": {
                '''
                return_dict = {'aclList': {}}

                for acls in legal_json_diff.values():
                    for acl_number, sequences in acls.items():
                        for seq_number in sequences.values():
                            for values in seq_number.values():
                                if values.get('ruleFilter'):
                                    return_dict['aclList'][acl_number] = {'sequence': {}}
                                    return_dict['aclList'][acl_number]['sequence'] = seq_number

                                    return return_dict


        def replace(string, test):
            ''' jsondiff editing for returning legal json format'''
            substitutions = {
                '\'':'\"',
                'insert':'"insert"',
                'delete':'"delete"',
                'True':'true',
                'False':'false',
                '(':'[',
                ')':']',
            }

            skip_list = [
                'vrf',
            ]

            substrings = sorted(substitutions, key=len, reverse=True)
            regex = re.compile('|'.join(map(re.escape, substrings)))
            sub_applied = regex.sub(lambda match: substitutions[match.group(0)], string)

            if test not in skip_list:
                # ('{', '28: ') (' ', '187: ')
                for integer in re.findall(r"(\s|{)(\d+:\s)", sub_applied):
                    int_replacement = integer[1][:-2]
                    merged_integer = integer[0] + integer[1]

                    if integer[0] == '{':
                        sub_applied = sub_applied.replace(
                            # double {{ required by format to exscape {
                            merged_integer, '{{"{0}": '.format(int_replacement)
                        )

                    else:
                        sub_applied = sub_applied.replace(
                            merged_integer, ' "{0}": '.format(int_replacement)
                        )

            return sub_applied


        if not os.path.exists(self.pwd_diff):
            os.makedirs(self.pwd_diff)


        try:
            before = open('{0}/{1}.json'.format(
                self.pwd_before,
                self.file_name[0],
                ), 'r')
        except FileNotFoundError as error:
            print(error)
            sys.exit(1)

        try:
            after = open('{0}/{1}.json'.format(
                self.pwd_after,
                self.file_name[1],
                ), 'r')
        except FileNotFoundError as error:
            print(error)
            sys.exit(1)

        json_diff = str(diff(before, after, load=True, syntax='symmetric'))
        legal_json_diff = replace(json_diff, self.test)
        diff_file_id = str((self.file_name[0] - self.file_name[1]) * -1)

        if not self.diff_filter:
            final_diff = json.loads(legal_json_diff)
        else:
            final_diff = CustomFilter().filter_jmespath(self.test, json.loads(legal_json_diff))

        with open('{0}/{1}.json'.format(
                self.pwd_diff,
                diff_file_id,
                ), 'w', encoding='utf-8') as file:
            json.dump(final_diff, file, ensure_ascii=False, indent=4)
            print('DIFF file ID for {} test: {}'.format(self.test.upper(), diff_file_id))


def bef_aft_com(**kwargs):
    '''
    arguments assertion, test groups definitation and run test
    '''

    test_run = list()

    if kwargs:
        before = kwargs.get('before')
        after = kwargs.get('after')
        compare = kwargs.get('compare')
        file_name = kwargs.get('file_name')
        diff_filter = kwargs.get('diff_filter')
        group = kwargs.get('group')

        nodes = kwargs.get('node')
        assert isinstance(nodes, list), '''
        "nodes" must be of type list()'''

        test = kwargs.get('test')

        if test:
            assert isinstance(test, list), '''
            "test" must be of type list()'''

        if group:
            assert isinstance(group, list), '''
            "group" must be of type list()'''

    test_all = [
        'acl',
        'arp',
        'as_path',
        'bgp_evpn',
        'bgp_ipv4',
        'interface',
        'ip_route',
        'mac',
        'mlag',
        'ntp',
        'lldp',
        'prefix_list',
        'route_map',
        'snmp',
        'stp',
        'vlan',
        'vrf',
        'vxlan',
    ]

    if group:
        if 'mgmt' in group:
            group.remove('mgmt')
            test_run.extend(('ntp', 'snmp'))
        elif 'routing' in group:
            group.remove('routing')
            test_run.extend(('bgp_evpn', 'bgp_ipv4', 'ip_route'))
        elif 'layer2' in group:
            group.remove('layer2')
            test_run.extend(('stp', 'vlan', 'vxlan', 'lldp', 'arp', 'mac'))
        elif 'ctrl' in group:
            group.remove('ctrl')
            test_run.extend(('acl', 'as_path', 'prefix_list', 'route_map'))
        elif 'all' in group:
            group.remove('all')
            test_run = test_all

    if test:
        test_run.extend(test)

    for node in nodes:
        for test in list(set(test_run)):

            if test not in test_all:
                raise ValueError('''Test name not valid. Please check if the test liist
                available under plugins/ folder and test_all list updated.''')

            if before:
                wr_before = WriteFile(test, node)
                wr_before.write_before(eval(test)(connect_to(node)).show)
            elif after:
                wr_after = WriteFile(test, node)
                wr_after.write_after(eval(test)(connect_to(node)).show)
            elif compare:
                wr_diff = WriteFile(test, node, diff_filter, file_name)
                wr_diff.write_diff()


def pyateos(**kwargs):
    '''
    Application entry point.

    kwargs:
        before:
            options: [ False, Ture ]
            type: bool()
            required: False

        after:
            options: [ False, Ture ]
            type: bool()
            required: False

        compare:
            options: [ False, Ture ]
            type: bool()
            required: False

        file_name:
            type: int()
            required: if compare is True

        nodes:
            type: list()
            required: True

        group:
            type: list

        test:
            type: list()
            options: [
                'acl',
                'arp',
                'as_path',
                'bgp_evpn',
                'bgp_ipv4',
                'interface',
                'ip_route',
                'mac',
                'mlag',
                'ntp',
                'lldp',
                'prefix_list',
                'route_map',
                'snmp',
                'stp',
                'vlan',
                'vrf',
                'vxlan',
            ]
    '''

    if kwargs.get('inventory'):
        inventory = kwargs.get('inventory')
    else:
        inventory = 'eos_inventory.ini'

    load_config(inventory)

    if kwargs.get('test') or kwargs.get('group'):
        bef_aft_com(**kwargs)


if __name__ == '__main__':

    if flags(arguments()):
        kwargs = flags(arguments())

    pyateos(**kwargs)

#f20200302 v0.1.6
