capture - draft commit

Signed-off-by: imarom <imarom@cisco.com>

4 files changed, 189 insertions, 26 deletions

diff --git a/scripts/automation/trex_control_plane/stl/console/trex_console.py b/scripts/automation/trex_control_plane/stl/console/trex_console.py
index 38a1fca4..b0ab70e0 100755
--- a/scripts/automation/trex_control_plane/stl/console/trex_console.py
+++ b/scripts/automation/trex_control_plane/stl/console/trex_console.py
@@ -348,8 +348,8 @@ class TRexConsole(TRexGeneralCmd):
         
     @verify_connected
     def do_capture (self, line):
-        '''Start PCAP capturing on port'''
-        self.stateless_client.start_capture_line(line)
+        '''Manage PCAP captures'''
+        self.stateless_client.capture_line(line)
 
     def help_capture (self):
         self.do_capture("-h")
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 1b57218f..d75c554e 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
@@ -26,6 +26,7 @@ import random
 import json
 import traceback
 import os.path
+import argparse
 
 ############################     logger     #############################
 ############################                #############################
@@ -2961,7 +2962,7 @@ class STLClient(object):
             Resolves ports (ARP resolution)
 
             :parameters:
-                ports          - for which ports to apply a unique sniffer (each port gets a unique file)
+                ports          - which ports to resolve
                 retires        - how many times to retry on each port (intervals of 100 milliseconds)
                 verbose        - log for each request the response
             :raises:
@@ -3022,7 +3023,7 @@ class STLClient(object):
             
         self.logger.pre_cmd("Starting PCAP capturing up to {0} packets".format(limit))
         
-        rc = self._transmit("start_capture", params = {'limit': limit, 'tx': tx_ports, 'rx': rx_ports})
+        rc = self._transmit("capture", params = {'command': 'start', 'limit': limit, 'tx': tx_ports, 'rx': rx_ports})
         self.logger.post_cmd(rc)
 
 
@@ -3032,24 +3033,82 @@ class STLClient(object):
 
 
     @__api_check(True)
-    def stop_capture (self, ports = None):
+    def stop_capture (self, capture_id, output_filename):
         """
-            Removes RX sniffer from port(s)
+            Stops an active capture
+
+            :parameters:
+                capture_id        - an active capture ID to stop
+                output_filename   - output filename to save capture
 
             :raises:
                 + :exe:'STLError'
 
         """
-        ports = ports if ports is not None else self.get_acquired_ports()
-        ports = self._validate_port_list(ports)
 
-        self.logger.pre_cmd("Removing RX sniffers on port(s) {0}:".format(ports))
-        rc = self.__remove_rx_sniffer(ports)
+        
+        
+        # stopping a capture requires:
+        # 1. stopping
+        # 2. fetching
+        # 3. saving to file
+        
+        # stop
+        
+        self.logger.pre_cmd("Stopping PCAP capture {0}".format(capture_id))
+        rc = self._transmit("capture", params = {'command': 'stop', 'capture_id': capture_id})
         self.logger.post_cmd(rc)
+        if not rc:
+            raise STLError(rc)
+        
+        # pkt count
+        pkt_count = rc.data()['pkt_count']
+            
+        if not output_filename or pkt_count == 0:
+            return
+        
+        self.logger.pre_cmd("Writing {0} packets to '{1}'".format(pkt_count, output_filename))
+
+        # create a PCAP file
+        writer = RawPcapWriter(output_filename, linktype = 1)
+        writer._write_header(None)
+        
+        # fetch
+        while True:
+            rc = self._transmit("capture", params = {'command': 'fetch', 'capture_id': capture_id, 'pkt_limit': 50})
+            if not rc:
+                self.logger.post_cmd(rc)
+                raise STLError(rc)
+        
+            pkts = rc.data()['pkts']
+            for pkt in pkts:
+                ts = pkt['ts']
+                pkt_bin = base64.b64decode(pkt['binary'])
+                writer._write_packet(pkt_bin, sec = 0, usec = 0)
+                
+            if rc.data()['pending'] == 0:
+                break
+        
+        self.logger.post_cmd(rc)
+        
+
+    # get capture status
+    @__api_check(True)
+    def get_capture_status (self):
+        """
+            returns a list of all active captures
+            each element in the list is an object containing
+            info about the capture
+
+        """
+
+        rc = self._transmit("capture", params = {'command': 'status'})
 
         if not rc:
             raise STLError(rc)
 
+        return rc.data()
+
     
     @__api_check(True)
     def set_rx_queue (self, ports = None, size = 1000):
@@ -3766,23 +3825,71 @@ class STLClient(object):
 
              
     @__console
-    def start_capture_line (self, line):
-        '''Starts PCAP recorder on port(s)'''
+    def capture_line (self, line):
+        '''Manage PCAP recorders'''
 
-        parser = parsing_opts.gen_parser(self,
-                                         "capture",
-                                         self.start_capture_line.__doc__,
-                                         parsing_opts.TX_PORT_LIST,
-                                         parsing_opts.RX_PORT_LIST,
-                                         parsing_opts.LIMIT)
+        # default
+        if not line:
+            line = "show"
+
+        parser = parsing_opts.gen_parser(self, "capture", self.capture_line.__doc__)
+        subparsers = parser.add_subparsers(title = "commands", dest="commands")
+
+        # start
+        start_parser = subparsers.add_parser('start', help = "starts a new capture")
+        start_parser.add_arg_list(parsing_opts.TX_PORT_LIST,
+                                  parsing_opts.RX_PORT_LIST,
+                                  parsing_opts.LIMIT)
+
+        # stop
+        stop_parser = subparsers.add_parser('stop', help = "stops an active capture")
+        stop_parser.add_arg_list(parsing_opts.CAPTURE_ID,
+                                 parsing_opts.OUTPUT_FILENAME)
+
+        # show
+        show_parser = subparsers.add_parser('show', help = "show all active captures")
+
+        opts = parser.parse_args(line.split())
 
-        opts = parser.parse_args(line.split(), default_ports = self.get_acquired_ports(), verify_acquired = True)
         if not opts:
             return opts
 
-        self.start_capture(opts.tx_port_list, opts.rx_port_list, opts.limit)
+        # start
+        if opts.commands == 'start':
+            if not opts.tx_port_list and not opts.rx_port_list:
+                start_parser.formatted_error('please provide either --tx or --rx')
+                return
+
+            self.start_capture(opts.tx_port_list, opts.rx_port_list, opts.limit)
+
+        # stop
+        elif opts.commands == 'stop':
+            self.stop_capture(opts.capture_id, opts.output_filename)
+
+        # show
+        else:
+            data = self.get_capture_status()
+
+            stats_table = text_tables.TRexTextTable()
+            stats_table.set_cols_align(["c"] * 6)
+            stats_table.set_cols_width([15] * 6)
+
+            for elem in data:
+                row = [elem['id'],
+                       elem['state'],
+                       '[{0}/{1}]'.format(elem['count'], elem['limit']),
+                       format_num(elem['bytes'], suffix = 'B'),
+                       bitfield_to_str(elem['filter']['tx']),
+                       bitfield_to_str(elem['filter']['rx'])]
+                
+                stats_table.add_rows([row], header=False)
+
+            stats_table.header(['ID', 'Status', 'Count', 'Bytes', 'TX Ports', 'RX Ports'])
+            text_tables.print_table_with_header(stats_table, "Captures")
+            
 
         return RC_OK()
+
         
 
     @__console
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
index cbbacb27..c386451b 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
@@ -107,4 +107,18 @@ def list_remove_dup (l):
             
     return tmp
             
-        
+def bitfield_to_list (bf):
+    rc = []
+    bitpos = 0
+
+    while bf > 0:
+        if bf & 0x1:
+            rc.append(bitpos)
+        bitpos += 1
+        bf = bf >> 1
+
+    return rc
+
+def bitfield_to_str (bf):
+    lst = bitfield_to_list(bf)
+    return "-" if not lst else ', '.join([str(x) for x in lst])
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
index 265c43fb..cb594ef4 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
@@ -69,6 +69,8 @@ RX_PORT_LIST
 SRC_IPV4
 DST_IPV4
 
+CAPTURE_ID
+
 GLOBAL_STATS
 PORT_STATS
 PORT_STATUS
@@ -81,12 +83,14 @@ EXTENDED_INC_ZERO_STATS
 
 STREAMS_MASK
 CORE_MASK_GROUP
+CAPTURE_PORTS_GROUP
 
 # ALL_STREAMS
 # STREAM_LIST_WITH_ALL
 
 # list of ArgumentGroup types
 MUTEX
+NON_MUTEX
 
 '''
 
@@ -392,7 +396,6 @@ OPTIONS_DB = {MULTIPLIER: ArgumentPack(['-m', '--multiplier'],
                                             {'help': 'Output PCAP filename',
                                              'dest': 'output_filename',
                                              'default': None,
-                                             'required': True,
                                              'type': str}),
 
 
@@ -612,6 +615,12 @@ OPTIONS_DB = {MULTIPLIER: ArgumentPack(['-m', '--multiplier'],
                                           'help': 'A list of ports to capture on the RX side',
                                           'default': []}),
                 
+              CAPTURE_ID: ArgumentPack(['-i', '--id'],
+                                  {'help': "capture ID to remove",
+                                   'dest': "capture_id",
+                                   'type': int,
+                                   'required': True}),
+
               # advanced options
               PORT_LIST_WITH_ALL: ArgumentGroup(MUTEX, [PORT_LIST,
                                                         ALL_PORTS],
@@ -636,6 +645,7 @@ OPTIONS_DB = {MULTIPLIER: ArgumentPack(['-m', '--multiplier'],
                                                       CORE_MASK],
                                               {'required': False}),
 
+              CAPTURE_PORTS_GROUP: ArgumentGroup(NON_MUTEX, [TX_PORT_LIST, RX_PORT_LIST], {}),
               }
 
 class _MergeAction(argparse._AppendAction):
@@ -654,12 +664,30 @@ class _MergeAction(argparse._AppendAction):
 
 class CCmdArgParser(argparse.ArgumentParser):
 
-    def __init__(self, stateless_client, *args, **kwargs):
+    def __init__(self, stateless_client = None, x = None, *args, **kwargs):
         super(CCmdArgParser, self).__init__(*args, **kwargs)
         self.stateless_client = stateless_client
         self.cmd_name = kwargs.get('prog')
         self.register('action', 'merge', _MergeAction)
 
+
+    def add_arg_list (self, *args):
+        populate_parser(self, *args)
+
+    def add_subparsers(self, *args, **kwargs):
+        sub = super(CCmdArgParser, self).add_subparsers(*args, **kwargs)
+
+        add_parser = sub.add_parser
+        stateless_client = self.stateless_client
+
+        def add_parser_hook (self, *args, **kwargs):
+            parser = add_parser(self, *args, **kwargs)
+            parser.stateless_client = stateless_client
+            return parser
+
+        sub.add_parser = add_parser_hook
+        return sub
+
     # hook this to the logger
     def _print_message(self, message, file=None):
         self.stateless_client.logger.log(message)
@@ -730,13 +758,15 @@ class CCmdArgParser(argparse.ArgumentParser):
             # recover from system exit scenarios, such as "help", or bad arguments.
             return RC_ERR("'{0}' - {1}".format(self.cmd_name, "no action"))
 
+    def formatted_error (self, msg):
+        self.print_usage()
+        self.stateless_client.logger.log(msg)
+
 
 def get_flags (opt):
     return OPTIONS_DB[opt].name_or_flags
 
-def gen_parser(stateless_client, op_name, description, *args):
-    parser = CCmdArgParser(stateless_client, prog=op_name, conflict_handler='resolve',
-                           description=description)
+def populate_parser (parser, *args):
     for param in args:
         try:
 
@@ -752,6 +782,12 @@ def gen_parser(stateless_client, op_name, description, *args):
                     for sub_argument in argument.args:
                         group.add_argument(*OPTIONS_DB[sub_argument].name_or_flags,
                                            **OPTIONS_DB[sub_argument].options)
+
+                elif argument.type == NON_MUTEX:
+                    group = parser.add_argument_group(**argument.options)
+                    for sub_argument in argument.args:
+                        group.add_argument(*OPTIONS_DB[sub_argument].name_or_flags,
+                                           **OPTIONS_DB[sub_argument].options)
                 else:
                     # ignore invalid objects
                     continue
@@ -764,6 +800,12 @@ def gen_parser(stateless_client, op_name, description, *args):
         except KeyError as e:
             cause = e.args[0]
             raise KeyError("The attribute '{0}' is missing as a field of the {1} option.\n".format(cause, param))
+
+def gen_parser(stateless_client, op_name, description, *args):
+    parser = CCmdArgParser(stateless_client, prog=op_name, conflict_handler='resolve',
+                           description=description)
+
+    populate_parser(parser, *args)
     return parser