#!/usr/bin/python3
#
# Copyright (c) 2017 Mellanox Technologies. All rights reserved.
#
# This Software is licensed under one of the following licenses:
#
# 1) under the terms of the "Common Public License 1.0" a copy of which is
#    available from the Open Source Initiative, see
#    http://www.opensource.org/licenses/cpl.php.
#
# 2) under the terms of the "The BSD License" a copy of which is
#    available from the Open Source Initiative, see
#    http://www.opensource.org/licenses/bsd-license.php.
#
# 3) under the terms of the "GNU General Public License (GPL) Version 2" a
#    copy of which is available from the Open Source Initiative, see
#    http://www.opensource.org/licenses/gpl-license.php.
#
# Licensee has the right to choose one of the above licenses.
#
# Redistributions of source code must retain the above copyright
# notice and one of the license notices.
#
# Redistributions in binary form must reproduce both the above copyright
# notice, one of the license notices in the documentation
# and/or other materials provided with the distribution.
#

import sys
import os
import subprocess
import array

if os.path.exists('/usr/share/pyshared'):
	sys.path.append('/usr/share/pyshared')
if os.path.exists('/usr/lib/python2.7/dist-packages'):
	sys.path.append('/usr/lib/python2.7/dist-packages')
if os.path.exists('/usr/lib/python2.7/site-packages'):
	sys.path.append('/usr/lib/python2.7/site-packages')
from optparse import OptionParser
from dcbnetlink import DcbController
from collections import defaultdict
from subprocess import Popen, PIPE

DCB_CAP_DCBX_HOST = 0x1
DCB_CAP_DCBX_LLD_MANAGED = 0x2
DCB_CAP_DCBX_VER_CEE = 0x4
DCB_CAP_DCBX_VER_IEEE = 0x8
DCB_CAP_DCBX_STATIC = 0x10

IEEE_8021QAZ_TSA_STRICT = 0
IEEE_8021QAZ_TSA_CB_SHAPER = 1
IEEE_8021QAZ_TSA_ETS = 2
IEEE_8021QAZ_TSA_VENDOR = 255

IEEE_8021QAZ_APP_SEL_ETHERTYPE	= 1
IEEE_8021QAZ_APP_SEL_STREAM	= 2
IEEE_8021QAZ_APP_SEL_DGRAM	= 3
IEEE_8021QAZ_APP_SEL_ANY	= 4
IEEE_8021QAZ_APP_SEL_DSCP       = 5

parser = OptionParser(usage="%prog -i <interface> [options]", version="%prog 1.2")

parser.add_option("-f", "--pfc", dest="pfc",
		  help="Set priority flow control for each priority. LIST is " +
			"comma separated value for each priority starting from 0 to 7. " +
			"Example: 0,0,0,0,1,1,1,1 enable PFC on TC4-7",
		  metavar="LIST")
parser.add_option("-p", "--prio_tc", dest="prio_tc",
		  help="maps UPs to TCs. LIST is 8 comma separated TC numbers. " +
			"Example: 0,0,0,0,1,1,1,1 maps UPs 0-3 to TC0, and UPs 4-7 to " +
			"TC1",
		  metavar="LIST")
parser.add_option("-s", "--tsa", dest="tsa",
		  help="Transmission algorithm for " +
			"each TC. LIST is comma separated algorithm names for each TC. " +
			"Possible algorithms: strict, ets and vendor. " +
			"Example: vendor,strict,ets,ets,ets,ets,ets,ets sets " +
			"TC0 to vendor, TC1 to strict, TC2-7 to ets.",
		  metavar="LIST")
parser.add_option("-t", "--tcbw", dest="tc_bw",
		  help="Set minimal guaranteed %BW for ETS TCs. LIST is comma " +
			"separated percents for each TC. Values set to TCs that are " +
			"not configured to ETS algorithm must be zero. " +
			"Example: if TC0,TC2 are set to ETS, then 10,0,90,0,0,0,0,0 " +
			"will set TC0 to 10% and TC2 to 90%. Percents must sum to " +
			"100.",
		  metavar="LIST")
parser.add_option("-r", "--ratelimit", dest="ratelimit",
		  help="Rate limit for TCs (in Gbps). LIST is a comma separated " +
			"Gbps limit for each TC. Example: 1,8,8 will limit TC0 to " +
			"1Gbps, and TC1,TC2 to 8 Gbps each.",
		  metavar="LIST")
parser.add_option("-d", "--dcbx", dest="dcbx",
		  help="set dcbx mode to firmware controlled(fw) or " +
		       "OS controlled(os). Note, when in OS mode, mlnx_qos should not be used " +
		       "in parallel with other dcbx tools such as lldptool")
parser.add_option("--trust", dest="trust",
                  help="set priority trust state to pcp or dscp")
parser.add_option("--dscp2prio", dest="dscp2prio",
		  action="append",
		  help="set/del a (dscp,prio) mapping, or 'flush' to delete all. Example 'set,30,2' maps dscp 30 to priority 2. " +
		       "'del,30,2' deletes the mapping. Deleting last dscp or using 'flush' will automatically change trust to pcp.")
parser.add_option("--defltprio", dest="defltprio",
                  help="set/del default priority. Example 'set,2' sets default priority to 2. " +
                       "'del,2' deletes this setting.")
parser.add_option("--cable_len", dest="cable_len",
                  help="set cable_len for buffer's xoff and xon thresholds")
parser.add_option("--prio2buffer", dest="prio2buffer",
		  help="maps priority to receive buffer. " +
			"Example: 0,2,5,7,1,2,3,6 maps priorities 0,1,2,3,4,5,6,7 to receive buffer " +
			"0,2,5,7,1,2,3,6",
		  metavar="LIST")
parser.add_option("--buffer_size", dest="buffer_size",
		  help="Set receive buffer size LIST is comma separated percents for each buffer." +
			"For pfc enabled buffer, the buffer size must be larger than the xoff_threshold. " +
			"Example: 87296,87296,0,87296,0,0,0,0 sets receive buffer size for buffer " +
			"0,1,2,3,4,5,6,7 respectively" ,
		  metavar="LIST")
parser.add_option("-i", "--interface", dest="intf",
		  help="Interface name")

parser.add_option("-a", action="store_true", dest="printall", default=False,
		  help="Show all interface's TCs")

(options, args) = parser.parse_args()

if len(args) > 0:
	print("Bad arguments")
	parser.print_usage()
	sys.exit(1)

if (options.intf == None):
	print("Interface name is required")
	parser.print_usage()

	sys.exit(1)

ratelimit_path = "/sys/class/net/" + options.intf + "/qos/maxrate"
trust_path = "/sys/class/net/" + options.intf + "/qos/trust"
dscp2prio_path = "/sys/class/net/" + options.intf + "/qos/dscp2prio"
buffer_size_path = "/sys/class/net/" + options.intf + "/qos/buffer_size"
prio2buffer_path = "/sys/class/net/" + options.intf + "/qos/prio2buffer"

is_trust_sysfs = os.path.exists(trust_path);
is_buffer_sysfs = os.path.exists(buffer_size_path);


if (not is_trust_sysfs):
	from dcbnetlink import DcbApp, DcbAppTable

class Maxrate:
	def get(self):
		pass
	def set(self, ratelimit):
		pass

	def prepare(self, ratelimit):
		old_ratelimit = self.get()
		ratelimit += old_ratelimit[len(ratelimit):8]
		return ratelimit

class MaxrateNL(Maxrate):
	def __init__(self, ctrl):
		self.ctrl = ctrl
	def get(self):
		return ctrl.get_ieee_maxrate()

	def set(self, ratelimit):
		ratelimit = self.prepare(ratelimit)
		ctrl.set_ieee_maxrate(ratelimit)

class MaxrateSysfs(Maxrate):
	def __init__(self, path):
		self.path = path

	def get(self):
		ratelimit = []
		f = open(self.path, "r")
		for item in f.read().split():
			ratelimit.append(float(item))
		f.close()

		return ratelimit

	def set(self, ratelimit):
		ratelimit = self.prepare(ratelimit)
		f = open(self.path, "w")
		f.write(" ".join(str(r) for r in ratelimit))
		f.close()

class Trust:
	def __init__(self):
		self.trust = "none"

	def getTrust(self, ctrl):
		if (is_trust_sysfs):
			f = open(trust_path, "r")
			self.trust = f.read().replace('\n', '')
			f.close()
		else:
			appTable = ctrl.get_ieee_app_table()
			if appTable.countAppSelector(IEEE_8021QAZ_APP_SEL_DSCP) == 0:
				self.trust = "pcp"
			else:
				self.trust = "dscp"

	def setTrust(self, ctrl, optionTrust):
		if self.trust == optionTrust:
			return
		if (is_trust_sysfs):
			f = open(trust_path, "w")
			f.write(optionTrust)
			f.close()
		else:
			appTable = ctrl.get_ieee_app_table()
			if optionTrust == "pcp":
				appTable.delAppEntry(ctrl, IEEE_8021QAZ_APP_SEL_DSCP)
			elif optionTrust == "dscp":
				appTable.setDefaultAppEntry(ctrl, IEEE_8021QAZ_APP_SEL_DSCP, 64)

def get_prio2buffer_sysfs():
	f = open(prio2buffer_path, "r")
	buffer = array.array('B', [0,0,0,0,0,0,0,0])
	next(f)
	for line in f:
		data = line.split()
		prio = int(data[0])
		buf = int(data[1])
		buffer[prio] = buf
	f.close()
	return buffer

def set_prio2buffer_sysfs(buffer):
	f = open(prio2buffer_path, "w")
	msg = ""
	for prio in range(8):
		msg += "%1d," % (buffer[prio])
	f.write("%s" % (msg))
	f.close()

def get_buffer_size_sysfs():
	f = open(buffer_size_path, "r")
	buffer = array.array('I', [0,0,0,0,0,0,0,0,0])
	first_line =  f.readline()
	data = first_line.split('=')
	total_size = int(data[1])
	for line in f:
		if not line[0].isdigit():
			continue
		data = line.split()
		buf = int(data[0])
		size = int(data[1])
		buffer[buf] = size
	f.close()
	buffer[8] = total_size
	return buffer

def set_buffer_size_sysfs(buffer):
	f = open(buffer_size_path, "w")
	msg = ""
	for buf in range(8):
		msg += "%1d," % (buffer[buf])
	f.write("%s" % (msg))
	f.close()

def print_dscp2prio_sysfs():
	f = open(dscp2prio_path, "r")
	s = ["","","","","","","",""]

	next(f)
	for line in f:
		data = line.split()
		dscp = int(data[0])
		prio = int(data[1])
		s[prio] += '%02d,' % dscp

	for i in range(8):
		temp = ""
		pad = "\tprio:%d dscp:" % i
		while (len(s[i]) > 24):
			temp += pad + s[i][:24] + "\n"
			s[i] = s[i][24:]
		if s[i] != "":
			temp += pad + s[i]
		if temp != "":
			print(temp)
	f.close()

def pretty_print(prio_tc, tsa, tcbw, ratelimit, pfc_en, trust, pfc_delay, buffer_size, tot_size, prio2buffer):
	if (ctrl.get_dcbx() & DCB_CAP_DCBX_HOST):
		print("DCBX mode: OS controlled")
	else:
		print("DCBX mode: Firmware controlled")

	print ("Priority trust state: " + trust)
	if trust == "dscp":
		print("dscp2prio mapping:")
		if (is_trust_sysfs):
			print_dscp2prio_sysfs()
		else:
			appTable.printAppSelector(IEEE_8021QAZ_APP_SEL_DSCP)
	print("default priority:")
	appTable.printDefaultPriority()

	if (not buffer_error):
		msg = "Receive buffer size (bytes): "
		for buf in range(8):
			msg += "%1d," % (buffer_size[buf])
		if tot_size is not None:
			msg += "max_buffer_size=%d" % tot_size
		print(msg)

	tc2up = defaultdict(list)

	if (printall == True):
		for i in range(8):
			tc2up.setdefault(i,[])

	print("Cable len: %d" % pfc_delay)

	print("PFC configuration:")
	print("\tpriority    0   1   2   3   4   5   6   7")
	msg = "\tenabled     "
	for up in range(8):
		msg += "%1d   " % ((pfc_en >> up) & 0x01)
	if (not buffer_error):
		msg += "\n\tbuffer      "
		for up in range(8):
			msg += "%1d   " % (prio2buffer[up])
	print(msg)

	for up in range(len(prio_tc)):
		tc = prio_tc[up]
		tc2up[int(tc)].append(up)

	for tc in sorted(tc2up):
		r = "unlimited"
		msg = ""
		try:
			if ratelimit[tc] > 0:
				r = "%.1f Gbps" % (float(ratelimit[tc] / 1000)/1000)
			msg = "tc: %d ratelimit: %s, tsa: " % (tc, r)
		except Exception as err:
			pass
		try:
			if (tsa[tc] == IEEE_8021QAZ_TSA_ETS):
				msg +="ets, bw: %s%%" % (tcbw[tc])
			elif (tsa[tc] == IEEE_8021QAZ_TSA_STRICT):
				msg += "strict"
			elif (tsa[tc] == IEEE_8021QAZ_TSA_VENDOR):
				msg += "vendor"
			else:
				msg += "unknown"
		except Exception as err:
			pass

		if msg:
			print(msg)

		try:
			for up in tc2up[tc]:
				print(("\t priority:  %s" % up))
		except Exception as err:
			pass

def parse_int(str, min, max, description):
	try:
		v = int(str)

		if (v < min or v > max):
			raise ValueError("%d is not in the range %d..%d" % (v, min, max))

		return v
	except ValueError as e:
		print(("Bad value for %s: %s" % (description, e)))
		parser.print_usage()
		sys.exit(1)


pfc_en = 0
pfc_delay = 0
tsa = [IEEE_8021QAZ_TSA_STRICT, IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT,IEEE_8021QAZ_TSA_STRICT]
tc_bw = [0, 0, 0, 0, 0, 0, 0, 0]
prio_tc = [0, 0, 0, 0, 0, 0, 0, 0]
printall = False

res = subprocess.Popen(["ps", "-e"], stdout=subprocess.PIPE)
output = res.communicate()[0]
if b'lldpad' in output:
	print('****** WARNING: lldpad service is running and may overwrite your settings ******\n')

ctrl = DcbController(options.intf)

# ********* dcbx mode command **************************************
if (options.dcbx != None):
	if (options.dcbx == "os"):
		ctrl.set_dcbx(ctrl.get_dcbx() | DCB_CAP_DCBX_HOST);
	elif (options.dcbx == "fw"):
		ctrl.set_dcbx(0);
	elif (options.dcbx != "get"):
		print ("Invalid dcbx mode command. Refer to the help.")
		sys.exit(1)

# ********* pfc command ******************************
if options.cable_len:
	pfc_delay = parse_int(options.cable_len, 0, 0xffff, "cable_len")

if options.pfc:
	i = 0
	pfc_en = 0

	for t in options.pfc.split(","):
		if i >= 8:
			print("Too many items for PFC")
			sys.exit(1)

		temp = parse_int(t, 0, 1, "PFC")
		pfc_en |= (temp << i)

		i += 1

	if i != 8:
		print("pfc list must have 8 items")
		sys.exit(1)

if options.cable_len or options.pfc:
	try:
		ctrl.set_ieee_pfc(_pfc_en = pfc_en, _delay = pfc_delay)
	except OSError as e:
		print(e)
		sys.exit(1)

# ********* dscp2prio command ******************************
dscp2prio = []
if options.dscp2prio and options.trust:
	print ("trust and dscp2prio commands cannot be used at the same time.")
	sys.exit(1)

try:
	if options.dscp2prio:
		for opt in options.dscp2prio:
			if opt == "flush":
				dscp2prio.append((opt, None, None))
				continue

			action, dscp, prio = opt.split(",")
			dscp = int(dscp)
			prio = int(prio)

			if ((action != "set") and (action != "del")) or (dscp > 63) or (prio > 7):
				sys.exit(1)
			dscp2prio.append((action, dscp, prio))

except:
	print("Invalid dscp2prio command. Refer to the help.")
	sys.exit(1)

try:
	for action, dscp, prio in dscp2prio:
		if action == "set":
			if (is_trust_sysfs):
				f = open(dscp2prio_path, "w")
				f.write("%d %d " % (dscp, prio))
				f.close()
			else:
				ctrl.set_ieee_app(IEEE_8021QAZ_APP_SEL_DSCP,prio,dscp)
		elif action == "del":
			if (not is_trust_sysfs):
				ctrl.del_ieee_app(IEEE_8021QAZ_APP_SEL_DSCP,prio,dscp)
		elif action == "flush":
			if (not is_trust_sysfs):
				appTable = ctrl.get_ieee_app_table()
				for app in appTable.apps.values():
					if app.selector != IEEE_8021QAZ_APP_SEL_DSCP:
						continue
					ctrl.del_ieee_app(app.selector, app.priority, app.protocol)

	if (not is_trust_sysfs):
		appTable = ctrl.get_ieee_app_table()

except:
	if options.dscp2prio:
		print("dscp2prio command failed")
		sys.exit(1)
	else:
		if (not is_trust_sysfs):
			appTable = DcbAppTable()

# ********* defltprio command ******************************
try:
	if options.defltprio:
		action, prio = options.defltprio.split(",")
		prio = int(prio)

		if action not in ("set", "del") or prio < 0 or prio > 7:
			raise Exception()
except:
	print("Invalid defltprio command. Refer to the help.")
	sys.exit(1)

try:
	if options.defltprio:
		if action == "set":
			ctrl.set_ieee_app(IEEE_8021QAZ_APP_SEL_ETHERTYPE,prio,0)
		elif action == "del":
			ctrl.del_ieee_app(IEEE_8021QAZ_APP_SEL_ETHERTYPE,prio,0)

	appTable = ctrl.get_ieee_app_table()
except:
	if options.defltprio:
		print("defltprio command failed")
		sys.exit(1)
	else:
		appTable = DcbAppTable()

# ********* trust command ******************************
if options.trust:
	if (options.trust != "dscp") and (options.trust != "pcp"):
		print("Invalid trust state command. Refer to the help.")
		sys.exit(1)
try:
	trustObj = Trust()
	trustObj.getTrust(ctrl)

	if options.trust:
		trustObj.setTrust(ctrl, options.trust)
		trustObj.getTrust(ctrl)
		if (not is_trust_sysfs):
			appTable = ctrl.get_ieee_app_table()

except:
	print ("Priority trust state is not supported on your system")
	if options.trust:
		sys.exit(1)
	else:
		trustObj.trust = "none"

# ********* buffer_size and prio2buffer command ************
buffer_error = 0

try:
	if (is_buffer_sysfs):
		prio2buffer = get_prio2buffer_sysfs()
		buffer_size = get_buffer_size_sysfs()
		tot_size = buffer_size[8]
	else:
		prio2buffer, buffer_size, tot_size = ctrl.get_ieee_dcb_buffer()

except:
	buffer_error = 1
	buffer_size = [0, 0, 0, 0, 0, 0, 0, 0]
	prio2buffer = [0, 0, 0, 0, 0, 0, 0, 0]
	tot_size = None
	print ("Buffers commands are not supported on your system")
	if options.buffer_size or options.prio2buffer:
		sys.exit(1)

try:
	if options.prio2buffer:
		i = 0

		for t in options.prio2buffer.split(","):
			if i >= 8:
				print("Too many items for prio2buffer")
				sys.exit(1)

			prio2buffer[i] = parse_int(t, 0, 7, "prio buffer")
			i += 1

	if options.buffer_size:
		i = 0

		for t in options.buffer_size.split(","):
			if i >= 8:
				print("Too many items for buffer_size")
				sys.exit(1)

			buffer_size[i] = int(t)
			i += 1

	if options.buffer_size or options.prio2buffer:
		if (is_buffer_sysfs):
			if (options.buffer_size):
				set_buffer_size_sysfs(buffer_size)
				buffer_size = get_buffer_size_sysfs()
				tot_size = buffer_size[8]
			else:
				set_prio2buffer_sysfs(prio2buffer)
				prio2buffer = get_prio2buffer_sysfs()
		else:
			ctrl.set_dcb_buffer(_prio2buffer = prio2buffer, _buffer_size = buffer_size,
					    _tot_size = tot_size)
			prio2buffer, buffer_size, tot_size = ctrl.get_ieee_dcb_buffer()
except:
	buffer_error = 1
	if options.buffer_size or options.prio2buffer:
		print ("Buffers commands error")
		sys.exit(1)

# ********* ratelimit command ******************************
try:
	ratelimit = []
	maxrate = None
	if (not os.path.exists(ratelimit_path)):
		maxrate = MaxrateNL(ctrl)
	else:
		maxrate = MaxrateSysfs(ratelimit_path)

	if options.ratelimit:
		i = 0
		for r in options.ratelimit.split(","):
			if i >=8:
				print ("Too many items for ratelimit")
				sys.exit(1)

			r = parse_int(r, 0, 1000000, "ratelimit")
	
			ratelimit += [r * 1000 * 1000]
			i += 1

		if i != 8:
			print("ratelimit must have 8 items")
			sys.exit(1)

		try:
			maxrate.set(ratelimit)
		except:
			print("Rate limit is not supported on your system!")

	try:
		ratelimit = maxrate.get()
	except:
		print("Rate limit is not supported on your system!")
except:
	if options.ratelimit:
		sys.exit(1)
	else:
		ratelimit = []

try:
	prio_tc, tsa, tc_bw = ctrl.get_ieee_ets()
	pfc_en = ctrl.get_ieee_pfc_en()
	pfc_delay = ctrl.get_ieee_pfc_delay()

except:
	print("ETS features are not supported on your system")
	sys.exit(1)

if options.printall:
	printall = True

# ********* ets (tsa, tc_bw, prio_tc) command ******************************
if (options.tsa):
	i = 0
	for t in options.tsa.split(","):
		if i >= 8:
			print("Too many items for TSA")
			sys.exit(1)

		if (t == "strict"):
			tsa[i] = IEEE_8021QAZ_TSA_STRICT
		elif (t == 'ets'):
			tsa[i] = IEEE_8021QAZ_TSA_ETS
		elif (t == 'vendor'):
			tsa[i] = IEEE_8021QAZ_TSA_VENDOR
		else:
			print(("Bad TSA value: ", t))
			parser.print_usage()
			sys.exit(1)
		i += 1

	if i != 8:
		print("tsa list must have 8 items")
		sys.exit(1)

if options.tc_bw:
	i = 0
	for t in options.tc_bw.split(","):
		if i >= 8:
			print("Too many items for ETS BW")
			sys.exit(1)

		bw = parse_int(t, 0, 100, "ETS BW")

		if tsa[i] != IEEE_8021QAZ_TSA_ETS and bw != 0:
			print("ETS BW for a strict/vendor TC must be 0")
			parser.print_usage()
			sys.exit(1)

		tc_bw[i] = bw
		i += 1

	if i != 8:
		print("tcbw list must have 8 items")
		sys.exit(1)

if options.prio_tc:
	i = 0
	for t in options.prio_tc.split(","):
		if i >= 8:
			print("Too many items in UP => TC mapping")
			sys.exit(1)

		prio_tc[i] = parse_int(t, 0, 7, "UP => TC mapping")
		i += 1

	if i != 8:
		print("prio_tc list must have 8 items")
		sys.exit(1)

if options.tsa or options.tc_bw or options.prio_tc:
	try:
		ctrl.set_ieee_ets(_prio_tc = prio_tc, _tsa = tsa, _tc_bw = tc_bw)
	except OSError as e:
		print(e)
		sys.exit(1)

pretty_print(prio_tc, tsa, tc_bw, ratelimit, pfc_en, trustObj.trust, pfc_delay, buffer_size, tot_size, prio2buffer)