#!/router/bin/python import trex_root_path from client_utils.packet_builder import CTRexPktBuilder from trex_stateless_client import STLClient from common.trex_streams import * from client_utils.general_utils import id_count_gen import dpkt class CTRexHltApi(object): def __init__(self): self.trex_client = None self.connected = False # self._stream_db = CStreamList() self._port_data = {} # ----- session functions ----- # # sync = RPC, async = ZMQ def connect(self, device, port_list, username, sync_port = 4501, async_port = 4500, reset=False, break_locks=False): ret_dict = {"status": 0} self.trex_client = STLClient(username, device, sync_port, async_port) rc = self.trex_client.connect() if rc.bad(): self.trex_client = None ret_dict.update({"log": rc.err()}) return ret_dict # arrived here, connection successfully created with server # next, try acquiring ports of TRex port_list = self.parse_port_list(port_list) response = self.trex_client.acquire(port_list, force=break_locks) res_ok, log = CTRexHltApi.process_response(port_list, response) if not res_ok: self.trex_client.disconnect() self.trex_client = None ret_dict.update({"log": log}) # TODO: should revert taken ports?! return ret_dict # arrived here, all desired ports were successfully acquired print log if reset: # remove all port traffic configuration from TRex response = self.trex_client.remove_all_streams(port_list) res_ok, log = CTRexHltApi.process_response(port_list, response) if not res_ok: self.trex_client.disconnect() self.trex_client = None ret_dict.update({"log": log}) return ret_dict print log port_handle = {device: {port: port # since only supporting single TRex at the moment, 1:1 map for port in port_list} } ret_dict.update({"status": 1, "log": None, "port_handle": port_handle, "offline": 0}) # ports are online self.connected = True self._port_data_init(port_list) return ret_dict def cleanup_session(self, port_list, maintain_lock=False): ret_dict = {"status": 0} if not maintain_lock: # release taken ports if port_list == "all": port_list = self.trex_client.get_acquired_ports() else: port_list = self.parse_port_list(port_list) response = self.trex_client.release(port_list) res_ok, log = CTRexHltApi.process_response(port_list, response) if not res_ok: ret_dict.update({"log": log}) return ret_dict ret_dict.update({"status": 1, "log": None}) self.trex_client.disconnect() self.trex_client = None self.connected = False return ret_dict def interface_config(self, port_handle, mode="config"): ALLOWED_MODES = ["config", "modify", "destroy"] if mode not in ALLOWED_MODES: raise ValueError("mode must be one of the following values: {modes}".format(modes=ALLOWED_MODES)) # pass this function for now... return {"status": 1, "log": None} # ----- traffic functions ----- # def traffic_config(self, mode, port_list, l2_encap="ethernet_ii", mac_src="00:00:01:00:00:01", mac_dst="00:00:00:00:00:00", l3_protocol="ipv4", ip_src_addr="0.0.0.0", ip_dst_addr="192.0.0.1", l3_length=110, transmit_mode="continuous", rate_pps=100, **kwargs): if type(port_list) is not list(): port_list = [port_list] ALLOWED_MODES = ["create", "modify", "remove", "enable", "disable", "reset"] if mode not in ALLOWED_MODES: raise ValueError("mode must be one of the following values: {modes}".format(modes=ALLOWED_MODES)) if mode == "create": # create a new stream with desired attributes, starting by creating packet try: packet = CTRexHltApi.generate_stream(l2_encap, mac_src, mac_dst, l3_protocol, ip_src_addr, ip_dst_addr, l3_length) # set transmission attributes tx_mode = CTxMode(transmit_mode, rate_pps, **kwargs) # set rx_stats rx_stats = CRxStats() # defaults with disabled # join the generated data into stream stream_obj = CStream() stream_obj_params = {"enabled": True, "self_start": True, "next_stream_id": -1, "isg": 0.0, "mode": tx_mode, "rx_stats": rx_stats, "packet": packet} # vm is excluded from this list since CTRexPktBuilder obj is passed stream_obj.load_data(**stream_obj_params) except Exception as e: # some exception happened during the stream creation return {"status": 0, "log": str(e)} # try adding the stream per port, until free stream_id is found for port_id in port_list: port_data = self._port_data.get(port_id) id_candidate = None # TODO: change this to better implementation while True: id_candidate = port_data["stream_id_gen"].next() response = self.trex_client.add_stream(stream_id=id_candidate, stream_obj=stream_obj.dump(), port_id_list=port_id) res_ok, log = CTRexHltApi.process_response(port_id, response) if res_ok: # found non-taken stream_id on server # save it for modifying needs port_data["streams"].update({id_candidate: stream_obj}) break else: print log # proceed to another iteration to use another id print 'need another iteration?' continue return {"status": 1, "stream_id": id_candidate, "log": None} else: raise NotImplementedError("mode '{0}' is not supported yet on TRex".format(mode)) def traffic_control(self, action, port_handle, **kwargs): ALLOWED_ACTIONS = ["clear_stats", "run", "stop", "sync_run"] if action not in ALLOWED_ACTIONS: raise ValueError("action must be one of the following values: {actions}".format(actions=ALLOWED_ACTIONS)) # ret_dict = {"status": 0, "stopped": 1} port_list = self.parse_port_list(port_handle) if type(port_list) is not list(): port_list = [port_list] if action == "run": if not set(kwargs.keys()) >= set(['mul', 'duration']): raise ValueError("For 'run' action should be specified mul and duration arguments") response = self.trex_client.start_traffic(kwargs['mul'], kwargs['duration'], port_id_list=port_list) res_ok, log = CTRexHltApi.process_response(port_list, response) if res_ok: return {"status": 1, "stopped": 0, "log": None} else: print log elif action == "stop": response = self.trex_client.stop_traffic(port_id_list=port_list) res_ok, log = CTRexHltApi.process_response(port_list, response) if res_ok: return {"status": 1, "stopped": 1, "log": None} else: raise NotImplementedError("action '{0}' is not supported yet on TRex".format(action)) # if we arrived here, this means that operation FAILED! return {"status": 0, "log": log} def traffic_stats(self, port_handle, mode): ALLOWED_MODES = ["aggregate", "streams", "all"] if mode not in ALLOWED_MODES: raise ValueError("mode must be one of the following values: {modes}".format(modes=ALLOWED_MODES)) # pass this function for now... if mode == "aggregate": # create a new stream with desired attributes, starting by creating packet try: packet = CTRexHltApi.generate_stream(l2_encap, mac_src, mac_dst, l3_protocol, ip_src_addr, ip_dst_addr, l3_length) # set transmission attributes tx_mode = CTxMode(transmit_mode, rate_pps, **kwargs) # set rx_stats rx_stats = CRxStats() # defaults with disabled # join the generated data into stream stream_obj = CStream() stream_obj_params = {"enabled": True, "self_start": True, "next_stream_id": -1, "isg": 0.0, "mode": tx_mode, "rx_stats": rx_stats, "packet": packet} # vm is excluded from this list since CTRexPktBuilder obj is passed stream_obj.load_data(**stream_obj_params) except Exception as e: # some exception happened during the stream creation return {"status": 0, "log": str(e)} # try adding the stream, until free stream_id is found port_data = self._port_data.get(port_handle) id_candidate = None # TODO: change this to better implementation while True: id_candidate = port_data["stream_id_gen"].next() response = self.trex_client.add_stream(stream_id=id_candidate, stream_obj=stream_obj, port_id=port_handle) res_ok, log = CTRexHltApi.process_response(port_handle, response) if res_ok: # found non-taken stream_id on server # save it for modifying needs port_data["streams"].update({id_candidate: stream_obj}) break else: # proceed to another iteration to use another id continue return {"status": 1, "stream_id": id_candidate, "log": None} else: raise NotImplementedError("mode '{0}' is not supported yet on TRex".format(mode)) def get_aggregate_port_stats(self, port_handle): return self.traffic_stats(port_handle, mode="aggregate") def get_stream_stats(self, port_handle): return self.traffic_stats(port_handle, mode="streams") # ----- internal functions ----- # def _port_data_init(self, port_list): for port in port_list: self._port_data[port] = {"stream_id_gen": id_count_gen(), "streams": {}} @staticmethod def process_response(port_list, response): log = response.data() if response.good() else response.err() if isinstance(port_list, list): log = CTRexHltApi.join_batch_response(log) return response.good(), log @staticmethod def parse_port_list(port_list): if isinstance(port_list, str): return [int(port) for port in port_list.split()] elif isinstance(port_list, list): return [int(port) for port in port_list] else: return port_list @staticmethod def join_batch_response(responses): if type(responses) is list(): return "\n". join([str(response) for response in responses]) return responses @staticmethod def generate_stream(l2_encap, mac_src, mac_dst, l3_protocol, ip_src_addr, ip_dst_addr, l3_length): ALLOWED_L3_PROTOCOL = {"ipv4": dpkt.ethernet.ETH_TYPE_IP, "ipv6": dpkt.ethernet.ETH_TYPE_IP6, "arp": dpkt.ethernet.ETH_TYPE_ARP} ALLOWED_L4_PROTOCOL = {"tcp": dpkt.ip.IP_PROTO_TCP, "udp": dpkt.ip.IP_PROTO_UDP, "icmp": dpkt.ip.IP_PROTO_ICMP, "icmpv6": dpkt.ip.IP_PROTO_ICMP6, "igmp": dpkt.ip.IP_PROTO_IGMP, "rtp": dpkt.ip.IP_PROTO_IRTP, "isis": dpkt.ip.IP_PROTO_ISIS, "ospf": dpkt.ip.IP_PROTO_OSPF} pkt_bld = CTRexPktBuilder() if l2_encap == "ethernet_ii": pkt_bld.add_pkt_layer("l2", dpkt.ethernet.Ethernet()) # set Ethernet layer attributes pkt_bld.set_eth_layer_addr("l2", "src", mac_src) pkt_bld.set_eth_layer_addr("l2", "dst", mac_dst) else: raise NotImplementedError("l2_encap does not support the desired encapsulation '{0}'".format(l2_encap)) # set l3 on l2 if l3_protocol not in ALLOWED_L3_PROTOCOL: raise ValueError("l3_protocol must be one of the following: {0}".format(ALLOWED_L3_PROTOCOL)) pkt_bld.set_layer_attr("l2", "type", ALLOWED_L3_PROTOCOL[l3_protocol]) # set l3 attributes if l3_protocol == "ipv4": pkt_bld.add_pkt_layer("l3", dpkt.ip.IP()) pkt_bld.set_ip_layer_addr("l3", "src", ip_src_addr) pkt_bld.set_ip_layer_addr("l3", "dst", ip_dst_addr) pkt_bld.set_layer_attr("l3", "len", l3_length) else: raise NotImplementedError("l3_protocol '{0}' is not supported by TRex yet.".format(l3_protocol)) pkt_bld.dump_pkt_to_pcap("stream_test.pcap") return pkt_bld if __name__ == "__main__": pass