PCAP API - added packet hook

example is much simpler now

4 files changed, 69 insertions, 41 deletions

diff --git a/scripts/automation/trex_control_plane/stl/examples/stl_pcap.py b/scripts/automation/trex_control_plane/stl/examples/stl_pcap.py
index 45ddc24b..eae0f18b 100644
--- a/scripts/automation/trex_control_plane/stl/examples/stl_pcap.py
+++ b/scripts/automation/trex_control_plane/stl/examples/stl_pcap.py
@@ -3,62 +3,52 @@ from trex_stl_lib.api import *
 import argparse
 import sys
 
-def create_vm (ip_start, ip_end):
-     vm =[
 
-            # dest
-            STLVmFlowVar(name="dst", min_value = ip_start, max_value = ip_end, size = 4, op = "inc"),
-            STLVmWrFlowVar(fv_name="dst",pkt_offset= "IP.dst"),
+def packet_hook_generator (remove_fcs, vlan_id):
 
-            # checksum
-            STLVmFixIpv4(offset = "IP")
+    def packet_hook (packet):
+        packet = Ether(packet)
 
-            ]
-
-     return vm
-
-# warning: might make test slow
-def alter_streams(streams, remove_fcs, vlan_id):
-    for stream in streams:
-        packet = Ether(stream.pkt)
         if vlan_id >= 0 and vlan_id <= 4096:
             packet_l3 = packet.payload
             packet = Ether() / Dot1Q(vlan = vlan_id) / packet_l3
+
         if remove_fcs and packet.lastlayer().name == 'Padding':
             packet.lastlayer().underlayer.remove_payload()
-        packet = STLPktBuilder(packet)
-        stream.fields['packet'] = packet.dump_pkt()
-        stream.pkt = base64.b64decode(stream.fields['packet']['binary'])
+
+        return str(packet)
+
+    return packet_hook
+
 
 def inject_pcap (pcap_file, server, port, loop_count, ipg_usec, use_vm, remove_fcs, vlan_id):
 
     # create client
     c = STLClient(server = server)
-    
-    try:
-        if use_vm:
-            vm = create_vm("10.0.0.1", "10.0.0.254")
-        else:
-            vm = None
 
-        profile = STLProfile.load_pcap(pcap_file, ipg_usec = ipg_usec, loop_count = loop_count, vm = vm)
+    if remove_fcs or vlan_id:
+        packet_hook = packet_hook_generator(remove_fcs, vlan_id)
+    else:
+        packet_hook = None
 
-        print("Loaded pcap {0} with {1} packets...\n".format(pcap_file, len(profile)))
-        streams = profile.get_streams()
-        if remove_fcs or (vlan_id >= 0 and vlan_id <= 4096):
-            alter_streams(streams, remove_fcs, vlan_id)
+    try:
 
-        # uncomment this for simulator run
-        #STLSim().run(profile.get_streams(), outfile = '/auto/srg-sce-swinfra-usr/emb/users/ybrustin/out.pcap')
+        vm = STLIPRange(dst = {'start': '10.0.0.1', 'end': '10.0.0.254', 'step' : 1}) if use_vm else None
 
         c.connect()
         c.reset(ports = [port])
-        stream_ids = c.add_streams(streams, ports = [port])
 
         c.clear_stats()
+        d = c.push_pcap(pcap_file,
+                        ipg_usec = ipg_usec,
+                        count = loop_count,
+                        vm = vm,
+                        packet_hook = packet_hook)
+
+        STLSim().run(d, outfile = 'test.cap')
+
+        c.wait_on_traffic()
 
-        c.start()
-        c.wait_on_traffic(ports = [port])
 
         stats = c.get_stats()
         opackets = stats[port]['opackets']
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
index 7f6f3c0d..010d966c 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
@@ -1935,7 +1935,9 @@ class STLClient(object):
                    speedup = 1.0,
                    count = 1,
                    duration = -1,
-                   force = False):
+                   force = False,
+                   vm = None,
+                   packet_hook = None):
         """
             Push a local PCAP to the server
             This is equivalent to loading a PCAP file to a profile
@@ -1965,6 +1967,12 @@ class STLClient(object):
                 force: bool
                     Ignore file size limit - push any file size to the server
 
+                vm: list of VM instructions
+                    VM instructions to apply for every packet
+
+                packet_hook : Callable or function
+                    Will be applied to every packet
+
             :raises:
                 + :exc:`STLError`
 
@@ -1977,7 +1985,8 @@ class STLClient(object):
         validate_type('speedup',  speedup, (float, int))
         validate_type('count',  count, int)
         validate_type('duration', duration, (float, int))
-
+        validate_type('vm', vm, (list, type(None)))
+        
         # no support for > 1MB PCAP - use push remote
         if not force and os.path.getsize(pcap_filename) > (1024 * 1024):
             raise STLError("PCAP size of {:} is too big for local push - consider using remote push or provide 'force'".format(format_num(os.path.getsize(pcap_filename), suffix = 'B')))
@@ -1987,8 +1996,9 @@ class STLClient(object):
         profile = STLProfile.load_pcap(pcap_filename,
                                        ipg_usec,
                                        speedup,
-                                       count)
-
+                                       count,
+                                       vm = vm,
+                                       packet_hook = packet_hook)
 
         id_list = self.add_streams(profile.get_streams(), ports)
 
@@ -2077,7 +2087,6 @@ class STLClient(object):
         if not type(clear_global) is bool:
             raise STLArgumentError('clear_global', clear_global)
 
-
         rc = self.__clear_stats(ports, clear_global, clear_flow_stats)
         if not rc:
             raise STLError(rc)
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_packet_builder_scapy.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_packet_builder_scapy.py
index f8517a47..8c711563 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_packet_builder_scapy.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_packet_builder_scapy.py
@@ -1479,4 +1479,27 @@ class STLPktBuilder(CTrexPktBuilderInterface):
         pass;
 
 
+def STLIPRange (src = None,
+                dst = None,
+                fix_chksum = True):
+
+    vm = []
+
+    if src:
+        vm += [
+               STLVmFlowVar(name="src", min_value = src['start'], max_value = src['end'], size = 4, op = "inc", step = src['step']),
+               STLVmWrFlowVar(fv_name="src",pkt_offset= "IP.src")
+              ]
+
+    if dst:
+        vm += [
+                STLVmFlowVar(name="dst", min_value = dst['start'], max_value = dst['end'], size = 4, op = "inc", step = dst['step']),
+                STLVmWrFlowVar(fv_name="dst",pkt_offset= "IP.dst")
+              ]
+
+    if fix_chksum:
+        vm.append( STLVmFixIpv4(offset = "IP"))
+
+
+    return vm
 
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_streams.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_streams.py
index 92598312..a7fd3026 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_streams.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_streams.py
@@ -928,7 +928,7 @@ class STLProfile(object):
     
     # loop_count = 0 means loop forever
     @staticmethod
-    def load_pcap (pcap_file, ipg_usec = None, speedup = 1.0, loop_count = 1, vm = None):
+    def load_pcap (pcap_file, ipg_usec = None, speedup = 1.0, loop_count = 1, vm = None, packet_hook = None):
         """ Convert a pcap file with a number of packets to a list of connected streams.  
 
         packet1->packet2->packet3 etc 
@@ -950,6 +950,9 @@ class STLProfile(object):
                   vm        :  list 
                         List of Field engine instructions 
 
+                  packet_hook : Callable or function
+                        will be applied to every packet
+
                  :return: STLProfile
 
         """
@@ -973,7 +976,10 @@ class STLProfile(object):
         except Scapy_Exception as e:
             raise STLError("failed to open PCAP file '{0}'".format(pcap_file))
 
-        
+        if packet_hook:
+            pkts = [(packet_hook(cap), meta) for (cap, meta) in pkts]
+
+
         for i, (cap, meta) in enumerate(pkts, start = 1):
             # IPG - if not provided, take from cap
             if ipg_usec == None: