#!/usr/bin/env python3
# Written by Jonatan Langlet for Direct Telemetry Access
# 用于向 Translator 注入合成的 DTA Report 数据包，但是速度太慢，不太适合基准测试，但比较适合测试转换能力

#from scapy.all import send, IP, ICMP
from scapy.all import *
import random
import sys
import binascii
import struct
import argparse
import time

parser = argparse.ArgumentParser(description='Inject a DTA Key-Write packet into the Tofino ASIC.')
parser.add_argument('operation',  type=str, nargs='+', choices=["keywrite", "append"], help='The DTA operation')
parser.add_argument('--data',  type=int, nargs='+', help='The telemetry data')
parser.add_argument('--key',  type=int, nargs='+', help='The telemetry key for KeyWrite operations')
parser.add_argument('--redundancy', type=int, nargs='+', help='The telemetry redundancy for KeyWrite operations')
parser.add_argument('--listID', type=int, nargs='+', help='The telemetry list ID for Append operations')
parser.add_argument('--loop', action='store_true', help='Indicates that the script should loop, generating traffic continuously')
parser.add_argument('--increment_data', action='store_true', help='Indicates that the data value should increment, if looping is enabled')
parser.add_argument('--increment_key', action='store_true', help='Indicates that the key-write key should increment, if looping is enabled')
parser.add_argument('--ipg', type=float, default=0.0, help='The IPG to replay traffic at, if emitting multiple packets simultaneously')
parser.add_argument('--batchsize', type=int, default=1, help='The batch size to use when --loop is enabled')

#args = vars(parser.parse_args())
args = parser.parse_args()
print(args)


class dta_base(Packet):
	name = "dtaBase"
	fields_desc = [ 
		XByteField("opcode",		0x01),
		XByteField("seqnum",		0),            # DTA sequence number
		BitField("immediate",		0,		1),
		BitField("retransmitable",	0,		1),
		BitField("reserved", 		0,		6)
	]

class dta_keyWrite(Packet):
	name = "dtaKeyWrite"
	fields_desc = [ 
		ByteField("redundancy",	0x02),
		# IntField("key",		0),
		IntField("key",		0),
		IntField("data", 	0)
	]

class dta_append(Packet):
	name = "dtaAppend"
	fields_desc = [ 
		IntField("listID",		0),
		IntField("data", 	0)
	]

def craft_dta_keywrite(key, data, redundancy):
	''' 使用 key, data, redundancy 创建一个 DTA Key-Write 数据包 '''

	print("Crafting a keywrite packet with key:%i data:%i, redundancy:%i" %(key,data,redundancy))
	
	key_bin = struct.pack(">I", key)
	
        pkt = Ether(dst="08:c0:eb:58:92:89")\
	/IP(src="192.168.1.91", dst="192.168.4.3")\
	/UDP(sport=40041, dport=40040)\
	/dta_base(opcode=0x01)\
	/dta_keyWrite(redundancy=redundancy, key=RawVal(key_bin), data=data)
	#/dta_keyWrite(redundancy=redundancy, key=key, data=data)
	
	return pkt


def craft_dta_append(listID, data):
	''' 使用 listID, data 创建一个 DTA Append 数据包 '''

	print("Crafting an append packet with listID:%i data:%i" %(listID,data))
	
    pkt = Ether(dst="08:c0:eb:58:92:89")\
	/IP(src="192.168.1.91", dst="192.168.4.3")\
	/UDP(sport=40041, dport=40040)\
	/dta_base(opcode=0x02)\
	/dta_append(listID=listID, data=data)
	
	return pkt


def emitPacket(pkts):
	''' 以 IPG(Inter Packet Gap)(如果定义了的话) 为时间间隔来发送传入的数据包 '''

	ipg = args.ipg
	print("Sending %i DTA packet(s) at ipg:%.3f" % (len(pkts), ipg))
	sendp(pkts, inter=ipg, iface="ma1")


# DTA Key-Write 数据包测试
if args.operation[0] == "keywrite":
	print("Crafting a DTA KeyWrite packet...")
	
	assert args.key, "No telemetry key specified!"
	assert args.redundancy, "No telemetry redundancy specified!"
	assert args.data, "No telemetry data specified!"
	
	# 首先传入 Key-Write 数据包制作所需的参数
	key = args.key[0]
	data = args.data[0]
	redundancy = args.redundancy[0]
	doLoop = args.loop
	increment_data = args.increment_data
	increment_key = args.increment_key
	
	# 如果启用循环的话，则持续递增，用于一次性的可靠性测试
	if doLoop:
		# 单个批次中元素的数量
		batchSize = args.batchsize
		print("Looping enabled with batchsize %i" % batchSize)
		
		# 依次制作 batchSize 个 Key-Write 数据包并一起发送
		pkts = []
		while True: 
			pkt = craft_dta_keywrite(key=key, data=data, redundancy=redundancy)
			
			pkts.append(pkt)
			
			if len(pkts) >= batchSize:
				emitPacket(pkts)
				pkts = []
				print("sleeping 1 sec before next batch...")
				time.sleep(1)
			
			# 对 key 和 data 进行依次递增操作
			if increment_key:
				key = key + 1
			if increment_data:
				data = data + 1
			print("Incremented key:%i and data:%i" %(key,data))
	# 这是默认功能，用于制作和发送单个数据包
	else:
		print("Looping disabled, sending single packet")
		pkt = craft_dta_keywrite(key=key, data=data, redundancy=redundancy)
		emitPacket(pkt)

# DTA Append 数据包测试
if args.operation[0] == "append":
	print("Crafting a DTA Append packet...")
	
	assert args.listID, "No telemetry list ID specified!"
	assert args.data, "No telemetry data specified!"

	# 首先传入 Append 数据包制作所需的参数
	doLoop = args.loop
	listID = args.listID[0]
	data = args.data[0]
	increment_data = args.increment_data
	
	# 如果启用循环的话，则持续递增，用于一次性的可靠性测试
	if doLoop:
		batchSize = args.batchsize
		print("Looping enabled with batchsize %i" %batchSize)
		
		pkts = []
		while True: 
			pkt = craft_dta_append(listID=listID, data=data)
			
			pkts.append(pkt)
			
			if len(pkts) >= batchSize:
				emitPacket(pkts)
				pkts = []
			
			if increment_data:
				data = data + 1
			print("Incremented data:%i" %(data))
	# 这是默认功能，用于制作和发送单个数据包
	else:
		pkt = craft_dta_append(listID=listID, data=data)
		emitPacket(pkt)