add /ssh-exec command

src/ssh.rs

@@ -1,4 +1,11 @@
-use russh::client::Handler;
+use std::{fmt::Debug, sync::Arc, time::Duration};
+
+use anyhow::{Context as _, Result as AnyhowResult};
+use russh::{
+    client::{AuthResult, Handle, Handler},
+    keys::PrivateKeyWithHashAlg,
+};
+use tokio::net::ToSocketAddrs;
 
 pub(crate) struct ClientHandler {}
 
@@ -12,3 +19,121 @@ impl Handler for ClientHandler {
         Ok(true)
     }
 }
+
+pub(crate) async fn connect_ssh<A: ToSocketAddrs + Debug, K: AsRef<[u8]>, P: Into<String>>(
+    connect_timeout: Duration,
+    auth_timeout: Duration,
+    addr: A,
+    username: &str,
+    key: Option<K>,
+    password: Option<P>,
+) -> AnyhowResult<Handle<ClientHandler>> {
+    let ssh_config = Arc::new(russh::client::Config {
+        inactivity_timeout: None,
+        ..Default::default()
+    });
+    log::info!(
+        "Connecting to SSH {:?} {} key & {} password",
+        &addr,
+        if key.is_some() { "with" } else { "without" },
+        if password.is_some() {
+            "with"
+        } else {
+            "without"
+        }
+    );
+    let mut connection = tokio::time::timeout(
+        connect_timeout,
+        russh::client::connect(ssh_config, addr, crate::ssh::ClientHandler {}),
+    )
+    .await
+    .context("Connection to host timed out")?
+    .context("Cannot connect to host")?;
+
+    let key = match key {
+        Some(key) => {
+            let privkey = russh::keys::PrivateKey::from_openssh(key).context("Invalid SSH key")?;
+            Some(PrivateKeyWithHashAlg::new(
+                Arc::new(privkey),
+                connection.best_supported_rsa_hash().await?.flatten(),
+            ))
+        }
+        None => None,
+    };
+
+    let mut auth_success = false;
+    if !auth_success {
+        if let Some(key) = key {
+            log::info!("Trying to authenticate using key");
+            match tokio::time::timeout(
+                auth_timeout,
+                connection.authenticate_publickey(username, key),
+            )
+            .await
+            {
+                Ok(Ok(AuthResult::Success)) => {
+                    auth_success = true;
+                }
+                Ok(Ok(AuthResult::Failure {
+                    remaining_methods,
+                    partial_success,
+                })) => {
+                    if partial_success || !remaining_methods.contains(&russh::MethodKind::Password)
+                    {
+                        if partial_success {
+                            anyhow::bail!("SSH auth failed: multi-factor auth is not supported");
+                        } else {
+                            anyhow::bail!("SSH auth failed: no auth methods left");
+                        }
+                    }
+                    log::warn!("SSH public key rejected");
+                }
+                Ok(Err(e)) => {
+                    log::warn!("SSH public key auth failed: {e}");
+                }
+                Err(_) => {
+                    log::warn!("SSH public key auth timed out");
+                }
+            }
+        } else {
+            log::info!("No key provided");
+        }
+    }
+    if !auth_success {
+        if let Some(password) = password {
+            log::info!("Trying to authenticate using password");
+            match tokio::time::timeout(
+                auth_timeout,
+                connection.authenticate_password(username, password),
+            )
+            .await
+            {
+                Ok(Ok(AuthResult::Success)) => {
+                    auth_success = true;
+                }
+                Ok(Ok(AuthResult::Failure {
+                    remaining_methods: _,
+                    partial_success,
+                })) => {
+                    if partial_success {
+                        anyhow::bail!("SSH auth failed: multi-factor auth is not supported");
+                    } else {
+                        anyhow::bail!("SSH auth failed: no auth methods left");
+                    }
+                }
+                Ok(Err(e)) => {
+                    log::warn!("SSH password auth failed: {e}");
+                }
+                Err(_) => {
+                    log::warn!("SSH password auth timed out");
+                }
+            }
+        } else {
+            log::info!("No password provided");
+        }
+    }
+    if !auth_success {
+        anyhow::bail!("SSH auth failed: no auth methods left");
+    }
+    Ok(connection)
+}