GPL/traffic_scripts/TrafficScriptArg.py

   1 # Copyright (c) 2021 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
  16 # Apache 2.
  17 #
  18 # Unless required by applicable law or agreed to in writing, software
  19 # distributed under the License is distributed on an "AS IS" BASIS,
  20 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  21 # See the License for the specific language governing permissions and
  22 # limitations under the License.
  23
  24 """Traffic scripts argument parser library."""
  25
  26 import argparse
  27
  28
  29 class TrafficScriptArg:
  30     """Traffic scripts argument parser.
  31
  32     Parse arguments for traffic script. Default has two arguments '--tx_if'
  33     and '--rx_if'. You can provide more arguments. All arguments have string
  34     representation of the value. You can add also optional arguments. Default
  35     value for optional arguments is empty string.
  36
  37     :param more_args: List of additional arguments (optional).
  38     :param opt_args: List of optional arguments (optional).
  39     :type more_args: list
  40     :type opt_args: list
  41
  42     :Example:
  43
  44     >>> from TrafficScriptArg import TrafficScriptArg
  45     >>> args = TrafficScriptArg(['src_mac', 'dst_mac', 'src_ip', 'dst_ip'])
  46     """
  47
  48     def __init__(self, more_args=None, opt_args=None):
  49         parser = argparse.ArgumentParser()
  50         parser.add_argument(u"--tx_if", help=u"interface that sends traffic")
  51         parser.add_argument(u"--rx_if", help=u"interface that receives traffic")
  52
  53         if more_args is not None:
  54             for arg in more_args:
  55                 arg_name = f"--{arg}"
  56                 parser.add_argument(arg_name)
  57
  58         if opt_args is not None:
  59             for arg in opt_args:
  60                 arg_name = f"--{arg}"
  61                 parser.add_argument(arg_name, nargs=u"?", default=u"")
  62
  63         self._parser = parser
  64         self._args = vars(parser.parse_args())
  65
  66     def get_arg(self, arg_name):
  67         """Get argument value.
  68
  69         :param arg_name: Argument name.
  70         :type arg_name: str
  71         :returns: Argument value.
  72         :rtype: str
  73         """
  74         arg_val = self._args.get(arg_name)
  75         if arg_val is None:
  76             raise Exception(f"Argument '{arg_name}' not found")
  77
  78         return arg_val