1 files changed, 91 insertions, 0 deletions
diff --git a/resources/libraries/python/TrafficScriptExecutor.py b/resources/libraries/python/TrafficScriptExecutor.py
new file mode 100644
index 0000000000..2e65a520d0
--- /dev/null
+++ b/resources/libraries/python/TrafficScriptExecutor.py
@@ -0,0 +1,91 @@
+# Copyright (c) 2016 Cisco and/or its affiliates.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Traffic script executor library."""
+
+from constants import Constants
+from ssh import SSH
+from robot.api import logger
+
+__all__ = ['TrafficScriptExecutor']
+
+
+class TrafficScriptExecutor(object):
+    """Traffic script executor utilities."""
+
+    @staticmethod
+    def _escape(string):
+        """Escape quotation mark and dollar mark for shell command.
+
+           :param string: String to escape.
+           :type string: str
+           :return: Escaped string.
+           :rtype: str
+        """
+        return string.replace('"', '\\"').replace("$", "\\$")
+
+    @staticmethod
+    def run_traffic_script_on_node(script_file_name, node, script_args,
+                                   timeout=10):
+        """Run traffic script on the TG node.
+
+           :param script_file_name: Traffic script name
+           :param node: Node to run traffic script on.
+           :param script_args: Traffic scripts arguments.
+           :param timeout: Timeout (optional).
+           :type script_file_name: str
+           :type node: dict
+           :type script_args: str
+           :type timeout: int
+        """
+        logger.trace("{}".format(timeout))
+        ssh = SSH()
+        ssh.connect(node)
+        cmd = ("cd {}; virtualenv env && " +
+               "export PYTHONPATH=${{PWD}}; " +
+               ". ${{PWD}}/env/bin/activate; " +
+               "resources/traffic_scripts/{} {}") \
+                  .format(Constants.REMOTE_FW_DIR, script_file_name,
+                          script_args)
+        (ret_code, stdout, stderr) = ssh.exec_command_sudo(
+            'sh -c "{}"'.format(TrafficScriptExecutor._escape(cmd)),
+            timeout=timeout)
+        logger.debug("stdout: {}".format(stdout))
+        logger.debug("stderr: {}".format(stderr))
+        logger.debug("ret_code: {}".format(ret_code))
+        if ret_code != 0:
+            raise Exception("Traffic script execution failed")
+
+    @staticmethod
+    def traffic_script_gen_arg(rx_if, tx_if, src_mac, dst_mac, src_ip, dst_ip):
+        """Generate traffic script basic arguments string.
+
+           :param rx_if: Interface that sends traffic.
+           :param tx_if: Interface that receives traffic.
+           :param src_mac: Source MAC address.
+           :param dst_mac: Destination MAC address.
+           :param src_ip: Source IP address.
+           :param dst_ip: Destination IP address.
+           :type rx_if: str
+           :type tx_if: str
+           :type src_mac: str
+           :type dst_mac: str
+           :type src_ip: str
+           :type dst_ip: str
+           :return: Traffic script arguments string.
+           :rtype: str
+        """
+        args = '--rx_if {0} --tx_if {1} --src_mac {2} --dst_mac {3} --src_ip' \
+            ' {4} --dst_ip {5}'.format(rx_if, tx_if, src_mac, dst_mac, src_ip,
+                                       dst_ip)
+        return args