#!/usr/bin/python
# Copyright 2014,2016 Cumulus Networks, Inc.
#
# portsamp
#
# Utility to configure packet sampling on physical interfaces.
#
import sys
import cumulus.porttab as porttab
import cumulus.platforms
import os.path
import subprocess

try:
    import bcmshell
    HAS_BCM_SHELL = True
except ImportError:
    HAS_BCM_SHELL = False


base_incr = 0


def usage():
    sys.stderr.write("Usage: %s (swp##|swp*) ingr_rate egr_rate "
                     "nflog_base_group\n" % sys.argv[0])
    sys.stderr.write("  On average, every 1/rate packets will be sampled.  A\n")
    sys.stderr.write("  rate of 0 means no sampling, 1 means all packets sampled.\n")
    sys.stderr.write("  Most recommend monitoring only ingress traffic, but if\n")
    sys.stderr.write("  ingr_rate and egr_rate are specified an nflog_base_group\n")
    sys.stderr.write("  + 1 is allocated for egress samples.\n")
    sys.exit(10)

if len(sys.argv) != 5:
    usage()

intf = sys.argv[1]
try:
    ingr_rate = int(sys.argv[2])
    egr_rate = int(sys.argv[3])
    nflog_base_group = int(sys.argv[4])
except:
    sys.stderr.write("Rates and ulog_group must be integers.\n\n")
    usage()
if nflog_base_group < 0 or nflog_base_group > 65535:
    sys.stderr.write("nflog_base_group must be between 1 and 65535\n\n")
    usage()

pt = porttab.porttab()
if intf == "swp*":
    interfaces = pt.get_linux_ports()
else :
    interfaces = [intf]

platform_object = cumulus.platforms.probe()
if platform_object.switch.chip.sw_base == 'mlx':

    if egr_rate > 0:
        sys.stderr.write("Non-zero egr_rate not supported on Mellanox platforms\n\n")
        sys.exit(10)

    with open('/proc/mlx/sx_netdev/sample_groups', 'w') as f:
        if (ingr_rate != 0):
            f.write('rx=%s\n' % nflog_base_group)
        else:
            f.write('rx=%s\n' % 0)

    for swp_intf in interfaces:
        with open('/cumulus/switchd/config/traffic/sflow/portsamprate_set', 'w') as f:
            f.write('%s: %d, %d\n' % (swp_intf, ingr_rate, egr_rate));


elif HAS_BCM_SHELL:

    if os.path.isfile('/proc/bcm/knet/sample_groups'):
        with open('/proc/bcm/knet/sample_groups', 'w') as f:
            if (ingr_rate != 0):
                f.write('rx=%s\n' % nflog_base_group)
                base_incr = 1
            else:
                f.write('rx=%s\n' % 0)

        with open('/proc/bcm/knet/sample_groups', 'w') as f:
            if egr_rate != 0 and nflog_base_group < 65535:
                f.write('tx=%d\n' % (int(nflog_base_group) + base_incr))
            else:
                f.write('tx=%s\n' % 0)
                if egr_rate != 0 and nflog_base_group >= 65535:
                    sys.stderr.write("nflog group for egress traffic out of range 0-65535.\n\n")
                    egr_rate = 0

        for swp_intf in interfaces:
            with open('/cumulus/switchd/config/traffic/sflow/portsamprate_set', 'w') as f:
                f.write('%s: %d, %d\n' % (swp_intf, ingr_rate, egr_rate));

    else:
        with open('/cumulus/switchd/ctrl/sample/ulog_channel', 'w') as f:
            f.write('%s\n' % nflog_base_group)

        # hardware specific code to enable sampling goes here
else:
    sys.exit("Platform is not supported.")
