resources/libraries/python/TrafficScriptExecutor.py

   1 # Copyright (c) 2016 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 script executor library."""
  15
  16 from constants import Constants
  17 from ssh import SSH
  18 from robot.api import logger
  19
  20 __all__ = ['TrafficScriptExecutor']
  21
  22
  23 class TrafficScriptExecutor(object):
  24     """Traffic script executor utilities."""
  25
  26     @staticmethod
  27     def _escape(string):
  28         """Escape quotation mark and dollar mark for shell command.
  29
  30            :param string: String to escape.
  31            :type string: str
  32            :return: Escaped string.
  33            :rtype: str
  34         """
  35         return string.replace('"', '\\"').replace("$", "\\$")
  36
  37     @staticmethod
  38     def run_traffic_script_on_node(script_file_name, node, script_args,
  39                                    timeout=10):
  40         """Run traffic script on the TG node.
  41
  42            :param script_file_name: Traffic script name
  43            :param node: Node to run traffic script on.
  44            :param script_args: Traffic scripts arguments.
  45            :param timeout: Timeout (optional).
  46            :type script_file_name: str
  47            :type node: dict
  48            :type script_args: str
  49            :type timeout: int
  50         """
  51         logger.trace("{}".format(timeout))
  52         ssh = SSH()
  53         ssh.connect(node)
  54         cmd = ("cd {}; virtualenv env && " +
  55                "export PYTHONPATH=${{PWD}}; " +
  56                ". ${{PWD}}/env/bin/activate; " +
  57                "resources/traffic_scripts/{} {}") \
  58                   .format(Constants.REMOTE_FW_DIR, script_file_name,
  59                           script_args)
  60         (ret_code, stdout, stderr) = ssh.exec_command_sudo(
  61             'sh -c "{}"'.format(TrafficScriptExecutor._escape(cmd)),
  62             timeout=timeout)
  63         logger.debug("stdout: {}".format(stdout))
  64         logger.debug("stderr: {}".format(stderr))
  65         logger.debug("ret_code: {}".format(ret_code))
  66         if ret_code != 0:
  67             raise Exception("Traffic script execution failed")
  68
  69     @staticmethod
  70     def traffic_script_gen_arg(rx_if, tx_if, src_mac, dst_mac, src_ip, dst_ip):
  71         """Generate traffic script basic arguments string.
  72
  73            :param rx_if: Interface that receives traffic.
  74            :param tx_if: Interface that sends traffic.
  75            :param src_mac: Source MAC address.
  76            :param dst_mac: Destination MAC address.
  77            :param src_ip: Source IP address.
  78            :param dst_ip: Destination IP address.
  79            :type rx_if: str
  80            :type tx_if: str
  81            :type src_mac: str
  82            :type dst_mac: str
  83            :type src_ip: str
  84            :type dst_ip: str
  85            :return: Traffic script arguments string.
  86            :rtype: str
  87         """
  88         args = '--rx_if {0} --tx_if {1} --src_mac {2} --dst_mac {3} --src_ip' \
  89             ' {4} --dst_ip {5}'.format(rx_if, tx_if, src_mac, dst_mac, src_ip,
  90                                        dst_ip)
  91         return args