Merge pull request #499 from AutomationSolutionz/ssh-tunnel

Start and Stop SSH Tunnel via Node

Framework/Built_In_Automation/Database/BuiltInFunctions.py

@@ -163,7 +163,7 @@ def db_get_connection(session_name):
         db_con = None
 
         # Get the values
-        db_type = db_params.get(DB_TYPE)
+        db_type = db_params.get(DB_TYPE).lower()
         db_name = db_params.get(DB_NAME)
         db_user_id = db_params.get(DB_USER_ID)
         db_password = db_params.get(DB_PASSWORD)
@@ -368,14 +368,18 @@ def db_select(data_set):
         query = None
         session_name = 'default'
         variable_name = None
+        return_type = None
         for left, mid, right in data_set:
             if left == "query":
                 # Get the and query, and remove any whitespaces
                 query = right.strip()
             if "action" in mid:
                 variable_name = right.strip()
             if 'session' in left.lower():
-                session_name = right.strip() 
+                session_name = right.strip()
+            if 'return' in left.lower():
+                return_type = "records"
+
 
         if variable_name is None:
             CommonUtil.ExecLog(sModuleInfo, "Variable name must be provided.", 3)
@@ -403,6 +407,10 @@ def db_select(data_set):
                     db_rows.append(list(db_row))
 
                 # Set the rows as a shared variable
+                if return_type == 'records':
+                    column_headers = [i[0] for i in db_cursor.description]
+                    db_rows = [dict(zip(column_headers, row)) for row in db_rows]
+
                 sr.Set_Shared_Variables(variable_name, db_rows)
 
         db_con.close()

Framework/Built_In_Automation/Sequential_Actions/action_declarations/common.py

@@ -54,6 +54,8 @@
     { "name": "disable showing value",                       "function": "disable_showing_value",                 "screenshot": "none" },
     { "name": "skip testcases",                              "function": "skip_testcases",                        "screenshot": "none" },
     { "name": "global sleep",                                "function": "global_sleep",                          "screenshot": "none" },
+    { "name": "start ssh tunnel",                            "function": "start_ssh_tunnel",                      "screenshot": "none" },
+    { "name": "stop ssh tunnel",                             "function": "stop_ssh_tunnel",                       "screenshot": "none" },
 
     # Mail actions
     { "name": "send mail",                                   "function": "send_mail",                             "screenshot": "none" },

Framework/Built_In_Automation/Sequential_Actions/common_functions.py

@@ -49,7 +49,7 @@
 from datetime import timedelta
 from .utility import send_email, check_latest_received_email, delete_mail, save_mail,random_mail_factory
 import re
-
+import subprocess
 months = [
     "Unknown",
     "January",
@@ -6653,4 +6653,91 @@ def text_to_speech(data_set):
         tts.save(full_filepath)
         return "passed"
     except:
-        return CommonUtil.Exception_Handler(sys.exc_info())
+        return CommonUtil.Exception_Handler(sys.exc_info())
+
+@logger
+def start_ssh_tunnel(data_set):
+
+    sModuleInfo = inspect.currentframe().f_code.co_name + " : " + MODULE_NAME
+
+    host_ip = None
+    ssh_user = None
+    private_key_path = None
+    remote_bind_host = "127.0.0.1"
+    local_bind_host = "127.0.0.1"
+    local_port = None
+    remote_port = None
+    tunnel_var = None
+
+    for left, mid, right in data_set:
+        if left == 'host ip':
+            host_ip = right.strip()
+        elif left == 'ssh user':
+            ssh_user = right.strip()
+        elif left == 'private key path':
+            private_key_path = right.strip()
+        elif left == 'remote_bind_host':
+            remote_bind_host = right.strip()
+        elif left == 'local_bind_host':
+            local_bind_host = right.strip()
+        elif left == 'local port':
+            local_port = int(right.strip())
+        elif left == 'remote port':
+            remote_port = int(right.strip())
+        elif left == 'start ssh tunnel':
+            tunnel_var = right.strip()
+
+    missing = [x for x in [host_ip,ssh_user,private_key_path,local_port, remote_port, tunnel_var] if x == None]
+    if missing:
+        CommonUtil.ExecLog(sModuleInfo, f"{', '.join(missing)} needs to be provided", 3)
+        return "zeuz_failed"
+
+    tunneling_script_path = Path(os.getcwd()) / 'Utilities' / 'ssh_tunnel.py'
+
+    command = [
+        "python",
+        str(tunneling_script_path),
+        "--host_ip", host_ip,
+        "--ssh_user", ssh_user,
+        "--private_key_path", private_key_path,
+        "--remote_bind_host", "127.0.0.1",
+        "--local_bind_host", "127.0.0.1",
+        "--local_port", str(local_port),
+        "--remote_port", str(remote_port)
+    ]
+
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    from time import sleep
+    import socket
+    sleep(5)
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.settimeout(5)
+        try:
+            s.connect((local_bind_host, local_port))
+        except (socket.timeout, socket.error):
+            CommonUtil.ExecLog(sModuleInfo, f"SSH Tunnel can not be established", 3)
+            return "zeuz_failed"
+
+        sr.Set_Shared_Variables(tunnel_var, process.pid)
+        CommonUtil.ExecLog(sModuleInfo, f"SSH Tunnel is established", 1)
+        return "passed"
+
+
+
+@logger
+def stop_ssh_tunnel(data_set):
+
+    sModuleInfo = inspect.currentframe().f_code.co_name + " : " + MODULE_NAME
+
+    tunnel_pid = None
+    for left, mid, right in data_set:
+        if left == 'stop ssh tunnel':
+            tunnel_pid = int(right.strip())
+
+    subprocess.run(['kill', str(tunnel_pid)], check=True)
+    CommonUtil.ExecLog(sModuleInfo, f"SSH Tunnel is closed", 1)
+
+    return "passed"
+