13696dee8d2036f4a50f930f96daa0b74a8f7a65
[csit.git] / 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 resources.libraries.python.constants import Constants
17 from resources.libraries.python.ssh import SSH
18
19 __all__ = ['TrafficScriptExecutor']
20
21
22 class TrafficScriptExecutor(object):
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         :return: Escaped string.
32         :rtype: str
33         """
34         return string.replace('"', '\\"').replace("$", "\\$")
35
36     @staticmethod
37     def run_traffic_script_on_node(script_file_name, node, script_args,
38                                    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 = ("cd {}; " +
59                "virtualenv --system-site-packages --never-download env && " +
60                "export PYTHONPATH=${{PWD}}; " +
61                ". ${{PWD}}/env/bin/activate; " +
62                "resources/traffic_scripts/{} {}") \
63                   .format(Constants.REMOTE_FW_DIR, script_file_name,
64                           script_args)
65         ret_code, stdout, stderr = ssh.exec_command_sudo(
66             'sh -c "{cmd}"'.format(cmd=TrafficScriptExecutor._escape(cmd)),
67             timeout=timeout)
68         if ret_code != 0:
69             if "RuntimeError: ICMP echo Rx timeout" in stderr:
70                 raise RuntimeError("ICMP echo Rx timeout")
71             elif "RuntimeError: DHCP REQUEST Rx timeout" in stderr:
72                 raise RuntimeError("DHCP REQUEST Rx timeout")
73             elif "RuntimeError: DHCP DISCOVER Rx timeout" in stderr:
74                 raise RuntimeError("DHCP DISCOVER Rx timeout")
75             elif "RuntimeError: TCP/UDP Rx timeout" in stderr:
76                 raise RuntimeError("TCP/UDP Rx timeout")
77             elif "Error occurred: ARP reply timeout" in stdout:
78                 raise RuntimeError("ARP reply timeout")
79             elif "RuntimeError: ESP packet Rx timeout" in stderr:
80                 raise RuntimeError("ESP packet Rx timeout")
81             else:
82                 raise RuntimeError("Traffic script execution failed")
83
84     @staticmethod
85     def traffic_script_gen_arg(rx_if, tx_if, src_mac, dst_mac, src_ip, dst_ip):
86         """Generate traffic script basic arguments string.
87
88         :param rx_if: Interface that receives traffic.
89         :param tx_if: Interface that sends traffic.
90         :param src_mac: Source MAC address.
91         :param dst_mac: Destination MAC address.
92         :param src_ip: Source IP address.
93         :param dst_ip: Destination IP address.
94         :type rx_if: str
95         :type tx_if: str
96         :type src_mac: str
97         :type dst_mac: str
98         :type src_ip: str
99         :type dst_ip: str
100         :return: Traffic script arguments string.
101         :rtype: str
102         """
103         args = ('--rx_if {0} --tx_if {1} --src_mac {2} --dst_mac {3} --src_ip'
104                 ' {4} --dst_ip {5}').format(rx_if, tx_if, src_mac, dst_mac,
105                                             src_ip, dst_ip)
106         return args

©2016 FD.io a Linux Foundation Collaborative Project. All Rights Reserved.
Linux Foundation is a registered trademark of The Linux Foundation. Linux is a registered trademark of Linus Torvalds.
Please see our privacy policy and terms of use.