feat: ssh-agent configuration to authenticate

veeso · veeso · commit 0100625c83a5 · 2024-07-09T11:17:14.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,8 @@ Released on 09/07/2024
 
 - Fix: resolved_host from configuration wasn't used to connect
 - `SshOpts::method` now requires `KeyMethod` and `MethodType` to setup key method
+- Feat: Implemented `SshAgentIdentity` to specify the ssh agent configuration to be used to authenticate.
+  - use `SshOpts.ssh_agent_identity()` to set the option
 
 ## 0.2.1
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -59,7 +59,8 @@ extern crate log;
 
 mod ssh;
 pub use ssh::{
-    KeyMethod, MethodType, ParseRule as SshConfigParseRule, ScpFs, SftpFs, SshKeyStorage, SshOpts,
+    KeyMethod, MethodType, ParseRule as SshConfigParseRule, ScpFs, SftpFs, SshAgentIdentity,
+    SshKeyStorage, SshOpts,
 };
 
 // -- utils
diff --git a/src/ssh/commons.rs b/src/ssh/commons.rs
@@ -13,6 +13,7 @@ use ssh2::{MethodType as SshMethodType, Session};
 
 use super::config::Config;
 use super::SshOpts;
+use crate::SshAgentIdentity;
 
 // -- connect
 
@@ -79,29 +80,45 @@ pub fn connect(opts: &SshOpts) -> RemoteResult<Session> {
         error!("SSH handshake failed: {}", err);
         return Err(RemoteError::new_ex(RemoteErrorType::ProtocolError, err));
     }
-    // Authenticate with password or key
-    match opts.key_storage.as_ref().and_then(|x| {
-        x.resolve(ssh_config.host.as_str(), ssh_config.username.as_str())
-            .or(x.resolve(
-                ssh_config.resolved_host.as_str(),
-                ssh_config.username.as_str(),
-            ))
-    }) {
-        Some(rsa_key) => {
-            session_auth_with_rsakey(
-                &mut session,
-                &ssh_config.username,
-                rsa_key.as_path(),
-                opts.password.as_deref(),
-                ssh_config.params.identity_file.as_deref(),
-            )?;
+
+    // if use_ssh_agent is enabled, try to authenticate with ssh agent
+    if let Some(ssh_agent_config) = &opts.ssh_agent_identity {
+        match session_auth_with_agent(&mut session, &ssh_config.username, ssh_agent_config) {
+            Ok(_) => {
+                info!("Authenticated with ssh agent");
+                return Ok(session);
+            }
+            Err(err) => {
+                error!("Could not authenticate with ssh agent: {}", err);
+            }
         }
-        None => {
-            session_auth_with_password(
-                &mut session,
-                &ssh_config.username,
-                opts.password.as_deref(),
-            )?;
+    }
+
+    // Authenticate with password or key
+    if !session.authenticated() {
+        match opts.key_storage.as_ref().and_then(|x| {
+            x.resolve(ssh_config.host.as_str(), ssh_config.username.as_str())
+                .or(x.resolve(
+                    ssh_config.resolved_host.as_str(),
+                    ssh_config.username.as_str(),
+                ))
+        }) {
+            Some(rsa_key) => {
+                session_auth_with_rsakey(
+                    &mut session,
+                    &ssh_config.username,
+                    rsa_key.as_path(),
+                    opts.password.as_deref(),
+                    ssh_config.params.identity_file.as_deref(),
+                )?;
+            }
+            None => {
+                session_auth_with_password(
+                    &mut session,
+                    &ssh_config.username,
+                    opts.password.as_deref(),
+                )?;
+            }
         }
     }
     // Return session
@@ -179,6 +196,58 @@ fn set_algo_prefs(session: &mut Session, opts: &SshOpts, config: &Config) -> Rem
     Ok(())
 }
 
+/// Authenticate on session with ssh agent
+fn session_auth_with_agent(
+    session: &mut Session,
+    username: &str,
+    ssh_agent_config: &SshAgentIdentity,
+) -> RemoteResult<()> {
+    let mut agent = session
+        .agent()
+        .map_err(|err| RemoteError::new_ex(RemoteErrorType::ConnectionError, err))?;
+
+    agent
+        .connect()
+        .map_err(|err| RemoteError::new_ex(RemoteErrorType::ConnectionError, err))?;
+
+    agent
+        .list_identities()
+        .map_err(|err| RemoteError::new_ex(RemoteErrorType::ConnectionError, err))?;
+
+    let mut connection_result = Err(RemoteError::new(RemoteErrorType::AuthenticationFailed));
+
+    for identity in agent
+        .identities()
+        .map_err(|err| RemoteError::new_ex(RemoteErrorType::ConnectionError, err))?
+    {
+        if ssh_agent_config.pubkey_matches(identity.blob()) {
+            debug!("Trying to authenticate with ssh agent with key: {identity:?}");
+        } else {
+            continue;
+        }
+        match agent.userauth(username, &identity) {
+            Ok(()) => {
+                connection_result = Ok(());
+                debug!("Authenticated with ssh agent with key: {identity:?}");
+                break;
+            }
+            Err(err) => {
+                debug!("SSH agent auth failed: {err}");
+                connection_result = Err(RemoteError::new_ex(
+                    RemoteErrorType::AuthenticationFailed,
+                    err,
+                ));
+            }
+        }
+    }
+
+    if let Err(err) = agent.disconnect() {
+        warn!("Could not disconnect from ssh agent: {err}");
+    }
+
+    connection_result
+}
+
 /// Authenticate on session with private key
 fn session_auth_with_rsakey(
     session: &mut Session,
diff --git a/src/ssh/mod.rs b/src/ssh/mod.rs
@@ -53,6 +53,39 @@ impl KeyMethod {
 
 // -- ssh options
 
+/// Ssh agent identity
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum SshAgentIdentity {
+    /// Try all identities
+    All,
+    /// Use a specific identity
+    Pubkey(Vec<u8>),
+}
+
+impl From<Vec<u8>> for SshAgentIdentity {
+    fn from(v: Vec<u8>) -> Self {
+        SshAgentIdentity::Pubkey(v)
+    }
+}
+
+impl From<&[u8]> for SshAgentIdentity {
+    fn from(v: &[u8]) -> Self {
+        SshAgentIdentity::Pubkey(v.to_vec())
+    }
+}
+
+impl SshAgentIdentity {
+    /// Check if the provided public key matches the identity
+    ///
+    /// If `All` is provided, this method will always return `true`
+    pub(crate) fn pubkey_matches(&self, blob: &[u8]) -> bool {
+        match self {
+            SshAgentIdentity::All => true,
+            SshAgentIdentity::Pubkey(v) => v == blob,
+        }
+    }
+}
+
 /// Ssh options;
 /// used to build and configure SCP/SFTP client.
 ///
@@ -85,6 +118,8 @@ pub struct SshOpts {
     methods: Vec<KeyMethod>,
     /// Ssh config parser ruleset
     parse_rules: ParseRule,
+    /// Ssh agent configuration for authentication
+    ssh_agent_identity: Option<SshAgentIdentity>,
 }
 
 impl SshOpts {
@@ -104,6 +139,7 @@ impl SshOpts {
             key_storage: None,
             methods: Vec::default(),
             parse_rules: ParseRule::STRICT,
+            ssh_agent_identity: None,
         }
     }
 
@@ -134,6 +170,17 @@ impl SshOpts {
         self
     }
 
+    /// Set configuration for ssh agent
+    ///
+    /// If `None` the ssh agent will be disabled
+    ///
+    /// If `Some(SshAgentIdentity::All)` all identities will be tried
+    /// Otherwise the provided public key will be used
+    pub fn ssh_agent_identity(mut self, ssh_agent_identity: Option<SshAgentIdentity>) -> Self {
+        self.ssh_agent_identity = ssh_agent_identity;
+        self
+    }
+
     /// Set SSH configuration file to read
     ///
     /// The supported options are:
@@ -229,6 +276,16 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_should_tell_whether_pubkey_matches() {
+        let identity = SshAgentIdentity::Pubkey(b"hello".to_vec());
+        assert!(identity.pubkey_matches(b"hello"));
+        assert!(!identity.pubkey_matches(b"world"));
+
+        let identity = SshAgentIdentity::All;
+        assert!(identity.pubkey_matches(b"hello"));
+    }
+
     #[test]
     fn should_initialize_ssh_opts() {
         let opts = SshOpts::new("localhost");
diff --git a/src/ssh/scp.rs b/src/ssh/scp.rs
@@ -1421,11 +1421,14 @@ mod test {
 
     #[cfg(feature = "with-containers")]
     fn setup_client() -> ScpFs {
+        use crate::SshAgentIdentity;
+
         let config_file = ssh_mock::create_ssh_config();
         let mut client = ScpFs::new(
             SshOpts::new("scp")
                 .key_storage(Box::new(ssh_mock::MockSshKeyStorage::default()))
-                .config_file(config_file.path(), ParseRule::ALLOW_UNKNOWN_FIELDS),
+                .config_file(config_file.path(), ParseRule::ALLOW_UNKNOWN_FIELDS)
+                .ssh_agent_identity(Some(SshAgentIdentity::All)),
         );
         assert!(client.connect().is_ok());
         // Create wrkdir
diff --git a/src/ssh/sftp.rs b/src/ssh/sftp.rs
@@ -1238,11 +1238,14 @@ mod test {
 
     #[cfg(feature = "with-containers")]
     fn setup_client() -> SftpFs {
+        use crate::SshAgentIdentity;
+
         let config_file = ssh_mock::create_ssh_config();
         let mut client = SftpFs::new(
             SshOpts::new("sftp")
                 .key_storage(Box::new(ssh_mock::MockSshKeyStorage::default()))
-                .config_file(config_file.path(), ParseRule::ALLOW_UNKNOWN_FIELDS),
+                .config_file(config_file.path(), ParseRule::ALLOW_UNKNOWN_FIELDS)
+                .ssh_agent_identity(Some(SshAgentIdentity::All)),
         );
         assert!(client.connect().is_ok());
         // Create wrkdir
