Skip to content

Commit

Permalink
Refactor SSH key and passphrase handling in Node class
Browse files Browse the repository at this point in the history
- Remove direct handling of SSH key and passphrase decrypt for both primary and jump host connections
- Move ssh key and passphrase handling to asynchssh
- Update asyncssh options configuration for more flexible SSH authentication
  • Loading branch information
Ryan Wolfe committed Jan 31, 2025
1 parent d89f2aa commit 6922031
Showing 1 changed file with 23 additions and 38 deletions.
61 changes: 23 additions & 38 deletions suzieq/poller/worker/nodes/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,27 +142,21 @@ async def initialize(self, **kwargs) -> TNode:
self.ssh_config_file = kwargs.get("ssh_config_file", None)
self.enable_password = kwargs.get('enable_password')

passphrase: str = kwargs.get("passphrase", None)
self.pvtkey_file = kwargs.get("ssh_keyfile", None)
self.jump_host_pvtkey_file = kwargs.get("jump_host_key_file", None)
self.passphrase: str = kwargs.get("passphrase", None)

jump_host = kwargs.get("jump_host", "")
if jump_host:
jump_result = urlparse(jump_host)
self.jump_user = jump_result.username or self.username
self.jump_host = jump_result.hostname
self.jump_host_key = None
if jump_result.port:
self.jump_port = jump_result.port
else:
self.jump_port = 22
pvtkey_file = kwargs.pop('jump_host_key_file')
if pvtkey_file:
self.jump_host_key = self._decrypt_pvtkey(pvtkey_file,
passphrase)
if not self.jump_host_key:
raise SqPollerConfError('Unable to read private key file'
f' at {pvtkey_file}')
else:
self.jump_host = None
self.jump_host_key = None

self.ignore_known_hosts = kwargs.get('ignore_known_hosts', False)
self.slow_host = kwargs.get('slow_host', False)
Expand All @@ -174,15 +168,6 @@ async def initialize(self, **kwargs) -> TNode:
# 4 is a number we picked to limit using up too many SSH sessions
# Many newer implementations allow upto 5 simultaneous SSH sessions
self.batch_size = 4
pvtkey_file = kwargs.get("ssh_keyfile", None)
if pvtkey_file:
self.pvtkey = self._decrypt_pvtkey(pvtkey_file, passphrase)
if not self.pvtkey:
self.logger.error("ERROR: Falling back to password for "
f"{self.address}:{self.port}")
self.pvtkey = None
else:
self.pvtkey = None

self._init_service_queue()

Expand Down Expand Up @@ -258,8 +243,8 @@ def _decrypt_pvtkey(self, pvtkey_file: str, passphrase: str) -> str:
passphrase)
except Exception as e: # pylint: disable=broad-except
self.logger.error(
f"ERROR: Unable to read private key file {pvtkey_file}"
f"for jump host due to {e}")
f"ERROR: Unable to read private key file {pvtkey_file} "
f"due to {e}")

return keydata

Expand Down Expand Up @@ -571,9 +556,10 @@ async def _init_jump_host_connection(
if self._tunnel:
return

if self.jump_host_key:
if self.jump_host_pvtkey_file:
jump_host_options = asyncssh.SSHClientConnectionOptions(
client_keys=self.jump_host_key,
client_keys=self.jump_host_pvtkey_file,
passphrase=self.passphrase,
connect_timeout=self.connect_timeout,
)

Expand Down Expand Up @@ -619,9 +605,7 @@ def _init_ssh_options(self) -> asyncssh.SSHClientConnectionOptions:
options = asyncssh.SSHClientConnectionOptions(
connect_timeout=self.connect_timeout,
username=self.username,
agent_identities=self.pvtkey if self.pvtkey else None,
client_keys=self.pvtkey if self.pvtkey else None,
password=self.password if not self.pvtkey else None,
password=self.password if not self.pvtkey_file else None,
kex_algs='+diffie-hellman-group1-sha1', # for older boxes
encryption_algs='+aes256-cbc', # for older boxes
)
Expand All @@ -637,19 +621,20 @@ def _init_ssh_options(self) -> asyncssh.SSHClientConnectionOptions:
)

if self.pvtkey_file:
cert_file = f"{self.pvtkey_file}-cert.pub"
if os.path.isfile(cert_file):
self.logger.debug(f"Using SSH cert file: {cert_file}")
client_keys = [self.pvtkey, cert_file]
else:
client_keys = self.pvtkey
else:
client_keys = None
# Giving just the filename let's asyncssh know to look for the
# corresponding cert file in the same directory.
# Ref: https://asyncssh.readthedocs.io/en/stable/api.html#specifying-private-keys
client_keys = self.pvtkey_file
options = asyncssh.SSHClientConnectionOptions(
options=options,
client_keys=client_keys,
)

options = asyncssh.SSHClientConnectionOptions(
options=options,
client_keys=client_keys,
)
if self.passphrase:
options = asyncssh.SSHClientConnectionOptions(
options=options,
passphrase=self.passphrase,
)

return options

Expand Down

0 comments on commit 6922031

Please sign in to comment.