path: root/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_streams.py

#!/router/bin/python

from trex_stl_exceptions import *
from trex_stl_packet_builder_interface import CTrexPktBuilderInterface
from trex_stl_packet_builder_scapy import CScapyTRexPktBuilder, Ether, IP, RawPcapReader
from collections import OrderedDict, namedtuple

from dpkt import pcap
import random
import yaml
import base64
import string
import traceback

def random_name (l):
    return ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(l))


# base class for TX mode
class STLTXMode(object):
    def __init__ (self):
        self.fields = {}

    def to_json (self):
        return self.fields


# continuous mode
class STLTXCont(STLTXMode):

    def __init__ (self, pps = 1):

        if not isinstance(pps, (int, float)):
            raise STLArgumentError('pps', pps)

        super(STLTXCont, self).__init__()

        self.fields['type'] = 'continuous'
        self.fields['pps']  = pps


# single burst mode
class STLTXSingleBurst(STLTXMode):

    def __init__ (self, pps = 1, total_pkts = 1):

        if not isinstance(pps, (int, float)):
            raise STLArgumentError('pps', pps)

        if not isinstance(total_pkts, int):
            raise STLArgumentError('total_pkts', total_pkts)

        super(STLTXSingleBurst, self).__init__()

        self.fields['type'] = 'single_burst'
        self.fields['pps']  = pps
        self.fields['total_pkts'] = total_pkts


# multi burst mode
class STLTXMultiBurst(STLTXMode):

    def __init__ (self,
                  pps = 1,
                  pkts_per_burst = 1,
                  ibg = 0.0,
                  count = 1):

        if not isinstance(pps, (int, float)):
            raise STLArgumentError('pps', pps)

        if not isinstance(pkts_per_burst, int):
            raise STLArgumentError('pkts_per_burst', pkts_per_burst)

        if not isinstance(ibg, (int, float)):
            raise STLArgumentError('ibg', ibg)

        if not isinstance(count, int):
            raise STLArgumentError('count', count)

        super(STLTXMultiBurst, self).__init__()

        self.fields['type'] = 'multi_burst'
        self.fields['pps'] = pps
        self.fields['pkts_per_burst'] = pkts_per_burst
        self.fields['ibg'] = ibg
        self.fields['count'] = count


class STLStream(object):

    def __init__ (self,
                  name = None,
                  packet = None,
                  mode = STLTXCont(1),
                  enabled = True,
                  self_start = True,
                  isg = 0.0,
                  rx_stats = None,
                  next = None,
                  stream_id = None):

        # type checking
        if not isinstance(mode, STLTXMode):
            raise STLArgumentError('mode', mode)

        if packet and not isinstance(packet, CTrexPktBuilderInterface):
            raise STLArgumentError('packet', packet)

        if not isinstance(enabled, bool):
            raise STLArgumentError('enabled', enabled)

        if not isinstance(self_start, bool):
            raise STLArgumentError('self_start', self_start)

        if not isinstance(isg, (int, float)):
            raise STLArgumentError('isg', isg)

        if (type(mode) == STLTXCont) and (next != None):
            raise STLError("continuous stream cannot have a next stream ID")

        # tag for the stream and next - can be anything
        self.name = name if name else random_name(8)
        self.next = next
        self.set_id(stream_id)

        self.fields = {}

        # basic fields
        self.fields['enabled'] = enabled
        self.fields['self_start'] = self_start
        self.fields['isg'] = isg

        # mode
        self.fields['mode'] = mode.to_json()

        self.fields['packet'] = {}
        self.fields['vm'] = {}

        if not packet:
            packet = CScapyTRexPktBuilder(pkt = Ether()/IP())

        # packet builder
        packet.compile()

        # packet and VM
        self.fields['packet'] = packet.dump_pkt()
        self.fields['vm']     = packet.get_vm_data()

        if not rx_stats:
            self.fields['rx_stats'] = {}
            self.fields['rx_stats']['enabled'] = False
        else:
            self.fields['rx_stats'] = rx_stats


    def __str__ (self):
        s =  "Stream Name: {0}\n".format(self.name)
        s += "Stream Next: {0}\n".format(self.next)
        s += "Stream JSON:\n{0}\n".format(json.dumps(self.fields, indent = 4, separators=(',', ': '), sort_keys = True))
        return s

    def to_json (self):
        return self.fields

    def get_id (self):
        return self.id

    def set_id (self, id):
        self.id = id

    def get_name (self):
        return self.name

    def get_next (self):
        return self.next


    def to_yaml (self):
        return {'name': self.name, 'stream': self.fields}
  
    def dump_to_yaml (self, yaml_file = None):
        yaml_dump = yaml.dump([self.to_yaml()], default_flow_style = False)

        # write to file if provided
        if yaml_file:
            with open(yaml_file, 'w') as f:
                f.write(yaml_dump)

        return yaml_dump

class YAMLLoader(object):

    def __init__ (self, yaml_file):
        self.yaml_path = os.path.dirname(yaml_file)
        self.yaml_file = yaml_file


    def __parse_packet (self, packet_dict):

        packet_type = set(packet_dict).intersection(['binary', 'pcap'])
        if len(packet_type) != 1:
            raise STLError("packet section must contain either 'binary' or 'pcap'")

        if 'binary' in packet_type:
            try:
                pkt_str = base64.b64decode(packet_dict['binary'])
            except TypeError:
                raise STLError("'binary' field is not a valid packet format")

            builder = CScapyTRexPktBuilder(pkt_buffer = pkt_str)

        elif 'pcap' in packet_type:
            pcap = os.path.join(self.yaml_path, packet_dict['pcap'])

            if not os.path.exists(pcap):
                raise STLError("'pcap' - cannot find '{0}'".format(pcap))

            builder = CScapyTRexPktBuilder(pkt = pcap)

        return builder


    def __parse_mode (self, mode_obj):

        mode_type = mode_obj.get('type')

        if mode_type == 'continuous':
            defaults = STLTXCont()
            mode = STLTXCont(pps = mode_obj.get('pps', defaults.fields['pps']))

        elif mode_type == 'single_burst':
            defaults = STLTXSingleBurst()
            mode = STLTXSingleBurst(pps         = mode_obj.get('pps', defaults.fields['pps']),
                                    total_pkts  = mode_obj.get('total_pkts', defaults.fields['total_pkts']))

        elif mode_type == 'multi_burst':
            defaults = STLTXMultiBurst()
            mode = STLTXMultiBurst(pps            = mode_obj.get('pps', defaults.fields['pps']),
                                   pkts_per_burst = mode_obj.get('pkts_per_burst', defaults.fields['pkts_per_burst']),
                                   ibg            = mode_obj.get('ibg', defaults.fields['ibg']),
                                   count          = mode_obj.get('count', defaults.fields['count']))

        else:
            raise STLError("mode type can be 'continuous', 'single_burst' or 'multi_burst")


        return mode


    def __parse_stream (self, yaml_object):
        s_obj = yaml_object['stream']

        # parse packet
        packet = s_obj.get('packet')
        if not packet:
            raise STLError("YAML file must contain 'packet' field")

        builder = self.__parse_packet(packet)


        # mode
        mode_obj = s_obj.get('mode')
        if not mode_obj:
            raise STLError("YAML file must contain 'mode' field")

        mode = self.__parse_mode(mode_obj)

        
        defaults = STLStream()

        # create the stream
        stream = STLStream(name       = yaml_object.get('name'),
                           packet     = builder,
                           mode       = mode,
                           enabled    = s_obj.get('enabled', defaults.fields['enabled']),
                           self_start = s_obj.get('self_start', defaults.fields['self_start']),
                           isg        = s_obj.get('isg', defaults.fields['isg']),
                           rx_stats   = s_obj.get('rx_stats', defaults.fields['rx_stats']),
                           next       = yaml_object.get('next'))

        # hack the VM fields for now
        if 'vm' in s_obj:
            stream.fields['vm'].update(s_obj['vm'])

        return stream


    def parse (self):
        with open(self.yaml_file, 'r') as f:
            # read YAML and pass it down to stream object
            yaml_str = f.read()

            try:
                objects = yaml.load(yaml_str)
            except yaml.parser.ParserError as e:
                raise STLError(str(e))

            streams = [self.__parse_stream(object) for object in objects]
            
            return streams


# profile class
class STLProfile(object):
    def __init__ (self, streams = None):
        if streams == None:
            streams = []

        if not type(streams) == list:
            streams = [streams]

        if not all([isinstance(stream, STLStream) for stream in streams]):
            raise STLArgumentError('streams', streams, valid_values = STLStream)

        self.streams = streams


    def get_streams (self):
        return self.streams

    def __str__ (self):
        return '\n'.join([str(stream) for stream in self.streams])


    @staticmethod
    def load_yaml (yaml_file):
        # check filename
        if not os.path.isfile(yaml_file):
            raise STLError("file '{0}' does not exists".format(yaml_file))

        yaml_loader = YAMLLoader(yaml_file)
        streams = yaml_loader.parse()

        return STLProfile(streams)


    @staticmethod
    def load_py (python_file):
        # check filename
        if not os.path.isfile(python_file):
            raise STLError("file '{0}' does not exists".format(python_file))

        basedir = os.path.dirname(python_file)
        sys.path.append(basedir)

        try:
            file    = os.path.basename(python_file).split('.')[0]
            module = __import__(file, globals(), locals(), [], -1)
            reload(module) # reload the update 

            streams = module.register().get_streams()

            return STLProfile(streams)

        except Exception as e:
            a, b, tb = sys.exc_info()
            x =''.join(traceback.format_list(traceback.extract_tb(tb)[1:])) + a.__name__ + ": " + str(b) + "\n"

            summary = "\nPython Traceback follows:\n\n" + x
            raise STLError(summary)


        finally:
            sys.path.remove(basedir)

    @staticmethod
    def load_pcap (pcap_file):
        # check filename
        if not os.path.isfile(pcap_file):
            raise STLError("file '{0}' does not exists".format(pcap_file))

        streams = []
        last_ts_usec = 0

        pkts = RawPcapReader(pcap_file).read_all()
        
        for i, (cap, meta) in enumerate(pkts, start = 1):
            ts_usec = meta[0] * 1e6 + meta[1]
            streams.append(STLStream(name = i,
                                     packet = CScapyTRexPktBuilder(pkt_buffer = cap),
                                     mode = STLTXSingleBurst(total_pkts = 1),
                                     self_start = True if (i == 1) else False,
                                     isg = (ts_usec - last_ts_usec),  # seconds to usec
                                     next = (i + 1) if i != len(pkts) else None))
            last_ts_usec = ts_usec

        return STLProfile(streams)

      

    @staticmethod
    def load (filename):
        x = os.path.basename(filename).split('.')
        suffix = x[1] if (len(x) == 2) else None

        if suffix == 'py':
            profile = STLProfile.load_py(filename)

        elif suffix == 'yaml':
            profile = STLProfile.load_yaml(filename)

        elif suffix in ['cap', 'pcap']:
            profile = STLProfile.load_pcap(filename)

        else:
            raise STLError("unknown profile file type: '{0}'".format(suffix))

        return profile


    def dump_to_yaml (self, yaml_file = None):
        yaml_list = [stream.to_yaml() for stream in self.streams]
        yaml_str = yaml.dump(yaml_list, default_flow_style = False)

        # write to file if provided
        if yaml_file:
            with open(yaml_file, 'w') as f:
                f.write(yaml_str)

        return yaml_str