Skip to content

Commit 7e212e9

Browse files
Ryan Wolfernwolfe
authored andcommitted
Refactor SSH key and passphrase handling in Node class
- 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
1 parent d89f2aa commit 7e212e9

File tree

1 file changed

+23
-38
lines changed

1 file changed

+23
-38
lines changed

suzieq/poller/worker/nodes/node.py

Lines changed: 23 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -142,27 +142,21 @@ async def initialize(self, **kwargs) -> TNode:
142142
self.ssh_config_file = kwargs.get("ssh_config_file", None)
143143
self.enable_password = kwargs.get('enable_password')
144144

145-
passphrase: str = kwargs.get("passphrase", None)
145+
self.pvtkey_file = kwargs.get("ssh_keyfile", None)
146+
self.jump_host_pvtkey_file = kwargs.get("jump_host_key_file", None)
147+
self.passphrase: str = kwargs.get("passphrase", None)
148+
146149
jump_host = kwargs.get("jump_host", "")
147150
if jump_host:
148151
jump_result = urlparse(jump_host)
149152
self.jump_user = jump_result.username or self.username
150153
self.jump_host = jump_result.hostname
151-
self.jump_host_key = None
152154
if jump_result.port:
153155
self.jump_port = jump_result.port
154156
else:
155157
self.jump_port = 22
156-
pvtkey_file = kwargs.pop('jump_host_key_file')
157-
if pvtkey_file:
158-
self.jump_host_key = self._decrypt_pvtkey(pvtkey_file,
159-
passphrase)
160-
if not self.jump_host_key:
161-
raise SqPollerConfError('Unable to read private key file'
162-
f' at {pvtkey_file}')
163158
else:
164159
self.jump_host = None
165-
self.jump_host_key = None
166160

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

187172
self._init_service_queue()
188173

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

264249
return keydata
265250

@@ -571,9 +556,10 @@ async def _init_jump_host_connection(
571556
if self._tunnel:
572557
return
573558

574-
if self.jump_host_key:
559+
if self.jump_host_pvtkey_file:
575560
jump_host_options = asyncssh.SSHClientConnectionOptions(
576-
client_keys=self.jump_host_key,
561+
client_keys=self.jump_host_pvtkey_file,
562+
passphrase=self.passphrase,
577563
connect_timeout=self.connect_timeout,
578564
)
579565

@@ -619,9 +605,7 @@ def _init_ssh_options(self) -> asyncssh.SSHClientConnectionOptions:
619605
options = asyncssh.SSHClientConnectionOptions(
620606
connect_timeout=self.connect_timeout,
621607
username=self.username,
622-
agent_identities=self.pvtkey if self.pvtkey else None,
623-
client_keys=self.pvtkey if self.pvtkey else None,
624-
password=self.password if not self.pvtkey else None,
608+
password=self.password if not self.pvtkey_file else None,
625609
kex_algs='+diffie-hellman-group1-sha1', # for older boxes
626610
encryption_algs='+aes256-cbc', # for older boxes
627611
)
@@ -637,19 +621,20 @@ def _init_ssh_options(self) -> asyncssh.SSHClientConnectionOptions:
637621
)
638622

639623
if self.pvtkey_file:
640-
cert_file = f"{self.pvtkey_file}-cert.pub"
641-
if os.path.isfile(cert_file):
642-
self.logger.debug(f"Using SSH cert file: {cert_file}")
643-
client_keys = [self.pvtkey, cert_file]
644-
else:
645-
client_keys = self.pvtkey
646-
else:
647-
client_keys = None
624+
# Giving just the filename let's asyncssh know to look for the
625+
# corresponding cert file in the same directory.
626+
# Ref: https://asyncssh.readthedocs.io/en/stable/api.html#specifying-private-keys
627+
client_keys = self.pvtkey_file
628+
options = asyncssh.SSHClientConnectionOptions(
629+
options=options,
630+
client_keys=client_keys,
631+
)
648632

649-
options = asyncssh.SSHClientConnectionOptions(
650-
options=options,
651-
client_keys=client_keys,
652-
)
633+
if self.passphrase:
634+
options = asyncssh.SSHClientConnectionOptions(
635+
options=options,
636+
passphrase=self.passphrase,
637+
)
653638

654639
return options
655640

0 commit comments

Comments
 (0)