From d46ae180459d07221e3073090fdf56ff3bbe3161 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Thu, 20 Sep 2018 14:53:17 -0600
Subject: [PATCH] Establish SSH tunnel before connecting to SQL DB

* Also added disconnect_from_db() function.
---
 .bin/Scripts/functions/hw_diags.py | 40 ++++++++++++++++++++++++++----
 .bin/Scripts/settings/main.py      |  1 +
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/.bin/Scripts/functions/hw_diags.py b/.bin/Scripts/functions/hw_diags.py
index d42f3a31..467d3563 100644
--- a/.bin/Scripts/functions/hw_diags.py
+++ b/.bin/Scripts/functions/hw_diags.py
@@ -10,7 +10,8 @@ from functions.common import *
 ost_db = {
     'Connection': None,
     'Cursor': None,
-    'Errors': False
+    'Errors': False,
+    'Tunnel': None,
     }
 
 # STATIC VARIABLES
@@ -92,13 +93,39 @@ TESTS = {
     }
 
 def connect_to_db():
-    """Connect to osTicket database."""
-    #TODO# Open SSH tunnel to DB server
-    ost_db['Connection'] = mariadb.connect(
-        user=DB_USER, password=DB_PASS, database=DB_NAME, host=DB_HOST)
+    """Connect to osTicket database via SSH tunnel."""
+    cmd = [
+        'ssh', '-N', '-p', '2222', '-L3306:127.0.0.1:3306',
+        '{user}@{host}'.format(user=SSH_USER, host=DB_HOST),
+        ]
+
+    # Establish SSH tunnel unless one already exists
+    if not ost_db['Tunnel']:
+        ost_db['Tunnel'] = popen_program(cmd)
+
+    # Establish SQL connection (try a few times in case SSH is slow)
+    for x in range(5):
+        sleep(3)
+        try:
+            ost_db['Connection'] = mariadb.connect(
+                user=DB_USER, password=DB_PASS, database=DB_NAME)
+        except:
+            # Just try again
+            pass
+        else:
+            break
     ost_db['Cursor'] = ost_db['Connection'].cursor()
     ost_db['Errors'] = False
 
+def disconnect_from_db():
+    """Disconnect from SQL DB and close SSH tunnel."""
+    if ost_db['Cursor']:
+        ost_db['Cursor'].close()
+    if ost_db['Connection']:
+        ost_db['Connection'].close()
+    if ost_db['Tunnel']:
+        ost_db['Tunnel'].kill()
+
 def generate_horizontal_graph(rates):
     """Generate two-line horizontal graph from rates, returns str."""
     line_top = ''
@@ -287,6 +314,9 @@ def menu_diags(*args):
         elif selection == 'Q':
             break
 
+    # Done
+    disconnect_from_db()
+
 def osticket_get_ticket_name(ticket_id):
     """Lookup ticket and return name as str."""
     ticket_name = 'Unknown'
diff --git a/.bin/Scripts/settings/main.py b/.bin/Scripts/settings/main.py
index 4e6d4083..0d5944a2 100644
--- a/.bin/Scripts/settings/main.py
+++ b/.bin/Scripts/settings/main.py
@@ -16,6 +16,7 @@ DB_HOST='127.0.0.1'
 DB_NAME='osticket'
 DB_USER='wizardkit'
 DB_PASS='Abracadabra'
+SSH_USER='sql_tunnel'
 # Live Linux
 MPRIME_LIMIT='7'                          # of minutes to run Prime95 during hw-diags
 ROOT_PASSWORD='Abracadabra'