Merge pull request #9 from dreadnode/feature/ssh

new: implemented optional execution via ssh

Author: evilsocket

Cargo.lock

@@ -14,6 +14,7 @@ actix-cors = "0.7.0"
 actix-web = "4.9.0"
 actix-web-lab = "0.23.0"
 anyhow = "1.0.90"
+async-ssh2-tokio = "0.8.12"
 camino = { version = "1.1.9", features = ["serde"] }
 clap = { version = "4.5.20", features = ["derive"] }
 env_logger = "0.11.5"
@@ -26,6 +27,7 @@ regex = "1.11.0"
 reqwest = "0.12.8"
 serde = { version = "1.0.211", features = ["derive"] }
 serde_yaml = "0.9.34"
+shell-escape = "0.1.5"
 shellexpand = { version = "3.1.0", features = ["full"] }
 tempfile = "3.13.0"
 tokio = { version = "1.41.0", features = ["full"] }

Cargo.toml

@@ -155,6 +155,18 @@ Repeat for multiple variables:
 robopages run -F function_name -A -D target=www.example.com -D foo=bar
 ```
 
+#### SSH
+
+The `run` and `serve` commands support an optional SSH connection string. If provided, commands will be executed over SSH on the given host.
+
+```bash
+robopages serve --ssh user@host:port --ssh-key ~/.ssh/id_ed25519
+```
+
+> [!IMPORTANT]
+> * Setting a SSH connection string will override any container configuration.
+> * If the function requires sudo, the remote host is expected to have passwordless sudo access. 
+
 ### Using with LLMs
 
 The examples folder contains integration examples for [Rigging](/examples/rigging_example.py), [OpenAI](/examples/openai_example.py), [Groq](/examples/groq_example.py), [OLLAMA](/examples/ollama_example.py) and [Nerve](/examples/nerve.md).

README.md

@@ -192,7 +192,7 @@ pub struct Book {
 
 impl Book {
     pub fn from_path(path: Utf8PathBuf, filter: Option<String>) -> anyhow::Result<Self> {
-        log::info!("Searching for pages in {:?}", path);
+        log::debug!("Searching for pages in {:?}", path);
         let mut page_paths = Vec::new();
 
         let path = Utf8PathBuf::from(
@@ -203,28 +203,28 @@ impl Book {
         .canonicalize_utf8()
         .map_err(|e| anyhow::anyhow!("failed to canonicalize path: {}", e))?;
 
-        log::info!("Canonicalized path: {:?}", path);
+        log::debug!("canonicalized path: {:?}", path);
 
         if path.is_file() {
-            log::info!("Path is a file");
+            log::debug!("path is a file");
             eval_if_in_filter!(path, filter, page_paths.push(path.to_path_buf()));
         } else if path.is_dir() {
-            log::info!("Path is a directory, searching for .yml files");
+            log::debug!("path is a directory, searching for .yml files");
             let glob_pattern = path.join("**/*.yml").as_str().to_string();
-            log::info!("Using glob pattern: {}", glob_pattern);
+            log::debug!("using glob pattern: {}", glob_pattern);
 
             for entry in glob(&glob_pattern)? {
                 match entry {
                     Ok(entry_path) => {
-                        log::debug!("Found file: {:?}", entry_path);
+                        log::debug!("found file: {:?}", entry_path);
                         // skip files in hidden directories (starting with .)
                         // but allow the root .robopages directory
                         if let Ok(relative_path) = entry_path.strip_prefix(&path) {
                             if relative_path.components().any(|component| {
                                 let comp_str = component.as_os_str().to_string_lossy();
                                 comp_str.starts_with(".") && comp_str != "." && comp_str != ".."
                             }) {
-                                log::debug!("Skipping hidden file/directory");
+                                log::debug!("skipping hidden file/directory");
                                 continue;
                             }
                         }
@@ -239,13 +239,13 @@ impl Book {
                         }
                     }
                     Err(e) => {
-                        log::error!("Error in glob: {}", e);
+                        log::error!("error in glob: {}", e);
                     }
                 }
             }
         }
 
-        log::info!("Found {} page paths", page_paths.len());
+        log::debug!("found {} page paths", page_paths.len());
 
         if page_paths.is_empty() {
             return Err(anyhow::anyhow!("no pages found in {:?}", path));

src/book/mod.rs

@@ -39,15 +39,6 @@ impl ExecutionFlavor {
         ExecutionFlavor::Error(message)
     }
 
-    pub fn to_string(&self) -> String {
-        match self {
-            Self::Shell(shell) => shell.to_string(),
-            Self::Sudo => "sudo".to_string(),
-            Self::Docker(image) => format!("docker {}", image),
-            Self::Error(message) => message.to_string(),
-        }
-    }
-
     fn get_current_shell() -> String {
         let shell_name = std::env::var("SHELL")
             .map(|s| s.split('/').last().unwrap_or("unknown").to_string())
@@ -111,6 +102,18 @@ impl ExecutionFlavor {
     }
 }
 
+impl std::fmt::Display for ExecutionFlavor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let s = match self {
+            Self::Shell(shell) => shell.to_string(),
+            Self::Sudo => "sudo".to_string(),
+            Self::Docker(image) => format!("docker {}", image),
+            Self::Error(message) => message.to_string(),
+        };
+        write!(f, "{}", s)
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum ExecutionContext {
     #[serde(rename = "cmdline")]
@@ -202,8 +205,8 @@ impl<'a> FunctionRef<'a> {
                         let env_var = std::env::var(&env_var_name);
                         let env_var_value = if let Ok(value) = env_var {
                             value
-                        } else if var_default.is_some() {
-                            var_default.unwrap().to_string()
+                        } else if let Some(def) = var_default {
+                            def.to_string()
                         } else {
                             return Err(anyhow::anyhow!(
                                 "environment variable {} not set",
@@ -217,8 +220,12 @@ impl<'a> FunctionRef<'a> {
                         env_var_value
                     } else if let Some(value) = arguments.get(var_name) {
                         // if the value is empty and there's a default value, use the default value
-                        if value.is_empty() && var_default.is_some() {
-                            var_default.unwrap().to_string()
+                        if value.is_empty() {
+                            if let Some(def) = var_default {
+                                def.to_string()
+                            } else {
+                                value.to_string()
+                            }
                         } else {
                             // otherwise, use the provided value
                             value.to_string()

src/book/runtime.rs

@@ -51,13 +51,14 @@ impl Template {
     }
 }
 
-impl ToString for Template {
-    fn to_string(&self) -> String {
-        match self {
+impl std::fmt::Display for Template {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let s = match self {
             Template::Basic => "basic".to_string(),
             Template::DockerImage => "docker_image".to_string(),
             Template::DockerBuild => "docker_build".to_string(),
-        }
+        };
+        write!(f, "{}", s)
     }
 }
 

src/book/templates.rs

@@ -92,6 +92,15 @@ pub(crate) struct ServeArgs {
     /// Maximum number of parallel calls to execute. Leave to 0 to use all available cores.
     #[clap(long, default_value = "0")]
     workers: usize,
+    /// Optional SSH connection string, if set commands will be executed over SSH on the given host.
+    #[clap(long)]
+    ssh: Option<String>,
+    /// SSH key to use for authentication if --ssh is set.
+    #[clap(long, default_value = "~/.ssh/id_ed25519")]
+    ssh_key: String,
+    /// SSH passphrase to unlock the key.
+    #[clap(long)]
+    ssh_key_passphrase: Option<String>,
 }
 
 #[derive(Debug, Args)]
@@ -108,6 +117,15 @@ pub(crate) struct RunArgs {
     /// Execute the function without user interaction.
     #[clap(long, short = 'A')]
     auto: bool,
+    /// Optional SSH connection string, if set commands will be executed over SSH on the given host.
+    #[clap(long)]
+    ssh: Option<String>,
+    /// SSH key to use for authentication if --ssh is set.
+    #[clap(long, default_value = "~/.ssh/id_ed25519")]
+    ssh_key: String,
+    /// SSH passphrase to unlock the key.
+    #[clap(long)]
+    ssh_key_passphrase: Option<String>,
 }
 
 #[derive(Debug, Args)]

src/cli/mod.rs

@@ -2,12 +2,24 @@ use std::{collections::BTreeMap, sync::Arc};
 
 use crate::{
     book::{flavors::openai, Book},
-    runtime::{self, prompt},
+    runtime::{self, prompt, ssh::SSHConnection},
 };
 
 use super::RunArgs;
 
 pub(crate) async fn run(args: RunArgs) -> anyhow::Result<()> {
+    // parse and validate SSH connection string if provided
+    let ssh = if let Some(ssh_str) = args.ssh {
+        // parse
+        let conn = SSHConnection::from_str(&ssh_str, &args.ssh_key, args.ssh_key_passphrase)?;
+        // make sure we can connect
+        conn.test_connection().await?;
+
+        Some(conn)
+    } else {
+        None
+    };
+
     let book = Arc::new(Book::from_path(args.path, None)?);
     let function = book.get_function(&args.function)?;
 
@@ -39,9 +51,7 @@ pub(crate) async fn run(args: RunArgs) -> anyhow::Result<()> {
         call_type: "function".to_string(),
     };
 
-    log::debug!("running function {:?}", function);
-
-    let result = runtime::execute_call(!args.auto, 10, book, call).await?;
+    let result = runtime::execute_call(ssh, !args.auto, 10, book, call).await?;
 
     println!("\n{}", result.content);
 

src/cli/run.rs

@@ -14,12 +14,14 @@ use crate::book::{
     Book,
 };
 use crate::runtime;
+use crate::runtime::ssh::SSHConnection;
 
 use super::ServeArgs;
 
 struct AppState {
     max_running_tasks: usize,
     book: Arc<Book>,
+    ssh: Option<SSHConnection>,
 }
 
 async fn not_found() -> actix_web::Result<HttpResponse> {
@@ -65,7 +67,15 @@ async fn process_calls(
     state: web::Data<Arc<AppState>>,
     calls: web::Json<Vec<openai::Call>>,
 ) -> actix_web::Result<HttpResponse> {
-    match runtime::execute(false, state.book.clone(), calls.0, state.max_running_tasks).await {
+    match runtime::execute(
+        state.ssh.clone(),
+        false,
+        state.book.clone(),
+        calls.0,
+        state.max_running_tasks,
+    )
+    .await
+    {
         Ok(resp) => Ok(HttpResponse::Ok().json(resp)),
         Err(e) => Err(actix_web::error::ErrorBadRequest(e)),
     }
@@ -76,6 +86,18 @@ pub(crate) async fn serve(args: ServeArgs) -> anyhow::Result<()> {
         log::warn!("external address specified, this is an unsafe configuration as no authentication is provided");
     }
 
+    // parse and validate SSH connection string if provided
+    let ssh = if let Some(ssh_str) = args.ssh {
+        // parse
+        let conn = SSHConnection::from_str(&ssh_str, &args.ssh_key, args.ssh_key_passphrase)?;
+        // make sure we can connect
+        conn.test_connection().await?;
+
+        Some(conn)
+    } else {
+        None
+    };
+
     let book = Arc::new(Book::from_path(args.path, args.filter)?);
     if !args.lazy {
         for page in book.pages.values() {
@@ -103,6 +125,7 @@ pub(crate) async fn serve(args: ServeArgs) -> anyhow::Result<()> {
     let app_state = Arc::new(AppState {
         max_running_tasks,
         book,
+        ssh,
     });
 
     HttpServer::new(move || {

src/cli/serve.rs

@@ -12,7 +12,7 @@ pub(crate) async fn view(args: ViewArgs) -> anyhow::Result<()> {
             println!("    * {} : {}", function_name, function.description);
             println!(
                 "         running with: {}",
-                ExecutionFlavor::for_function(&function)?.to_string()
+                ExecutionFlavor::for_function(&function)?
             );
             println!("         parameters:");
             for (parameter_name, parameter) in &function.parameters {