Authentication using ssh-agent (#248)

* Support authentication using ssh-agent. * Unify exception handling for ssh authentication. * Implement review comments. --------- Co-authored-by: dosas <[email protected]>
SatelliteQE · Oct 2, 2023 · 7f0040d · 7f0040d
1 parent ab3a662
commit 7f0040d
Showing 1 changed file with 19 additions and 8 deletions.
diff --git a/broker/session.py b/broker/session.py
@@ -54,18 +54,29 @@ def __init__(self, **kwargs):
         sock.settimeout(kwargs.get("timeout"))
         port = kwargs.get("port", 22)
         key_filename = kwargs.get("key_filename")
+        password = kwargs.get("password")
         timeout = kwargs.get("timeout", 60)
         helpers.simple_retry(sock.connect, [(host, port)], max_timeout=timeout)
         self.session = ssh2_Session()
         self.session.handshake(sock)
-        if key_filename:
-            if not Path(key_filename).exists():
-                raise FileNotFoundError(f"Key not found in '{key_filename}'")
-            self.session.userauth_publickey_fromfile(user, key_filename)
-        elif kwargs.get("password"):
-            self.session.userauth_password(user, kwargs["password"])
-        else:
-            raise exceptions.AuthenticationError("No password or key file provided.")
+        try:
+            if key_filename:
+                auth_type = "Key"
+                if not Path(key_filename).exists():
+                    raise FileNotFoundError(f"Key not found in '{key_filename}'")
+                self.session.userauth_publickey_fromfile(user, key_filename)
+            elif password:
+                auth_type = "Password"
+                self.session.userauth_password(user, password)
+            elif user:
+                auth_type = "Session"
+                self.session.agent_auth(user)
+            else:
+                raise exceptions.AuthenticationError("No password or key file provided.")
+        except Exception as err:  # noqa: BLE001
+            raise exceptions.AuthenticationError(
+                f"{auth_type}-based authentication failed."
+            ) from err
 
     @staticmethod
     def _read(channel):