#!/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 bcmshell
import os.path

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_bcm_ports()
else:
    try:
        interfaces = [ pt.linux2bcm(intf) ]
    except:
        sys.stderr.write("Interface %s couldn't be found.\n\n" % intf);
        usage()

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

    bcm = bcmshell.bcmshell(keepopen=True)
    for i in interfaces:
        bcm.run("portsamprate %s %s %s" % (i, ingr_rate, egr_rate))
    bcm.close()

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

    # hardware specific code to enable sampling goes here
