GPL/traffic_scripts/TrafficScriptArg.py

   1 # Copyright (c) 2020 Cisco and/or its affiliates.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
   4 #
   5 # Licensed under the Apache License 2.0 or
   6 # GNU General Public License v2.0 or later;  you may not use this file
   7 # except in compliance with one of these Licenses. You
   8 # may obtain a copy of the Licenses at:
   9 #
  10 #     http://www.apache.org/licenses/LICENSE-2.0
  11 #     https://www.gnu.org/licenses/old-licenses/gpl-2.0-standalone.html
  12 #
  13 # Note: If this file is linked with Scapy, which is GPLv2+, your use of it
  14 # must be under GPLv2+.  If at any point in the future it is no longer linked
  15 # with Scapy (or other GPLv2+ licensed software), you are free to choose Apache 2.
  16 #
  17 # Unless required by applicable law or agreed to in writing, software
  18 # distributed under the License is distributed on an "AS IS" BASIS,
  19 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  20 # See the License for the specific language governing permissions and
  21 # limitations under the License.
  22
  23 """Traffic scripts argument parser library."""
  24
  25 import argparse
  26
  27
  28 class TrafficScriptArg:
  29     """Traffic scripts argument parser.
  30
  31     Parse arguments for traffic script. Default has two arguments '--tx_if'
  32     and '--rx_if'. You can provide more arguments. All arguments have string
  33     representation of the value. You can add also optional arguments. Default
  34     value for optional arguments is empty string.
  35
  36     :param more_args: List of additional arguments (optional).
  37     :param opt_args: List of optional arguments (optional).
  38     :type more_args: list
  39     :type opt_args: list
  40
  41     :Example:
  42
  43     >>> from TrafficScriptArg import TrafficScriptArg
  44     >>> args = TrafficScriptArg(['src_mac', 'dst_mac', 'src_ip', 'dst_ip'])
  45     """
  46
  47     def __init__(self, more_args=None, opt_args=None):
  48         parser = argparse.ArgumentParser()
  49         parser.add_argument(u"--tx_if", help=u"interface that sends traffic")
  50         parser.add_argument(u"--rx_if", help=u"interface that receives traffic")
  51
  52         if more_args is not None:
  53             for arg in more_args:
  54                 arg_name = f"--{arg}"
  55                 parser.add_argument(arg_name)
  56
  57         if opt_args is not None:
  58             for arg in opt_args:
  59                 arg_name = f"--{arg}"
  60                 parser.add_argument(arg_name, nargs=u"?", default=u"")
  61
  62         self._parser = parser
  63         self._args = vars(parser.parse_args())
  64
  65     def get_arg(self, arg_name):
  66         """Get argument value.
  67
  68         :param arg_name: Argument name.
  69         :type arg_name: str
  70         :returns: Argument value.
  71         :rtype: str
  72         """
  73         arg_val = self._args.get(arg_name)
  74         if arg_val is None:
  75             raise Exception(f"Argument '{arg_name}' not found")
  76
  77         return arg_val