Skip to content

Commit

Permalink
Merge pull request #29 from doyle-lab-ucla/dev2
Browse files Browse the repository at this point in the history
enable connecting to Hoffman2 with rsa key
  • Loading branch information
beef-broccoli authored Apr 10, 2024
2 parents a8b4af2 + 9498a62 commit af04d50
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 2 deletions.
28 changes: 28 additions & 0 deletions autoqchem/helper_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
7 changes: 5 additions & 2 deletions autoqchem/sge_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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}.")

Expand Down

0 comments on commit af04d50

Please sign in to comment.