resources/libraries/python/TrafficScriptArg.py

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