resources/libraries/python/TrafficScriptExecutor.py

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