support ssh tunnel password

Ui_mainWindow.py

@@ -145,9 +145,29 @@ def setupUi(self, MainWindow):
         self.tunnelUserEdit.setText("")
         self.tunnelUserEdit.setObjectName("tunnelUserEdit")
         self.horizontalLayout.addWidget(self.tunnelUserEdit)
-        self.keyFileLabel = QtWidgets.QLabel(self.tunnelWidget)
-        self.keyFileLabel.setObjectName("keyFileLabel")
-        self.horizontalLayout.addWidget(self.keyFileLabel)
+        self.keyComboBox = QtWidgets.QComboBox(self.tunnelWidget)
+        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Fixed, QtWidgets.QSizePolicy.Fixed)
+        sizePolicy.setHorizontalStretch(0)
+        sizePolicy.setVerticalStretch(0)
+        sizePolicy.setHeightForWidth(self.keyComboBox.sizePolicy().hasHeightForWidth())
+        self.keyComboBox.setSizePolicy(sizePolicy)
+        self.keyComboBox.setMinimumSize(QtCore.QSize(110, 0))
+        self.keyComboBox.setMaximumSize(QtCore.QSize(110, 16777215))
+        self.keyComboBox.setObjectName("keyComboBox")
+        self.keyComboBox.addItem("")
+        self.keyComboBox.addItem("")
+        self.horizontalLayout.addWidget(self.keyComboBox)
+        self.pwdLineEdit = QtWidgets.QLineEdit(self.tunnelWidget)
+        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Maximum, QtWidgets.QSizePolicy.Fixed)
+        sizePolicy.setHorizontalStretch(0)
+        sizePolicy.setVerticalStretch(0)
+        sizePolicy.setHeightForWidth(self.pwdLineEdit.sizePolicy().hasHeightForWidth())
+        self.pwdLineEdit.setSizePolicy(sizePolicy)
+        self.pwdLineEdit.setMinimumSize(QtCore.QSize(188, 22))
+        self.pwdLineEdit.setMaximumSize(QtCore.QSize(188, 16777215))
+        self.pwdLineEdit.setEchoMode(QtWidgets.QLineEdit.Password)
+        self.pwdLineEdit.setObjectName("pwdLineEdit")
+        self.horizontalLayout.addWidget(self.pwdLineEdit)
         self.keyFileButton = QtWidgets.QPushButton(self.tunnelWidget)
         self.keyFileButton.setEnabled(True)
         sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Maximum, QtWidgets.QSizePolicy.Fixed)
@@ -225,7 +245,10 @@ def retranslateUi(self, MainWindow):
         self.serverLabel.setText(_translate("MainWindow", "PrestoServer:"))
         self.gatewayLabel.setText(_translate("MainWindow", "Gateway:"))
         self.tunnelUserLabel.setText(_translate("MainWindow", "User:"))
-        self.keyFileLabel.setText(_translate("MainWindow", "KeyFile:"))
+        self.keyComboBox.setCurrentText(_translate("MainWindow", "KeyFile"))
+        self.keyComboBox.setItemText(0, _translate("MainWindow", "KeyFile"))
+        self.keyComboBox.setItemText(1, _translate("MainWindow", "Password"))
+        self.pwdLineEdit.setPlaceholderText(_translate("MainWindow", "Enter Password..."))
         self.keyFileButton.setText(_translate("MainWindow", "Choose Key File..."))
         self.clearButton.setText(_translate("MainWindow", "Clear"))
         self.sqlEdit.setPlaceholderText(_translate("MainWindow", "Write your Presto SQL"))

mainWindow.ui

@@ -293,9 +293,65 @@
             </widget>
            </item>
            <item>
-            <widget class="QLabel" name="keyFileLabel">
-             <property name="text">
-              <string>KeyFile:</string>
+            <widget class="QComboBox" name="keyComboBox">
+             <property name="sizePolicy">
+              <sizepolicy hsizetype="Fixed" vsizetype="Fixed">
+               <horstretch>0</horstretch>
+               <verstretch>0</verstretch>
+              </sizepolicy>
+             </property>
+             <property name="minimumSize">
+              <size>
+               <width>110</width>
+               <height>0</height>
+              </size>
+             </property>
+             <property name="maximumSize">
+              <size>
+               <width>110</width>
+               <height>16777215</height>
+              </size>
+             </property>
+             <property name="currentText">
+              <string>KeyFile</string>
+             </property>
+             <item>
+              <property name="text">
+               <string>KeyFile</string>
+              </property>
+             </item>
+             <item>
+              <property name="text">
+               <string>Password</string>
+              </property>
+             </item>
+            </widget>
+           </item>
+           <item>
+            <widget class="QLineEdit" name="pwdLineEdit">
+             <property name="sizePolicy">
+              <sizepolicy hsizetype="Maximum" vsizetype="Fixed">
+               <horstretch>0</horstretch>
+               <verstretch>0</verstretch>
+              </sizepolicy>
+             </property>
+             <property name="minimumSize">
+              <size>
+               <width>188</width>
+               <height>22</height>
+              </size>
+             </property>
+             <property name="maximumSize">
+              <size>
+               <width>188</width>
+               <height>16777215</height>
+              </size>
+             </property>
+             <property name="echoMode">
+              <enum>QLineEdit::Password</enum>
+             </property>
+             <property name="placeholderText">
+              <string>Enter Password...</string>
              </property>
             </widget>
            </item>

tunnel.py

@@ -8,8 +8,8 @@ def __init__(self):
 
     @timeout(6, "Fail to establish tunnel")
     def start_tunnel(
-            self, private_server_url, ssh_username, ssh_pkey, gateway_url,
-            local_url):
+            self, private_server_url, gateway_url, ssh_username, use_keyfile,
+            ssh_pkey, ssh_pwd, local_url):
         self.stop_tunnel()
         gateway_url = gateway_url.split(':')
         private_server = private_server_url.split(':')
@@ -20,13 +20,23 @@ def start_tunnel(
         local = local_url.split(':')
         if len(local) < 2:
             raise Exception("Url Port Not Specified")
-        self.server = SSHTunnelForwarder(
-            (gateway_url[0], int(gateway_url[1])),
-            ssh_username=ssh_username,
-            ssh_pkey=ssh_pkey,
-            remote_bind_address=(private_server[0], int(private_server[1])),
-            local_bind_address=(local[0], int(local[1]))
-        )
+
+        if use_keyfile:
+            self.server = SSHTunnelForwarder(
+                (gateway_url[0], int(gateway_url[1])),
+                ssh_username=ssh_username,
+                ssh_pkey=ssh_pkey,
+                remote_bind_address=(private_server[0], int(private_server[1])),
+                local_bind_address=(local[0], int(local[1]))
+            )
+        else:
+            self.server = SSHTunnelForwarder(
+                (gateway_url[0], int(gateway_url[1])),
+                ssh_username=ssh_username,
+                ssh_password=ssh_pwd,
+                remote_bind_address=(private_server[0], int(private_server[1])),
+                local_bind_address=(local[0], int(local[1]))
+            )
         self.server.start()
 
     def stop_tunnel(self):

winApp.py

@@ -17,6 +17,8 @@
 GATEWAY = 'GATEWAY'
 TUNNEL_USER = 'TUNNEL_USER'
 KEYFILE = 'KEYFILE'
+KEYTYPE = 'KEYTYPE'
+TUNNEL_PWD = 'TUNNEL_PWD'
 
 
 class WinApp(QMainWindow, Ui_MainWindow):
@@ -32,6 +34,7 @@ def __init__(self):
         self.tunnelWidget.hide()
         self.keyFileButton.clicked.connect(self.choose_keyfile)
         self.clearButton.clicked.connect(self.clear_tunnel)
+        self.keyComboBox.activated[str].connect(self.set_key_type)
 
         self.db_engine = DBEngine()
         self.db_engine.result_signal.connect(self.show_results)
@@ -113,17 +116,20 @@ def __init__(self):
         self.completer = SQLCompleter(self)
         self.sqlEdit.setCompleter(self.completer)
 
+        self.set_key_type(self.keyComboBox.currentText())
         self.sqlEdit.setFocus()
 
     def connect(self):
+        self.tablesWidget.clear()
+        self.schemaView.clear()
         try:
             if self.serverEdit.text().strip() and self.gatewayEdit.text().strip() and self.tunnelUserEdit.text().strip():
                 self.statusbar.showMessage('starting tunnel')
                 QApplication.processEvents()
                 self.tunnel.start_tunnel(
-                    self.serverEdit.text().strip(),
-                    self.tunnelUserEdit.text().strip(),
-                    self.tunnel_keyfile, self.gatewayEdit.text().strip(),
+                    self.serverEdit.text().strip(), self.gatewayEdit.text().strip(),
+                    self.tunnelUserEdit.text().strip(), self.keyComboBox.currentText() == 'KeyFile',
+                    self.tunnel_keyfile, self.pwdLineEdit.text(),
                     self.urlEdit.text().strip())
 
             self.statusbar.showMessage('connecting')
@@ -134,7 +140,6 @@ def connect(self):
             self.statusbar.showMessage('fetching db info')
             QApplication.processEvents()
 
-            self.tablesWidget.clear()
             dbs = self.db_engine.dbs()
             for db in dbs:
                 db_item = QTreeWidgetItem([db])
@@ -322,6 +327,7 @@ def clear_tunnel(self):
         self.serverEdit.clear()
         self.gatewayEdit.clear()
         self.tunnelUserEdit.clear()
+        self.pwdLineEdit.clear()
 
     def close_all(self):
         self.save_settings()
@@ -337,6 +343,8 @@ def restore_settings(self):
         self.tunnel_keyfile = self.settings.value(KEYFILE, "")
         if self.tunnel_keyfile:
             self.keyFileButton.setText(self.tunnel_keyfile.split('/')[-1])
+        self.pwdLineEdit.setText(self.settings.value(TUNNEL_PWD, ""))
+        self.keyComboBox.setCurrentText(self.settings.value(KEYTYPE, "KeyFile"))
 
     def save_settings(self):
         self.settings.setValue(URL, self.urlEdit.text().strip())
@@ -345,7 +353,17 @@ def save_settings(self):
         self.settings.setValue(GATEWAY, self.gatewayEdit.text().strip())
         self.settings.setValue(TUNNEL_USER, self.tunnelUserEdit.text().strip())
         self.settings.setValue(KEYFILE, self.tunnel_keyfile)
+        self.settings.setValue(KEYTYPE, self.keyComboBox.currentText())
+        self.settings.setValue(TUNNEL_PWD, self.pwdLineEdit.text())
         self.settings.sync()
 
     def show_error(self, error):
         QMessageBox.critical(self, "Error", error)
+
+    def set_key_type(self, key_type):
+        if key_type == 'KeyFile':
+            self.pwdLineEdit.hide()
+            self.keyFileButton.show()
+        else:
+            self.keyFileButton.hide()
+            self.pwdLineEdit.show()