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