port mapping - hardening

1 files changed, 22 insertions, 22 deletions

diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_std.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_std.py
index 72a5ea52..e0b25b1d 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_std.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_std.py
@@ -4,7 +4,6 @@ from trex_stl_packet_builder_scapy import *
 # map ports
 # will destroy all streams/data on the ports
 def stl_map_ports (client, ports = None):
-
     # by default use all ports
     if ports == None:
         ports = client.get_all_ports()
@@ -15,12 +14,15 @@ def stl_map_ports (client, ports = None):
     # generate streams
     base_pkt = CScapyTRexPktBuilder(pkt = Ether()/IP())
     
+    tx_pkts = {}
     pkts = 1
     for port in ports:
+        tx_pkts[pkts] = port
         stream = STLStream(packet = base_pkt,
                            mode = STLTXSingleBurst(pps = 100000, total_pkts = pkts))
 
         client.add_streams(stream, [port])
+        
         pkts = pkts * 2
 
     # inject
@@ -33,35 +35,33 @@ def stl_map_ports (client, ports = None):
     # cleanup
     client.reset(ports = ports)
 
-    table = {}
-    for port in ports:
-        table[port] = None
+    table = {'map': {}, 'bi' : [], 'unknown': []}
 
+    # actual mapping
     for port in ports:
+
         ipackets = stats[port]["ipackets"]
+        table['map'][port] = None
 
-        exp = 1
-        while ipackets >= exp:
-            if ((ipackets & exp) == (exp)):
-                source = int(math.log(exp, 2))
-                table[source] = port
+        for pkts in tx_pkts.keys():
+            if ( (pkts & ipackets) == pkts ):
+                tx_port = tx_pkts[pkts]
+                table['map'][port] = tx_port
 
-            exp *= 2
 
-    if not all(x != None for x in table.values()):
-        raise STLError('unable to map ports')
+    unmapped = list(ports)
+    while len(unmapped) > 0:
+        port_a = unmapped.pop(0)
+        port_b = table['map'][port_a]
 
-    dir_a = set()
-    dir_b = set()
-    for src, dst in table.iteritems():
-        # src is not in
-        if src not in (dir_a, dir_b):
-            if dst in dir_a:
-                dir_b.add(src)
-            else:
-                dir_a.add(src)
+        # if unknown - add to the unknown list
+        if port_b == None:
+            table['unknown'].append(port_a)
 
-    table['dir'] = [list(dir_a), list(dir_b)]
+        # bi-directional ports
+        elif (table['map'][port_b] == port_a):
+            unmapped.remove(port_b)
+            table['bi'].append( (port_a, port_b) )
 
     return table