1#!/router/bin/python
2
3from .trex_stl_exceptions import *
4from .trex_stl_types import verify_exclusive_arg, validate_type
5from .trex_stl_packet_builder_interface import CTrexPktBuilderInterface
6from .trex_stl_packet_builder_scapy import *
7from collections import OrderedDict, namedtuple
8
9from scapy.utils import ltoa
10from scapy.error import Scapy_Exception
11import random
12import yaml
13import base64
14import string
15import traceback
16import copy
17import imp
18
19
20# base class for TX mode
21class STLTXMode(object):
22    """ mode rate speed """
23
24    def __init__ (self, pps = None, bps_L1 = None, bps_L2 = None, percentage = None):
25        """
26        Speed can be given in packets per second (pps), L2/L1 bps, or port percent
27        Use only one unit.
28        you can enter pps =10000 oe bps_L1=10
29
30        :parameters:
31            pps : float
32               Packets per second
33
34            bps_L1 : float
35               Bits per second L1 (with IPG)
36
37            bps_L2 : float
38               Bits per second L2 (Ethernet-FCS)
39
40            percentage : float
41               Link interface percent (0-100). Example: 10 is 10% of the port link setup
42
43        .. code-block:: python
44
45            # STLTXMode Example
46
47            mode = STLTXCont(pps = 10)
48
49            mode = STLTXCont(bps_L1 = 10000000) #10mbps L1
50
51            mode = STLTXCont(bps_L2 = 10000000) #10mbps L2
52
53            mode = STLTXCont(percentage = 10)   #10%
54
55        """
56
57        args = [pps, bps_L1, bps_L2, percentage]
58
59        # default
60        if all([x is None for x in args]):
61            pps = 1.0
62        else:
63            verify_exclusive_arg(args)
64
65        self.fields = {'rate': {}}
66
67        if pps is not None:
68            validate_type('pps', pps, [float, int])
69
70            self.fields['rate']['type']  = 'pps'
71            self.fields['rate']['value'] = pps
72
73        elif bps_L1 is not None:
74            validate_type('bps_L1', bps_L1, [float, int])
75
76            self.fields['rate']['type']  = 'bps_L1'
77            self.fields['rate']['value'] = bps_L1
78
79        elif bps_L2 is not None:
80            validate_type('bps_L2', bps_L2, [float, int])
81
82            self.fields['rate']['type']  = 'bps_L2'
83            self.fields['rate']['value'] = bps_L2
84
85        elif percentage is not None:
86            validate_type('percentage', percentage, [float, int])
87            if not (percentage > 0 and percentage <= 100):
88                raise STLArgumentError('percentage', percentage)
89
90            self.fields['rate']['type']  = 'percentage'
91            self.fields['rate']['value'] = percentage
92
93
94
95    def to_json (self):
96        return self.fields
97
98
99# continuous mode
100class STLTXCont(STLTXMode):
101    """ Continuous mode """
102
103    def __init__ (self, **kwargs):
104        """
105        Continuous mode
106
107         see :class:`trex_stl_lib.trex_stl_streams.STLTXMode` for rate
108
109        .. code-block:: python
110
111            # STLTXCont Example
112
113            mode = STLTXCont(pps = 10)
114
115        """
116        super(STLTXCont, self).__init__(**kwargs)
117
118
119        self.fields['type'] = 'continuous'
120
121    @staticmethod
122    def __str__ ():
123        return "Continuous"
124
125# single burst mode
126class STLTXSingleBurst(STLTXMode):
127    """ Single burst mode """
128
129    def __init__ (self, total_pkts = 1, **kwargs):
130        """
131        Single burst mode
132
133            :parameters:
134                 total_pkts : int
135                    Number of packets for this burst
136
137         see :class:`trex_stl_lib.trex_stl_streams.STLTXMode` for rate
138
139        .. code-block:: python
140
141            # STLTXSingleBurst Example
142
143            mode = STLTXSingleBurst( pps = 10, total_pkts = 1)
144
145        """
146
147
148        if not isinstance(total_pkts, int):
149            raise STLArgumentError('total_pkts', total_pkts)
150
151        super(STLTXSingleBurst, self).__init__(**kwargs)
152
153        self.fields['type'] = 'single_burst'
154        self.fields['total_pkts'] = total_pkts
155
156    @staticmethod
157    def __str__ ():
158        return "Single Burst"
159
160# multi burst mode
161class STLTXMultiBurst(STLTXMode):
162    """ Multi-burst mode """
163
164    def __init__ (self,
165                  pkts_per_burst = 1,
166                  ibg = 0.0,   # usec not SEC
167                  count = 1,
168                  **kwargs):
169        """
170        Multi-burst mode
171
172        :parameters:
173
174             pkts_per_burst: int
175                Number of packets per burst
176
177              ibg : float
178                Inter-burst gap in usec 1,000,000.0 is 1 sec
179
180              count : int
181                Number of bursts
182
183         see :class:`trex_stl_lib.trex_stl_streams.STLTXMode` for rate
184
185        .. code-block:: python
186
187            # STLTXMultiBurst Example
188
189            mode = STLTXMultiBurst(pps = 10, pkts_per_burst = 1,count 10, ibg=10.0)
190
191        """
192
193
194        if not isinstance(pkts_per_burst, int):
195            raise STLArgumentError('pkts_per_burst', pkts_per_burst)
196
197        if not isinstance(ibg, (int, float)):
198            raise STLArgumentError('ibg', ibg)
199
200        if not isinstance(count, int):
201            raise STLArgumentError('count', count)
202
203        super(STLTXMultiBurst, self).__init__(**kwargs)
204
205        self.fields['type'] = 'multi_burst'
206        self.fields['pkts_per_burst'] = pkts_per_burst
207        self.fields['ibg'] = ibg
208        self.fields['count'] = count
209
210    @staticmethod
211    def __str__ ():
212        return "Multi Burst"
213
214STLStreamDstMAC_CFG_FILE=0
215STLStreamDstMAC_PKT     =1
216STLStreamDstMAC_ARP     =2
217
218class STLFlowStatsInterface(object):
219    def __init__ (self, pg_id):
220        self.fields = {}
221        self.fields['enabled']         = True
222        self.fields['stream_id']       = pg_id
223
224    def to_json (self):
225        """ Dump as json"""
226        return dict(self.fields)
227
228    @staticmethod
229    def defaults ():
230        return {'enabled' : False}
231
232
233class STLFlowStats(STLFlowStatsInterface):
234    """ Define per stream basic stats
235
236    .. code-block:: python
237
238        # STLFlowStats Example
239
240        flow_stats = STLFlowStats(pg_id = 7)
241
242    """
243
244    def __init__(self, pg_id):
245        super(STLFlowStats, self).__init__(pg_id)
246        self.fields['rule_type'] = 'stats'
247
248
249class STLFlowLatencyStats(STLFlowStatsInterface):
250    """ Define per stream basic stats + latency, jitter, packet reorder/loss
251
252    .. code-block:: python
253
254        # STLFlowLatencyStats Example
255
256        flow_stats = STLFlowLatencyStats(pg_id = 7)
257
258    """
259
260    def __init__(self, pg_id):
261        super(STLFlowLatencyStats, self).__init__(pg_id)
262        self.fields['rule_type'] = 'latency'
263
264
265class STLStream(object):
266    """ One stream object. Includes mode, Field Engine mode packet template and Rx stats
267
268        .. code-block:: python
269
270            # STLStream Example
271
272
273            base_pkt =  Ether()/IP(src="16.0.0.1",dst="48.0.0.1")/UDP(dport=12,sport=1025)
274            pad = max(0, size - len(base_pkt)) * 'x'
275
276            STLStream( isg = 10.0, # star in delay
277                       name    ='S0',
278                       packet = STLPktBuilder(pkt = base_pkt/pad),
279                       mode = STLTXSingleBurst( pps = 10, total_pkts = 1),
280                       next = 'S1'), # point to next stream
281
282
283    """
284
285    def __init__ (self,
286                  name = None,
287                  packet = None,
288                  mode = STLTXCont(pps = 1),
289                  enabled = True,
290                  self_start = True,
291                  isg = 0.0,
292                  flow_stats = None,
293                  next = None,
294                  stream_id = None,
295                  action_count = 0,
296                  random_seed =0,
297                  mac_src_override_by_pkt=None,
298                  mac_dst_override_mode=None    #see  STLStreamDstMAC_xx
299                  ):
300        """
301        Stream object
302
303        :parameters:
304
305                  name  : string
306                       Name of the stream. Required if this stream is dependent on another stream, and another stream needs to refer to this stream by name.
307
308                  packet :  STLPktBuilder see :class:`trex_stl_lib.trex_stl_packet_builder_scapy.STLPktBuilder`
309                       Template packet and field engine program. Example: packet = STLPktBuilder(pkt = base_pkt/pad)
310
311                  mode :  :class:`trex_stl_lib.trex_stl_streams.STLTXCont` or :class:`trex_stl_lib.trex_stl_streams.STLTXSingleBurst`  or  :class:`trex_stl_lib.trex_stl_streams.STLTXMultiBurst`
312
313                  enabled : bool
314                      Indicates whether the stream is enabled.
315
316                  self_start : bool
317                      If False, another stream activates it.
318
319                  isg : float
320                     Inter-stream gap in usec. Time to wait until the stream sends the first packet.
321
322                  flow_stats : :class:`trex_stl_lib.trex_stl_streams.STLFlowStats`
323                      Per stream statistic object. See: STLFlowStats
324
325                  next : string
326                      Name of the stream to activate.
327
328                  stream_id :
329                        For use by HLTAPI.
330
331                  action_count : uint16_t
332                        If there is a next stream, number of loops before stopping. Default: 0 (unlimited).
333
334                  random_seed: uint16_t
335                       If given, the seed for this stream will be this value. Useful if you need a deterministic random value.
336
337                  mac_src_override_by_pkt : bool
338                        Template packet sets src MAC.
339
340                  mac_dst_override_mode=None : STLStreamDstMAC_xx
341                        Template packet sets dst MAC.
342        """
343
344
345        # type checking
346        validate_type('mode', mode, STLTXMode)
347        validate_type('packet', packet, (type(None), CTrexPktBuilderInterface))
348        validate_type('flow_stats', flow_stats, (type(None), STLFlowStatsInterface))
349        validate_type('enabled', enabled, bool)
350        validate_type('self_start', self_start, bool)
351        validate_type('isg', isg, (int, float))
352        validate_type('stream_id', stream_id, (type(None), int))
353        validate_type('random_seed',random_seed,int);
354
355        if (type(mode) == STLTXCont) and (next != None):
356            raise STLError("Continuous stream cannot have a next stream ID")
357
358        # tag for the stream and next - can be anything
359        self.name = name
360        self.next = next
361
362        self.mac_src_override_by_pkt = mac_src_override_by_pkt # save for easy construct code from stream object
363        self.mac_dst_override_mode = mac_dst_override_mode
364        self.id = stream_id
365
366
367        self.fields = {}
368
369        int_mac_src_override_by_pkt = 0;
370        int_mac_dst_override_mode   = 0;
371
372
373        if mac_src_override_by_pkt == None:
374            int_mac_src_override_by_pkt=0
375            if packet :
376                if packet.is_default_src_mac ()==False:
377                    int_mac_src_override_by_pkt=1
378
379        else:
380            int_mac_src_override_by_pkt = int(mac_src_override_by_pkt);
381
382        if mac_dst_override_mode == None:
383            int_mac_dst_override_mode   = 0;
384            if packet :
385                if packet.is_default_dst_mac ()==False:
386                    int_mac_dst_override_mode=STLStreamDstMAC_PKT
387        else:
388            int_mac_dst_override_mode = int(mac_dst_override_mode);
389
390
391        self.is_default_mac = not (int_mac_src_override_by_pkt or int_mac_dst_override_mode)
392
393        self.fields['flags'] = (int_mac_src_override_by_pkt&1) +  ((int_mac_dst_override_mode&3)<<1)
394
395        self.fields['action_count'] = action_count
396
397        # basic fields
398        self.fields['enabled'] = enabled
399        self.fields['self_start'] = self_start
400        self.fields['isg'] = isg
401
402        if random_seed !=0 :
403            self.fields['random_seed'] = random_seed # optional
404
405        # mode
406        self.fields['mode'] = mode.to_json()
407        self.mode_desc      = str(mode)
408
409
410        # packet
411        self.fields['packet'] = {}
412        self.fields['vm'] = {}
413
414        if not packet:
415            packet = STLPktBuilder(pkt = Ether()/IP())
416
417        self.scapy_pkt_builder = packet
418        # packet builder
419        packet.compile()
420
421        # packet and VM
422        self.fields['packet'] = packet.dump_pkt()
423        self.fields['vm']     = packet.get_vm_data()
424
425        self.pkt = base64.b64decode(self.fields['packet']['binary'])
426
427        # this is heavy, calculate lazy
428        self.packet_desc = None
429
430        if not flow_stats:
431            self.fields['flow_stats'] = STLFlowStats.defaults()
432        else:
433            self.fields['flow_stats'] = flow_stats.to_json()
434
435
436    def __str__ (self):
437        s =  "Stream Name: {0}\n".format(self.name)
438        s += "Stream Next: {0}\n".format(self.next)
439        s += "Stream JSON:\n{0}\n".format(json.dumps(self.fields, indent = 4, separators=(',', ': '), sort_keys = True))
440        return s
441
442    def to_json (self):
443        """
444        Return json format
445        """
446        return dict(self.fields)
447
448    def get_id (self):
449        """ Get the stream id after resolution  """
450        return self.id
451
452
453    def has_custom_mac_addr (self):
454        """ Return True if src or dst MAC were set as custom """
455        return not self.is_default_mac
456
457    def get_name (self):
458        """ Get the stream name """
459        return self.name
460
461    def get_next (self):
462        """ Get next stream object """
463        return self.next
464
465
466    def has_flow_stats (self):
467        """ Return True if stream was configured with flow stats """
468        return self.fields['flow_stats']['enabled']
469
470    def get_pkt (self):
471        """ Get packet as string """
472        return self.pkt
473
474    def get_pkt_len (self, count_crc = True):
475       """ Get packet number of bytes  """
476       pkt_len = len(self.get_pkt())
477       if count_crc:
478           pkt_len += 4
479
480       return pkt_len
481
482
483    def get_pkt_type (self):
484        """ Get packet description. Example: IP:UDP """
485        if self.packet_desc == None:
486            self.packet_desc = STLPktBuilder.pkt_layers_desc_from_buffer(self.get_pkt())
487
488        return self.packet_desc
489
490    def get_mode (self):
491        return self.mode_desc
492
493    @staticmethod
494    def get_rate_from_field (rate_json):
495        """ Get rate from json  """
496        t = rate_json['type']
497        v = rate_json['value']
498
499        if t == "pps":
500            return format_num(v, suffix = "pps")
501        elif t == "bps_L1":
502            return format_num(v, suffix = "bps (L1)")
503        elif t == "bps_L2":
504            return format_num(v, suffix = "bps (L2)")
505        elif t == "percentage":
506            return format_num(v, suffix = "%")
507
508    def get_rate (self):
509        return self.get_rate_from_field(self.fields['mode']['rate'])
510
511    def to_pkt_dump (self):
512        """ Print packet description from Scapy  """
513        if self.name:
514            print("Stream Name: ",self.name)
515        scapy_b = self.scapy_pkt_builder;
516        if scapy_b and isinstance(scapy_b,STLPktBuilder):
517            scapy_b.to_pkt_dump()
518        else:
519            print("Nothing to dump")
520
521
522
523    def to_yaml (self):
524        """ Convert to YAML  """
525        y = {}
526
527        if self.name:
528            y['name'] = self.name
529
530        if self.next:
531            y['next'] = self.next
532
533        y['stream'] = copy.deepcopy(self.fields)
534
535        # some shortcuts for YAML
536        rate_type  = self.fields['mode']['rate']['type']
537        rate_value = self.fields['mode']['rate']['value']
538
539        y['stream']['mode'][rate_type] = rate_value
540        del y['stream']['mode']['rate']
541
542        return y
543
544    # returns the Python code (text) to build this stream, inside the code it will be in variable "stream"
545    def to_code (self):
546        """ Convert to Python code as profile  """
547        packet = Ether(self.pkt)
548        layer = packet
549        imports_arr = []
550        # remove checksums, add imports if needed
551        while layer:
552            layer_class = layer.__class__.__name__
553            try: # check if class can be instantiated
554                eval('%s()' % layer_class)
555            except NameError: # no such layer
556                found_import = False
557                for module_path, module in sys.modules.items():
558                    import_string = 'from %s import %s' % (module_path, layer_class)
559                    if import_string in imports_arr:
560                        found_import = True
561                        break
562                    if not module_path.startswith(('scapy.layers', 'scapy.contrib')):
563                        continue
564                    check_layer = getattr(module, layer_class, None)
565                    if not check_layer:
566                        continue
567                    try:
568                        check_layer()
569                        imports_arr.append(import_string)
570                        found_import = True
571                        break
572                    except: # can't by instantiated
573                        continue
574                if not found_import:
575                    raise STLError('Could not determine import of layer %s' % layer.name)
576            for chksum_name in ('cksum', 'chksum'):
577                if chksum_name in layer.fields:
578                    del layer.fields[chksum_name]
579            layer = layer.payload
580        packet.hide_defaults()          # remove fields with default values
581        payload = packet.getlayer('Raw')
582        packet_command = packet.command()
583
584        imports = '\n'.join(imports_arr)
585        if payload:
586            payload.remove_payload() # fcs etc.
587            data = payload.fields.get('load', '')
588
589            good_printable = [c for c in string.printable if ord(c) not in range(32)]
590            good_printable.remove("'")
591
592            if type(data) is str:
593                new_data = ''.join([c if c in good_printable else r'\x{0:02x}'.format(ord(c)) for c in data])
594            else:
595                new_data = ''.join([chr(c) if chr(c) in good_printable else r'\x{0:02x}'.format(c) for c in data])
596
597            payload_start = packet_command.find("Raw(load=")
598            if payload_start != -1:
599                packet_command = packet_command[:payload_start-1]
600        layers = packet_command.split('/')
601
602        if payload:
603            if len(new_data) and new_data == new_data[0] * len(new_data):
604                layers.append("Raw(load='%s' * %s)" % (new_data[0], len(new_data)))
605            else:
606                layers.append("Raw(load='%s')" % new_data)
607
608        packet_code = 'packet = (' + (' / \n          ').join(layers) + ')'
609        vm_list = []
610        for inst in self.fields['vm']['instructions']:
611            if inst['type'] == 'flow_var':
612                vm_list.append("STLVmFlowVar(name='{name}', size={size}, op='{op}', init_value={init_value}, min_value={min_value}, max_value={max_value}, step={step})".format(**inst))
613            elif inst['type'] == 'write_flow_var':
614                vm_list.append("STLVmWrFlowVar(fv_name='{name}', pkt_offset={pkt_offset}, add_val={add_value}, is_big={is_big_endian})".format(**inst))
615            elif inst['type'] == 'write_mask_flow_var':
616                inst = copy.copy(inst)
617                inst['mask'] = hex(inst['mask'])
618                vm_list.append("STLVmWrMaskFlowVar(fv_name='{name}', pkt_offset={pkt_offset}, pkt_cast_size={pkt_cast_size}, mask={mask}, shift={shift}, add_value={add_value}, is_big={is_big_endian})".format(**inst))
619            elif inst['type'] == 'fix_checksum_ipv4':
620                vm_list.append("STLVmFixIpv4(offset={pkt_offset})".format(**inst))
621            elif inst['type'] == 'trim_pkt_size':
622                vm_list.append("STLVmTrimPktSize(fv_name='{name}')".format(**inst))
623            elif inst['type'] == 'tuple_flow_var':
624                inst = copy.copy(inst)
625                inst['ip_min'] = ltoa(inst['ip_min'])
626                inst['ip_max'] = ltoa(inst['ip_max'])
627                vm_list.append("STLVmTupleGen(name='{name}', ip_min='{ip_min}', ip_max='{ip_max}', port_min={port_min}, port_max={port_max}, limit_flows={limit_flows}, flags={flags})".format(**inst))
628            elif inst['type'] == 'flow_var_rand_limit':
629                vm_list.append("STLVmFlowVarRepetableRandom(name='{name}', size={size}, limit={limit}, seed={seed}, min_value={min_value}, max_value={max_value})".format(**inst))
630
631        vm_code = 'vm = STLScVmRaw([' + ',\n                 '.join(vm_list) + '], split_by_field = %s)' % STLStream.__add_quotes(self.fields['vm'].get('split_by_var'))
632        stream_params_list = []
633        stream_params_list.append('packet = STLPktBuilder(pkt = packet, vm = vm)')
634        if default_STLStream.name != self.name:
635            stream_params_list.append('name = %s' % STLStream.__add_quotes(self.name))
636        if default_STLStream.fields['enabled'] != self.fields['enabled']:
637            stream_params_list.append('enabled = %s' % self.fields['enabled'])
638        if default_STLStream.fields['self_start'] != self.fields['self_start']:
639            stream_params_list.append('self_start = %s' % self.fields['self_start'])
640        if default_STLStream.fields['isg'] != self.fields['isg']:
641            stream_params_list.append('isg = %s' % self.fields['isg'])
642        if default_STLStream.fields['flow_stats'] != self.fields['flow_stats']:
643            stream_params_list.append('flow_stats = STLFlowStats(%s)' % self.fields['flow_stats']['stream_id'])
644        if default_STLStream.next != self.next:
645            stream_params_list.append('next = %s' % STLStream.__add_quotes(self.next))
646        if default_STLStream.id != self.id:
647            stream_params_list.append('stream_id = %s' % self.id)
648        if default_STLStream.fields['action_count'] != self.fields['action_count']:
649            stream_params_list.append('action_count = %s' % self.fields['action_count'])
650        if 'random_seed' in self.fields:
651            stream_params_list.append('random_seed = %s' % self.fields.get('random_seed', 0))
652        if default_STLStream.mac_src_override_by_pkt != self.mac_src_override_by_pkt:
653            stream_params_list.append('mac_src_override_by_pkt = %s' % self.mac_src_override_by_pkt)
654        if default_STLStream.mac_dst_override_mode != self.mac_dst_override_mode:
655            stream_params_list.append('mac_dst_override_mode = %s' % self.mac_dst_override_mode)
656
657        mode_args = ''
658        for key, value in self.fields['mode'].items():
659            if key not in ('rate', 'type'):
660                mode_args += '%s = %s, ' % (key, value)
661        mode_args += '%s = %s' % (self.fields['mode']['rate']['type'], self.fields['mode']['rate']['value'])
662        if self.mode_desc == STLTXCont.__str__():
663            stream_params_list.append('mode = STLTXCont(%s)' % mode_args)
664        elif self.mode_desc == STLTXSingleBurst().__str__():
665            stream_params_list.append('mode = STLTXSingleBurst(%s)' % mode_args)
666        elif self.mode_desc == STLTXMultiBurst().__str__():
667            stream_params_list.append('mode = STLTXMultiBurst(%s)' % mode_args)
668        else:
669            raise STLError('Could not determine mode: %s' % self.mode_desc)
670
671        stream = "stream = STLStream(" + ',\n                   '.join(stream_params_list) + ')'
672        return '\n'.join([imports, packet_code, vm_code, stream])
673
674    # add quoted for string, or leave as is if other type
675    @staticmethod
676    def __add_quotes(arg):
677        if type(arg) is str:
678            return "'%s'" % arg
679        return arg
680
681    # used to replace non-printable characters with hex
682    @staticmethod
683    def __replchars_to_hex(match):
684        return r'\x{0:02x}'.format(ord(match.group()))
685
686    def dump_to_yaml (self, yaml_file = None):
687        """ Print as yaml  """
688        yaml_dump = yaml.dump([self.to_yaml()], default_flow_style = False)
689
690        # write to file if provided
691        if yaml_file:
692            with open(yaml_file, 'w') as f:
693                f.write(yaml_dump)
694
695        return yaml_dump
696
697class YAMLLoader(object):
698
699    def __init__ (self, yaml_file):
700        self.yaml_path = os.path.dirname(yaml_file)
701        self.yaml_file = yaml_file
702
703
704    def __parse_packet (self, packet_dict):
705
706        packet_type = set(packet_dict).intersection(['binary', 'pcap'])
707        if len(packet_type) != 1:
708            raise STLError("Packet section must contain either 'binary' or 'pcap'")
709
710        if 'binary' in packet_type:
711            try:
712                pkt_str = base64.b64decode(packet_dict['binary'])
713            except TypeError:
714                raise STLError("'binary' field is not a valid packet format")
715
716            builder = STLPktBuilder(pkt_buffer = pkt_str)
717
718        elif 'pcap' in packet_type:
719            pcap = os.path.join(self.yaml_path, packet_dict['pcap'])
720
721            if not os.path.exists(pcap):
722                raise STLError("'pcap' - cannot find '{0}'".format(pcap))
723
724            builder = STLPktBuilder(pkt = pcap)
725
726        return builder
727
728
729    def __parse_mode (self, mode_obj):
730        if not mode_obj:
731            return None
732
733        rate_parser = set(mode_obj).intersection(['pps', 'bps_L1', 'bps_L2', 'percentage'])
734        if len(rate_parser) != 1:
735            raise STLError("'rate' must contain exactly one from 'pps', 'bps_L1', 'bps_L2', 'percentage'")
736
737        rate_type  = rate_parser.pop()
738        rate = {rate_type : mode_obj[rate_type]}
739
740        mode_type = mode_obj.get('type')
741
742        if mode_type == 'continuous':
743            mode = STLTXCont(**rate)
744
745        elif mode_type == 'single_burst':
746            defaults = STLTXSingleBurst()
747            mode = STLTXSingleBurst(total_pkts  = mode_obj.get('total_pkts', defaults.fields['total_pkts']),
748                                    **rate)
749
750        elif mode_type == 'multi_burst':
751            defaults = STLTXMultiBurst()
752            mode = STLTXMultiBurst(pkts_per_burst = mode_obj.get('pkts_per_burst', defaults.fields['pkts_per_burst']),
753                                   ibg            = mode_obj.get('ibg', defaults.fields['ibg']),
754                                   count          = mode_obj.get('count', defaults.fields['count']),
755                                   **rate)
756
757        else:
758            raise STLError("mode type can be 'continuous', 'single_burst' or 'multi_burst")
759
760
761        return mode
762
763
764
765    def __parse_flow_stats (self, flow_stats_obj):
766
767        # no such object
768        if not flow_stats_obj or flow_stats_obj.get('enabled') == False:
769            return None
770
771        pg_id = flow_stats_obj.get('stream_id')
772        if pg_id == None:
773            raise STLError("Enabled RX stats section must contain 'stream_id' field")
774
775        return STLFlowStats(pg_id = pg_id)
776
777
778    def __parse_stream (self, yaml_object):
779        s_obj = yaml_object['stream']
780
781        # parse packet
782        packet = s_obj.get('packet')
783        if not packet:
784            raise STLError("YAML file must contain 'packet' field")
785
786        builder = self.__parse_packet(packet)
787
788
789        # mode
790        mode = self.__parse_mode(s_obj.get('mode'))
791
792        # rx stats
793        flow_stats = self.__parse_flow_stats(s_obj.get('flow_stats'))
794
795
796        defaults = default_STLStream
797        # create the stream
798        stream = STLStream(name       = yaml_object.get('name'),
799                           packet     = builder,
800                           mode       = mode,
801                           flow_stats   = flow_stats,
802                           enabled    = s_obj.get('enabled', defaults.fields['enabled']),
803                           self_start = s_obj.get('self_start', defaults.fields['self_start']),
804                           isg        = s_obj.get('isg', defaults.fields['isg']),
805                           next       = yaml_object.get('next'),
806                           action_count = s_obj.get('action_count', defaults.fields['action_count']),
807                           mac_src_override_by_pkt = s_obj.get('mac_src_override_by_pkt', 0),
808                           mac_dst_override_mode = s_obj.get('mac_src_override_by_pkt', 0)
809                           )
810
811        # hack the VM fields for now
812        if 'vm' in s_obj:
813            stream.fields['vm'].update(s_obj['vm'])
814
815        return stream
816
817
818    def parse (self):
819        with open(self.yaml_file, 'r') as f:
820            # read YAML and pass it down to stream object
821            yaml_str = f.read()
822
823            try:
824                objects = yaml.safe_load(yaml_str)
825            except yaml.parser.ParserError as e:
826                raise STLError(str(e))
827
828            streams = [self.__parse_stream(object) for object in objects]
829
830            return streams
831
832
833# profile class
834class STLProfile(object):
835    """ Describe a list of streams
836
837        .. code-block:: python
838
839            # STLProfile Example
840
841            profile =  STLProfile( [ STLStream( isg = 10.0, # star in delay
842                                        name    ='S0',
843                                        packet = STLPktBuilder(pkt = base_pkt/pad),
844                                        mode = STLTXSingleBurst( pps = 10, total_pkts = self.burst_size),
845                                        next = 'S1'), # point to next stream
846
847                             STLStream( self_start = False, # stream is  disabled enable trow S0
848                                        name    ='S1',
849                                        packet  = STLPktBuilder(pkt = base_pkt1/pad),
850                                        mode    = STLTXSingleBurst( pps = 10, total_pkts = self.burst_size),
851                                        next    = 'S2' ),
852
853                             STLStream(  self_start = False, # stream is  disabled enable trow S0
854                                         name   ='S2',
855                                         packet = STLPktBuilder(pkt = base_pkt2/pad),
856                                         mode = STLTXSingleBurst( pps = 10, total_pkts = self.burst_size )
857                                        )
858                            ]).get_streams()
859
860
861
862    """
863
864    def __init__ (self, streams = None):
865        """
866
867            :parameters:
868
869                  streams  : list of :class:`trex_stl_lib.trex_stl_streams.STLStream`
870                       a list of stream objects
871
872        """
873
874
875        if streams == None:
876            streams = []
877
878        if not type(streams) == list:
879            streams = [streams]
880
881        if not all([isinstance(stream, STLStream) for stream in streams]):
882            raise STLArgumentError('streams', streams, valid_values = STLStream)
883
884        self.streams = streams
885        self.meta = None
886
887
888    def get_streams (self):
889        """ Get the list of streams"""
890        return self.streams
891
892    def __str__ (self):
893        return '\n'.join([str(stream) for stream in self.streams])
894
895    def is_pauseable (self):
896        return all([x.get_mode() == "Continuous" for x in self.get_streams()])
897
898    def has_custom_mac_addr (self):
899        return any([x.has_custom_mac_addr() for x in self.get_streams()])
900
901    def has_flow_stats (self):
902        return any([x.has_flow_stats() for x in self.get_streams()])
903
904    @staticmethod
905    def load_yaml (yaml_file):
906        """ Load (from YAML file) a profile with a number of streams"""
907
908        # check filename
909        if not os.path.isfile(yaml_file):
910            raise STLError("file '{0}' does not exists".format(yaml_file))
911
912        yaml_loader = YAMLLoader(yaml_file)
913        streams = yaml_loader.parse()
914
915        profile = STLProfile(streams)
916        profile.meta = {'type': 'yaml'}
917
918        return profile
919
920    @staticmethod
921    def get_module_tunables(module):
922        # remove self and variables
923        func = module.register().get_streams
924        argc = func.__code__.co_argcount
925        tunables = func.__code__.co_varnames[1:argc]
926
927        # fetch defaults
928        defaults = func.__defaults__
929        if defaults is None:
930            return {}
931        if len(defaults) != (argc - 1):
932            raise STLError("Module should provide default values for all arguments on get_streams()")
933
934        output = {}
935        for t, d in zip(tunables, defaults):
936            output[t] = d
937
938        return output
939
940
941    @staticmethod
942    def load_py (python_file, direction = 0, port_id = 0, **kwargs):
943        """ Load from Python profile """
944
945        # check filename
946        if not os.path.isfile(python_file):
947            raise STLError("File '{0}' does not exist".format(python_file))
948
949        basedir = os.path.dirname(python_file)
950        sys.path.insert(0, basedir)
951
952        try:
953            file    = os.path.basename(python_file).split('.')[0]
954            module = __import__(file, globals(), locals(), [], 0)
955            imp.reload(module) # reload the update
956
957            t = STLProfile.get_module_tunables(module)
958            #for arg in kwargs:
959            #    if not arg in t:
960            #        raise STLError("Profile {0} does not support tunable '{1}' - supported tunables are: '{2}'".format(python_file, arg, t))
961
962            streams = module.register().get_streams(direction = direction,
963                                                    port_id = port_id,
964                                                    **kwargs)
965            profile = STLProfile(streams)
966
967            profile.meta = {'type': 'python',
968                            'tunables': t}
969
970            return profile
971
972        except Exception as e:
973            a, b, tb = sys.exc_info()
974            x =''.join(traceback.format_list(traceback.extract_tb(tb)[1:])) + a.__name__ + ": " + str(b) + "\n"
975
976            summary = "\nPython Traceback follows:\n\n" + x
977            raise STLError(summary)
978
979
980        finally:
981            sys.path.remove(basedir)
982
983
984    # loop_count = 0 means loop forever
985    @staticmethod
986    def load_pcap (pcap_file,
987                   ipg_usec = None,
988                   speedup = 1.0,
989                   loop_count = 1,
990                   vm = None,
991                   packet_hook = None,
992                   split_mode = None,
993                   min_ipg_usec = None):
994        """ Convert a pcap file with a number of packets to a list of connected streams.
995
996        packet1->packet2->packet3 etc
997
998                :parameters:
999
1000                  pcap_file  : string
1001                       Name of the pcap file
1002
1003                  ipg_usec   : float
1004                       Inter packet gap in usec. If IPG is None, IPG is taken from pcap file
1005
1006                  speedup   : float
1007                       When reading the pcap file, divide IPG by this "speedup" factor. Resulting IPG is sped up by this factor.
1008
1009                  loop_count : uint16_t
1010                       Number of loops to repeat the pcap file
1011
1012                  vm        :  list
1013                        List of Field engine instructions
1014
1015                  packet_hook : Callable or function
1016                        will be applied to every packet
1017
1018                  split_mode : str
1019                        should this PCAP be split to two profiles based on IPs / MACs
1020                        used for dual mode
1021                        can be 'MAC' or 'IP'
1022
1023                  min_ipg_usec   : float
1024                       Minumum inter packet gap in usec. Used to guard from too small IPGs.
1025
1026                 :return: STLProfile
1027
1028        """
1029
1030        # check filename
1031        if not os.path.isfile(pcap_file):
1032            raise STLError("file '{0}' does not exists".format(pcap_file))
1033        if speedup <= 0:
1034            raise STLError('Speedup should not be negative.')
1035        if min_ipg_usec and min_ipg_usec < 0:
1036            raise STLError('min_ipg_usec should not be negative.')
1037
1038
1039        # make sure IPG is not less than 0.001 usec
1040        if (ipg_usec is not None and (ipg_usec < 0.001 * speedup) and
1041                              (min_ipg_usec is None or min_ipg_usec < 0.001)):
1042            raise STLError("ipg_usec cannot be less than 0.001 usec: '{0}'".format(ipg_usec))
1043
1044        if loop_count < 0:
1045            raise STLError("'loop_count' cannot be negative")
1046
1047
1048        try:
1049
1050            if split_mode is None:
1051                pkts = PCAPReader(pcap_file).read_all()
1052                if len(pkts) == 0:
1053                    raise STLError("'{0}' does not contain any packets".format(pcap_file))
1054
1055                return STLProfile.__pkts_to_streams(pkts,
1056                                                    ipg_usec,
1057                                                    min_ipg_usec,
1058                                                    speedup,
1059                                                    loop_count,
1060                                                    vm,
1061                                                    packet_hook)
1062            else:
1063                pkts_a, pkts_b = PCAPReader(pcap_file).read_all(split_mode = split_mode)
1064                if (len(pkts_a) + len(pkts_b)) == 0:
1065                    raise STLError("'{0}' does not contain any packets".format(pcap_file))
1066
1067                # swap the packets if a is empty, or the ts of first packet in b is earlier
1068                if not pkts_a:
1069                    pkts_a, pkts_b = pkts_b, pkts_a
1070                elif (ipg_usec is None) and pkts_b:
1071                    meta = pkts_a[0][1]
1072                    start_time_a = meta[0] * 1e6 + meta[1]
1073                    meta = pkts_b[0][1]
1074                    start_time_b = meta[0] * 1e6 + meta[1]
1075                    if start_time_b < start_time_a:
1076                        pkts_a, pkts_b = pkts_b, pkts_a
1077
1078                profile_a = STLProfile.__pkts_to_streams(pkts_a,
1079                                                         ipg_usec,
1080                                                         min_ipg_usec,
1081                                                         speedup,
1082                                                         loop_count,
1083                                                         vm,
1084                                                         packet_hook,
1085                                                         start_delay_usec = 10000)
1086
1087                profile_b = STLProfile.__pkts_to_streams(pkts_b,
1088                                                         ipg_usec,
1089                                                         min_ipg_usec,
1090                                                         speedup,
1091                                                         loop_count,
1092                                                         vm,
1093                                                         packet_hook,
1094                                                         start_delay_usec = 10000)
1095
1096                return profile_a, profile_b
1097
1098
1099        except Scapy_Exception as e:
1100            raise STLError("failed to open PCAP file {0}: '{1}'".format(pcap_file, str(e)))
1101
1102
1103    @staticmethod
1104    def __pkts_to_streams (pkts, ipg_usec, min_ipg_usec, speedup, loop_count, vm, packet_hook, start_delay_usec = 0):
1105
1106        streams = []
1107        if packet_hook:
1108            pkts = [(packet_hook(cap), meta) for (cap, meta) in pkts]
1109
1110        for i, (cap, meta) in enumerate(pkts, start = 1):
1111            # IPG - if not provided, take from cap
1112            if ipg_usec is None:
1113                packet_time = meta[0] * 1e6 + meta[1]
1114                if i == 1:
1115                    prev_time = packet_time
1116                isg = (packet_time - prev_time) / float(speedup)
1117                if min_ipg_usec and isg < min_ipg_usec:
1118                    isg = min_ipg_usec
1119                prev_time = packet_time
1120            else: # user specified ipg
1121                if min_ipg_usec:
1122                    isg = min_ipg_usec
1123                else:
1124                    isg = ipg_usec / float(speedup)
1125
1126            # handle last packet
1127            if i == len(pkts):
1128                next = 1
1129                action_count = loop_count
1130            else:
1131                next = i + 1
1132                action_count = 0
1133
1134            streams.append(STLStream(name = i,
1135                                     packet = STLPktBuilder(pkt_buffer = cap, vm = vm),
1136                                     mode = STLTXSingleBurst(total_pkts = 1, percentage = 100),
1137                                     self_start = True if (i == 1) else False,
1138                                     isg = isg,  # usec
1139                                     action_count = action_count,
1140                                     next = next))
1141
1142
1143        profile = STLProfile(streams)
1144        profile.meta = {'type': 'pcap'}
1145
1146        return profile
1147
1148
1149
1150    @staticmethod
1151    def load (filename, direction = 0, port_id = 0, **kwargs):
1152        """ Load a profile by its type. Supported types are:
1153           * py
1154           * yaml
1155           * pcap file that converted to profile automaticly
1156
1157           :Parameters:
1158              filename  : string as filename
1159              direction : profile's direction (if supported by the profile)
1160              port_id   : which port ID this profile is being loaded to
1161              kwargs    : forward those key-value pairs to the profile
1162
1163        """
1164
1165        x = os.path.basename(filename).split('.')
1166        suffix = x[1] if (len(x) == 2) else None
1167
1168        if suffix == 'py':
1169            profile = STLProfile.load_py(filename, direction, port_id, **kwargs)
1170
1171        elif suffix == 'yaml':
1172            profile = STLProfile.load_yaml(filename)
1173
1174        elif suffix in ['cap', 'pcap']:
1175            profile = STLProfile.load_pcap(filename, speedup = 1, ipg_usec = 1e6)
1176
1177        else:
1178            raise STLError("unknown profile file type: '{0}'".format(suffix))
1179
1180        profile.meta['stream_count'] = len(profile.get_streams()) if isinstance(profile.get_streams(), list) else 1
1181        return profile
1182
1183    @staticmethod
1184    def get_info (filename):
1185        profile = STLProfile.load(filename)
1186        return profile.meta
1187
1188    def dump_as_pkt (self):
1189        """ Dump the profile as Scapy packet. If the packet is raw, convert it to Scapy before dumping it."""
1190        cnt=0;
1191        for stream in self.streams:
1192            print("=======================")
1193            print("Stream %d" % cnt)
1194            print("=======================")
1195            cnt = cnt +1
1196            stream.to_pkt_dump()
1197
1198    def dump_to_yaml (self, yaml_file = None):
1199        """ Convert the profile to yaml """
1200        yaml_list = [stream.to_yaml() for stream in self.streams]
1201        yaml_str = yaml.dump(yaml_list, default_flow_style = False)
1202
1203        # write to file if provided
1204        if yaml_file:
1205            with open(yaml_file, 'w') as f:
1206                f.write(yaml_str)
1207
1208        return yaml_str
1209
1210    def dump_to_code (self, profile_file = None):
1211        """ Convert the profile to Python native profile. """
1212        profile_dump = '''# !!! Auto-generated code !!!
1213from trex_stl_lib.api import *
1214
1215class STLS1(object):
1216    def get_streams(self, direction = 0, **kwargs):
1217        streams = []
1218'''
1219        for stream in self.streams:
1220            profile_dump += ' '*8 + stream.to_code().replace('\n', '\n' + ' '*8) + '\n'
1221            profile_dump += ' '*8 + 'streams.append(stream)\n'
1222        profile_dump += '''
1223        return streams
1224
1225def register():
1226    return STLS1()
1227'''
1228        # write to file if provided
1229        if profile_file:
1230            with open(profile_file, 'w') as f:
1231                f.write(profile_dump)
1232
1233        return profile_dump
1234
1235
1236
1237    def __len__ (self):
1238        return len(self.streams)
1239
1240
1241class PCAPReader(object):
1242    def __init__ (self, pcap_file):
1243        self.pcap_file = pcap_file
1244
1245    def read_all (self, split_mode = None):
1246        if split_mode is None:
1247            return RawPcapReader(self.pcap_file).read_all()
1248
1249        # we need to split
1250        self.pcap = rdpcap(self.pcap_file)
1251        self.graph = Graph()
1252
1253        self.pkt_groups = [ [], [] ]
1254
1255        if split_mode == 'MAC':
1256            self.generate_mac_groups()
1257        elif split_mode == 'IP':
1258            self.generate_ip_groups()
1259        else:
1260            raise STLError('unknown split mode for PCAP')
1261
1262        return self.pkt_groups
1263
1264
1265    # generate two groups based on MACs
1266    def generate_mac_groups (self):
1267        for i, pkt in enumerate(self.pcap):
1268            if not isinstance(pkt, (Ether, Dot3) ):
1269                raise STLError("Packet #{0} has an unknown L2 format: {1}".format(i, type(pkt)))
1270            mac_src = pkt.fields['src']
1271            mac_dst = pkt.fields['dst']
1272            self.graph.add(mac_src, mac_dst)
1273
1274        # split the graph to two groups
1275        mac_groups = self.graph.split()
1276
1277        for pkt in self.pcap:
1278            mac_src = pkt.fields['src']
1279            group = 1 if mac_src in mac_groups[1] else 0
1280
1281            time, raw = pkt.time, bytes(pkt)
1282            self.pkt_groups[group].append((raw, (time, 0)))
1283
1284
1285    # generate two groups based on IPs
1286    def generate_ip_groups (self):
1287        for pkt in self.pcap:
1288            if not isinstance(pkt, (Ether, Dot3) ):
1289                raise STLError("Packet #{0} has an unknown L2 format: {1}".format(i, type(pkt)))
1290            # skip non IP packets
1291            if not isinstance(pkt.payload, IP):
1292                continue
1293            ip_src = pkt.payload.fields['src']
1294            ip_dst = pkt.payload.fields['dst']
1295            self.graph.add(ip_src, ip_dst)
1296
1297        # split the graph to two groups
1298        ip_groups = self.graph.split()
1299
1300        for pkt in self.pcap:
1301            # default group - 0
1302            group = 0
1303
1304            # if the packet is IP and IP SRC is in group 1 - move to group 1
1305            if isinstance(pkt.payload, IP) and pkt.payload.fields['src'] in ip_groups[1]:
1306                group = 1
1307
1308            time, raw = pkt.time, bytes(pkt)
1309            self.pkt_groups[group].append((raw, (time, 0)))
1310
1311
1312
1313# a simple graph object - used to split to two groups
1314class Graph(object):
1315    def __init__ (self):
1316        self.db = OrderedDict()
1317        self.debug = False
1318
1319    def log (self, msg):
1320        if self.debug:
1321            print(msg)
1322
1323    # add a connection v1 --> v2
1324    def add (self, v1, v2):
1325        # init value for v1
1326        if not v1 in self.db:
1327            self.db[v1] = set()
1328
1329        # init value for v2
1330        if not v2 in self.db:
1331            self.db[v2] = set()
1332
1333        # ignore self to self edges
1334        if v1 == v2:
1335            return
1336
1337        # undirected - add two ways
1338        self.db[v1].add(v2)
1339        self.db[v2].add(v1)
1340
1341
1342    # create a 2-color of the graph if possible
1343    def split (self):
1344        color_a = set()
1345        color_b = set()
1346
1347        # start with all
1348        nodes = list(self.db.keys())
1349
1350        # process one by one
1351        while len(nodes) > 0:
1352            node = nodes.pop(0)
1353
1354            friends = self.db[node]
1355
1356            # node has never been seen - move to color_a
1357            if not node in color_a and not node in color_b:
1358                self.log("<NEW> {0} --> A".format(node))
1359                color_a.add(node)
1360
1361            # node color
1362            node_color, other_color = (color_a, color_b) if node in color_a else (color_b, color_a)
1363
1364            # check that the coloring is possible
1365            bad_friends = friends.intersection(node_color)
1366            if bad_friends:
1367                raise STLError("ERROR: failed to split PCAP file - {0} and {1} are in the same group".format(node, bad_friends))
1368
1369            # add all the friends to the other color
1370            for friend in friends:
1371                self.log("<FRIEND> {0} --> {1}".format(friend, 'A' if other_color is color_a else 'B'))
1372                other_color.add(friend)
1373
1374
1375        return color_a, color_b
1376
1377
1378default_STLStream = STLStream()
1379
1380