#!/usr/bin/python
#
# Copyright 2013-2014.  Cumulus Networks, Inc.
#
# cl-resource-query --
#    tool to report hardware route table occupancy
#
try:
    import argparse
    import json
    import sys
except ImportError as e:
    raise ImportError(str(e) + "- required module not found")

stat_path = '/cumulus/switchd/run/'


class stat:

    def __init__(self, name, key, count_file, max_file, alloc_file=None, rule_class_list = None, att = None):
        self.name = name
        self.key = key
        self.count_file = count_file
        self.max_file = max_file
        self.alloc_file = alloc_file
        self.rule_class_list = rule_class_list
        self.att = att
stat_list = [stat('Host 0 entries',  'host_0_entry',      'route_info/host/count_0',      'route_info/host/max_0'),
             stat('Host 1 entries',  'host_1_entry',      'route_info/host/count_1',      'route_info/host/max_1'),
             stat('IPv4 neighbors',  'host_v4_entry',     'route_info/host/count_v4',     None),
             stat('IPv6 neighbors',  'host_v6_entry',     'route_info/host/count_v6',     None),
             stat('Route 0 entries', 'route_0_entry',     'route_info/route/count_0',     'route_info/route/max_0'),
             stat('Route 1 entries', 'route_1_entry',     'route_info/route/count_1',     'route_info/route/max_1'),
             stat('IPv4 Routes',     'route_v4_entry',    'route_info/route/count_v4',    None),
             stat('IPv6 Routes',     'route_v6_entry',    'route_info/route/count_v6',    None),
             stat('Total Routes',    'route_total_entry', 'route_info/route/count_total', 'route_info/route/max_total'),
             stat('ECMP nexthops',   'ecmp_nh_entry',     'route_info/ecmp_nh/count',     'route_info/ecmp_nh/max', att = 'bcm'),
             stat('Unicast Adjacency entries',   'ecmp_nh_entry',     'route_info/ecmp_nh/count', 'route_info/ecmp_nh/max', att = 'mlx'),
             stat('ECMP entries',    'ecmp_entry',        'ecmp/count',                   'ecmp/max',  att = 'mlx'),
             stat('MAC entries',     'mac_entry',         'route_info/mac/count',         'route_info/mac/max', att = 'bcm'),
             stat('MAC entries',     'mac_entry', ['route_info/mac/count', 'route_info/route/router_mac_count'], 'route_info/mac/max', att = 'mlx'),
             stat('Total Mcast Routes', 'mroute_total_entry', 'route_info/mroute/hw_count_total', 'route_info/mroute/max_total'),
             stat('Ingress ACL entries', 'in_acl_entry',  'acl_info/ingress/entries',     'acl_info/ingress/entries_total'),
             stat('Ingress ACL counters','in_acl_counter','acl_info/ingress/counters',    'acl_info/ingress/counters_total', att = 'bcm'),
             stat('Ingress ACL meters',  'in_acl_meter',  'acl_info/ingress/meters',      'acl_info/ingress/meters_total', att = 'bcm'),
             stat('Ingress ACL slices',  'in_acl_slice',  'acl_info/ingress/slices',      'acl_info/ingress/slices_total', att = 'bcm'),
             stat('Egress ACL entries',  'eg_acl_entry',  'acl_info/egress/entries',      'acl_info/egress/entries_total'),
             stat('Egress ACL counters', 'eg_acl_counter','acl_info/egress/counters',     'acl_info/egress/counters_total', att = 'bcm'),
             stat('Egress ACL meters',   'eg_acl_meter',  'acl_info/egress/meters',       'acl_info/egress/meters_total', att = 'bcm'),
             stat('Egress ACL slices',   'eg_acl_slice',  'acl_info/egress/slices',       'acl_info/egress/slices_total', att = 'bcm'),

             stat('Ingress ACL ipv4_mac filter table',   'in_acl_v4mac_filter',  'acl_info/ingress/v4mac_filter/entries_used', 'acl_info/ingress/v4mac_filter/entries_max', 'acl_info/ingress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL ipv6 filter table',   'in_acl_v6_filter',  'acl_info/ingress/v6_filter/entries_used', 'acl_info/ingress/v6_filter/entries_max', 'acl_info/ingress/v6_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL mirror table',   'in_acl_mirror_filter',  'acl_info/ingress/mirror_filter/entries_used', 'acl_info/ingress/mirror_filter/entries_max', 'acl_info/ingress/mirror_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL 8021x filter table',   'in_acl_8021x_filter',  'acl_info/ingress/8021x_filter/entries_used', 'acl_info/ingress/8021x_filter/entries_max', 'acl_info/ingress/8021x_filter/entries_allocated', att = 'bcm'),
             stat('Ingress PBR ipv4_mac filter table',   'in_pbr_v4mac_filter',  'iprule/info/ingress/v4mac_filter/entries_used', 'iprule/info/ingress/v4mac_filter/entries_max', 'iprule/info/ingress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Ingress PBR ipv6 filter table',   'in_pbr_v6_filter',  'iprule/info/ingress/v6_filter/entries_used', 'iprule/info/ingress/v6_filter/entries_max', 'iprule/info/ingress/v6_filter/entries_allocated', att = 'bcm'),

             stat('Ingress ACL ipv4_mac mangle table',   'in_acl_v4mac_mangle',  'acl_info/ingress/v4mac_mangle/entries_used', 'acl_info/ingress/v4mac_mangle/entries_max', 'acl_info/ingress/v4mac_mangle/entries_allocated', att = 'bcm'),
             stat('Ingress ACL ipv6 mangle table',   'in_acl_v6_mangle',  'acl_info/ingress/v6_mangle/entries_used', 'acl_info/ingress/v6_mangle/entries_max', 'acl_info/ingress/v6_mangle/entries_allocated', att = 'bcm'),

             stat('Egress ACL ipv4_mac filter table',   'eg_acl_v4mac_filter',  'acl_info/egress/v4mac_filter/entries_used', 'acl_info/egress/v4mac_filter/entries_max', 'acl_info/egress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Egress ACL ipv6 filter table',   'eg_acl_v6_filter',  'acl_info/egress/v6_filter/entries_used', 'acl_info/egress/v6_filter/entries_max', 'acl_info/egress/v6_filter/entries_allocated', att = 'bcm'),

             stat('ACL L4 port range checkers',   'acl_l4_port_range_checkers',  'acl_info/l4_port_range_checkers/entries_used', 'acl_info//l4_port_range_checkers/entries_max', att = 'bcm'),

             stat('ACL Regions',   'acl_mlx_regions',  'acl_info/mlx_regions/used', 'acl_info/mlx_regions/max', att = 'mlx'),
             stat('ACL 18B Rules Key',   'acl_mlx_18b_rules_key',  'acl_info/mlx_18b_rules_key/entries_used', 'acl_info/mlx_18b_rules_key/entries_max', att = 'mlx'),
             stat('ACL 36B Rules Key',   'acl_mlx_36b_rules_key',  'acl_info/mlx_32b_rules_key/entries_used', 'acl_info/mlx_32b_rules_key/entries_max', att = 'mlx'),
             stat('ACL 54B Rules Key',   'acl_mlx_54b_rules_key',  'acl_info/mlx_54b_rules_key/entries_used', 'acl_info/mlx_54b_rules_key/entries_max', att = 'mlx'),
             stat('Ingress ACL mac filter table',   'in_acl_mac_filter',  'acl_info/ingress/mac_filter/entries_used', None, None, ['acl_info/ingress/mac_filter_18b/entries_used','acl_info/ingress/mac_filter_36b/entries_used','acl_info/ingress/mac_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv4 filter table',   'in_acl_v4_filter',  'acl_info/ingress/v4_filter/entries_used', None, None, ['acl_info/ingress/v4_filter_18b/entries_used','acl_info/ingress/v4_filter_36b/entries_used','acl_info/ingress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv6 filter table',   'in_acl_v6_filter',  'acl_info/ingress/v6_filter/entries_used', None, None, ['acl_info/ingress/v6_filter_18b/entries_used','acl_info/ingress/v6_filter_36b/entries_used','acl_info/ingress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL mac filter table',   'eg_acl_mac_filter',  'acl_info/egress/mac_filter/entries_used', None, None, ['acl_info/egress/mac_filter_18b/entries_used','acl_info/egress/mac_filter_36b/entries_used','acl_info/egress/mac_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv4 filter table',   'eg_acl_v4_filter',  'acl_info/egress/v4_filter/entries_used', None, None, ['acl_info/egress/v4_filter_18b/entries_used','acl_info/egress/v4_filter_36b/entries_used','acl_info/egress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv6 filter table',   'eg_acl_v6_filter',  'acl_info/egress/v6_filter/entries_used', None, None, ['acl_info/egress/v6_filter_18b/entries_used','acl_info/egress/v6_filter_36b/entries_used','acl_info/egress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv4 mangle table',   'in_acl_v4_mangle',  'acl_info/ingress/v4_mangle/entries_used', None, None,['acl_info/ingress/v4_mangle_18b/entries_used','acl_info/ingress/v4_mangle_36b/entries_used','acl_info/ingress/v4_mangle_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv6 mangle table',   'in_acl_v6_mangle',  'acl_info/ingress/v6_mangle/entries_used', None, None, ['acl_info/ingress/v6_mangle_18b/entries_used','acl_info/ingress/v6_mangle_36b/entries_used','acl_info/ingress/v6_mangle_54b/entries_used'], att = 'mlx'),
             stat('Ingress PBR ipv4 filter table',   'in_pbr_v4_filter',  'iprule/info/ingress/v4_filter/entries_used', None, None, ['iprule/info/ingress/v4_filter_18b/entries_used','iprule/info/ingress/v4_filter_36b/entries_used','iprule/info/ingress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress PBR ipv6 filter table',   'in_pbr_v6_filter',  'iprule/info/ingress/v6_filter/entries_used', None, None, ['iprule/info/ingress/v6_filter_18b/entries_used','iprule/info/ingress/v6_filter_36b/entries_used','iprule/info/ingress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Flow Counters',   'mlx_flow_counters',  'acl_info/mlx_flow_counters/used', 'acl_info/mlx_flow_counters/max', att = 'mlx'),
             stat('RIF Basic Counters',   'mlx_rif_counters_basic',  'acl_info/mlx_rif_counters_basic/used', 'acl_info/mlx_rif_counters_basic/max', att = 'mlx'),
             stat('RIF Enhanced Counters',   'mlx_rif_counters_enh',  'acl_info/mlx_rif_counters_enh/used', 'acl_info/mlx_rif_counters_enh/max', att = 'mlx'),

             stat('Static SNAT entries',          'static_snat_count',    'nat_info/static/snat_count',       'nat_info/static/snat_max'),
             stat('Static DNAT entries',          'static_dnat_count',    'nat_info/static/dnat_count',       'nat_info/static/dnat_max'),
             stat('Dynamic SNAT entries',         'dynamic_snat_count',   'nat_info/dynamic/snat_count',      'nat_info/dynamic/snat_max'),
             stat('Dynamic DNAT entries',         'dynamic_dnat_count',   'nat_info/dynamic/dnat_count',      'nat_info/dynamic/dnat_max'),
             stat('Dynamic Config SNAT entries',  'config_snat_count',    'nat_info/config_rule/snat_count',  'nat_info/config_rule/snat_max'),
             stat('Dynamic Config DNAT entries',  'config_dnat_count',    'nat_info/config_rule/dnat_count',  'nat_info/config_rule/dnat_max'),
             ]

class bcolors:
    HEADER  = '\033[95m'
    OKBLUE  = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL    = '\033[91m'
    ENDC    = '\033[0m'

    def disable(self):
        self.HEADER  = ''
        self.OKBLUE  = ''
        self.OKGREEN = ''
        self.WARNING = ''
        self.FAIL    = ''
        self.ENDC    = ''


def generate_report(key_value_flag, use_json):

    # fetch the route table mode
    route_mode_file = stat_path + 'route_info/route/mode'
    try:
        with open(route_mode_file, 'r') as f:
            route_mode = int(f.readline())
    except IOError:
        print 'route table mode is not available'
        return -1

    # fetch the host table mode
    host_mode_file = stat_path + 'route_info/host/mode'
    try:
        with open(host_mode_file, 'r') as f:
            host_mode = int(f.readline())
    except IOError:
        print 'host table mode is not available'
        return -1

    # fetch nat support state
    nat_support_file = stat_path + 'nat_support'
    try:
        with open(nat_support_file, 'r') as f:
            nat_support = (f.readline())
    except IOError:
        print 'nat support file is not available'
        return -1

    #fetch system/asic info
    system_info_file = '/cumulus/switchd/config/system_info'
    try:
        with open(system_info_file, 'r') as f:
            system_info = (f.readline())
    except IOError:
        print 'system info file is not available'
        return -1

    # Used to dump json output
    stats_dict = {}

    for stat in stat_list:

        if stat.key == 'route_0_entry':
            if route_mode == 1:
                stat.name = 'IPv4/IPv6 route entries'
            elif route_mode == 2:
                stat.name = 'IPv4 route entries'

        if stat.key == 'route_1_entry':
            if route_mode == 1:
                stat.name = 'Long IPv6 route entries'
            elif route_mode == 2:
                stat.name = 'IPv6 route entries'

        if stat.key == 'host_0_entry':
            if host_mode == 1:
                stat.name = 'IPv4/IPv6 host entries'
            elif host_mode == 2:
                stat.name = 'IPv4 host entries'

        if stat.key == 'host_1_entry':
            if host_mode == 1:
                # not used
                continue
            elif host_mode == 2:
                stat.name = 'IPv6 host entries'

        if 'nat' in stat.key:
            if 'TRUE' in nat_support:
                if (stat.key == 'static_snat_count' or \
                   stat.key == 'static_dnat_count') and \
                   'Spectrum' in system_info:
                    continue
                if (stat.key == 'config_snat_count' or \
                   stat.key == 'config_dnat_count') and \
                   'BCM' in system_info:
                    continue
            else:
                continue
        if (stat.att == 'mlx' and \
           'BCM' in system_info):
               continue

        if (stat.att == 'bcm' and \
           'Spectrum' in system_info):
               continue
        count_files = list()
        if type(stat.count_file) is list:
            for file in stat.count_file:
                count_files.append(stat_path + file)
        else:
            count_file = stat_path + stat.count_file
        max_file = stat_path + stat.max_file if stat.max_file else None
        alloc_file = stat_path + stat.alloc_file if stat.alloc_file else None
        percent_occupancy = 0
        rule_class_list = []
        rule_count_list = []

        if stat.rule_class_list:
            for _ in stat.rule_class_list:
                rule_class_list.append(stat_path + _)

        if ((stat.key == 'in_acl_entry' or stat.key == 'eg_acl_entry') and \
            'Spectrum' in system_info):
            max_file = None
            alloc_file = None

        count = 0
        if type(stat.count_file) is list:
            for item in count_files:
                try:
                    with open(item, 'r') as f:
                        count += int(f.readline())
                except IOError:
                    print '%s count is not available' % stat.name
                    return -1
        else:
            try:
                with open(count_file, 'r') as f:
                    count = int(f.readline())
            except IOError:
                print '%s count is not available' % stat.name
                return -1
 
        if max_file:
            try:
                with open(max_file, 'r') as f:
                    max_value = int(f.readline())
                    #Max value being -1 indicates that this resource
                    #is not available, as per switchd
                    #This check makes sure that non availability of
                    #resource should be displayed as 0
                    if max_value == -1:
                        max_value = 0

            except IOError:
                print '%s max value is not available' % stat.name
                return -1

            if float(max_value) > 0:
                percent_occupancy = float(count) / float(max_value)
        else:
            max_value = 0

        if alloc_file:
            try:
                with open(alloc_file, 'r') as f:
                    alloc_value = int(f.readline())
            except IOError:
                print '%s alloc value is not available' % stat.name
                return -1
        else:
            alloc_value = 0

        if len(rule_class_list) != 0:
            for fh in rule_class_list:
                try:
                    with open(fh,'r') as f:
                        t = int(f.readline())
                        rule_count_list.append(t)
                except IOError:
                    print '%s rule_count is not available' % fh
                    return -1

        if key_value_flag:
            print stat.key + '_count=' + '%d' % count
            if max_file:
                print stat.key + '_max=' + '%d' % max_value
            if alloc_file:
                print stat.key + '_allocated=' + '%d' % alloc_value

        elif use_json:
            if stat.rule_class_list:
                stats_dict[stat.key] = {
                    'count': int(count),
                    'max' : 0,
                    'name': stat.name,
                    '18B' : rule_count_list[0],
                    '36B' : rule_count_list[1],
                    '54B' : rule_count_list[2],
                }
            else:
                stats_dict[stat.key] = {
                    'count': int(count),
                    'max': int(max_value),
                    'allocated': int(alloc_value),
                    'name': stat.name,
                    'percentage': int(percent_occupancy * 100),
                }

        else:
            if max_file:
                if alloc_file:
                        output = '%-34s %6d, %3d%% of maximum value %6d (allocated: %d)' % (stat.name + ':',
                                                                            count,
                                                                            int(percent_occupancy * 100),
                                                                            max_value, alloc_value)
                else:
                        output = '%-34s %6d, %3d%% of maximum value %6d' % (stat.name + ':',
                                                                    count,
                                                                    int(percent_occupancy * 100),
                                                                    max_value)
                if percent_occupancy > 0.90:
                    color = bcolors.FAIL
                elif percent_occupancy > 0.75:
                    color = bcolors.WARNING
                else:
                    color = None
                if color:
                    print color + output + bcolors.ENDC
                else:
                    print output
            else:
                if len(rule_count_list) != 0:
                    print '%-34s %6d %6s :%2d 36B :%2d 54B :%2d ' % (stat.name + ':', count,'18B',
                                                                  rule_count_list[0],
                                                                  rule_count_list[1],
                                                                  rule_count_list[2],)
                else:
                    print '%-34s %6d' % (stat.name + ':', count)

    if use_json:
        print json.dumps(stats_dict, sort_keys=True, indent=4)


def main(argv):
    """ main function """
    descr = 'Cumulus Resource Reporting'

    arg_parser = argparse.ArgumentParser(description=descr)
    # Command line arg parser
    #
    option = arg_parser.add_mutually_exclusive_group(required=False)
    option.add_argument('-k', '--key-value', dest='key_value',
                        action='store_true',
                        default=False,
                        help='Report key=value pairs')
    option.add_argument('-j', '--json',
                        action='store_true',
                        default=False,
                        help='Display JSON output')

    # Parse command line arguments
    cmdline_args = arg_parser.parse_args()

    ret = generate_report(cmdline_args.key_value, cmdline_args.json)

if __name__ == "__main__":
    main(sys.argv[1:])
