From 9498a6203fe7a32c2eebe9a7a053f737836ef933 Mon Sep 17 00:00:00 2001
From: WinstonGee <wgee@g.ucla.edu>
Date: Thu, 4 Apr 2024 12:40:36 -0700
Subject: [PATCH] enable connecting to Hoffman2 with rsa key

---
 autoqchem/helper_functions.py | 28 ++++++++++++++++++++++++++++
 autoqchem/sge_manager.py      |  7 +++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/autoqchem/helper_functions.py b/autoqchem/helper_functions.py
index f3d8bee..a26ede3 100644
--- a/autoqchem/helper_functions.py
+++ b/autoqchem/helper_functions.py
@@ -64,6 +64,34 @@ def ssh_connect_password(host, user) -> fabric.Connection:
     return c
 
 
+def ssh_connect_pem(host, user, pem_path) -> fabric.Connection:
+    """Create ssh connection using fabric and paramiko, tries rsa key in .pem file, otherwise asks for password (without DUO authentication).
+
+    :param host: remote host
+    :param user: username to authenticate on remote host
+    :param pem_path: private rsa key - full path to .pem file (optional)
+    :return: fabric.Connection
+    """
+
+    client = paramiko.SSHClient()
+    client.load_system_host_keys()
+
+    try:
+        if pem_path:
+		client.connect(host,username=user,key_filename=pem_path)
+	else:
+		client.connect(host,username=user)
+    		client.get_transport().auth_password(username=user,password=getpass.getpass(f"{user}@{host}'s password:"))
+    except paramiko.ssh_exception.SSHException:
+        pass
+
+    c = fabric.Connection(host)
+    c.client = client
+    c.transport = client.get_transport()
+
+    return c
+
+
 def cleanup_directory_files(dir_path, types=()) -> None:
     """Remove files with specific extension(s) from a directory.
 
diff --git a/autoqchem/sge_manager.py b/autoqchem/sge_manager.py
index 201889d..c2fcd25 100644
--- a/autoqchem/sge_manager.py
+++ b/autoqchem/sge_manager.py
@@ -16,13 +16,15 @@
 class sge_manager(object):
     """SGE manager class."""
 
-    def __init__(self, user, host):
+    def __init__(self, user, host, pem_path=None):
         """Initialize sge manager and load the cache file.
 
         :param user: username at remote host
         :type user: str
         :param host: remote host name
         :type host: str
+        :param pem_path: private rsa key - full path to .pem file (optional)
+        :type host: str
         """
 
         # set workdir and cache file
@@ -39,6 +41,7 @@ def __init__(self, user, host):
 
         self.host = host
         self.user = user
+        self.pem_path = pem_path
         self.remote_dir = f"/u/scratch/{self.user[0]}/{self.user}/gaussian"
         self.connection = None
 
@@ -58,7 +61,7 @@ def connect(self) -> None:
             logger.info(f"Creating connection to {self.host} as {self.user}")
             create_new_connection = True
         if create_new_connection:
-            self.connection = ssh_connect_password(self.host, self.user)    # hoffman2 requires connecting with password (no DUO)
+            self.connection = ssh_connect_pem(self.host, self.user, self.pem_path)    # hoffman2 requires connecting with password or rsa key (no DUO)
             self.connection.run(f"mkdir -p {self.remote_dir}")
             logger.info(f"Connected to {self.host} as {self.user}.")