fix: decoupled runtime.Command from SSHConnection

Author: evilsocket

src/runtime/cmd.rs

@@ -1,7 +1,5 @@
 use std::{collections::BTreeMap, fmt};
 
-use super::SSHConnection;
-
 #[derive(Debug)]
 pub struct CommandLine {
     pub sudo: bool,
@@ -111,12 +109,7 @@ impl CommandLine {
         args
     }
 
-    pub async fn execute(&self, ssh: Option<SSHConnection>) -> anyhow::Result<String> {
-        // execyte via ssh
-        if let Some(ssh) = ssh {
-            return ssh.execute(self.sudo, &self.app, &self.args).await;
-        }
-
+    pub async fn execute(&self) -> anyhow::Result<String> {
         log::debug!("executing command: {}", self);
         log::debug!("full command details: {:?}", self);
 
@@ -226,7 +219,7 @@ mod tests {
             env: BTreeMap::new(),
             temp_env_file: None,
         };
-        let result = cmd.execute(None).await.unwrap();
+        let result = cmd.execute().await.unwrap();
         assert_eq!(result, "Hello, World!");
     }
 
@@ -240,7 +233,7 @@ mod tests {
             env: BTreeMap::new(),
             temp_env_file: None,
         };
-        let result = cmd.execute(None).await.unwrap();
+        let result = cmd.execute().await.unwrap();
         assert!(result.contains("EXIT CODE:"));
         assert!(result.contains("ERROR:"));
     }
@@ -258,7 +251,7 @@ mod tests {
             env: BTreeMap::new(),
             temp_env_file: None,
         };
-        let result = cmd.execute(None).await.unwrap();
+        let result = cmd.execute().await.unwrap();
         assert!(result.contains("Hello"));
         assert!(result.contains("Error"));
     }
@@ -273,7 +266,7 @@ mod tests {
             env: BTreeMap::new(),
             temp_env_file: None,
         };
-        let result = cmd.execute(None).await;
+        let result = cmd.execute().await;
         assert!(result.is_err());
     }
 

src/runtime/mod.rs

@@ -86,8 +86,9 @@ pub(crate) async fn execute_call(
     let mut needs_container = false;
     let mut can_ssh = false;
 
+    // if --ssh was provided
     if let Some(ssh) = ssh.as_ref() {
-        // if ssh is set we don't need a container
+        // check if the app is in $PATH on the ssh host
         can_ssh = ssh.app_in_path(&command_line.app).await?;
         if !can_ssh {
             log::warn!(
@@ -112,6 +113,7 @@ pub(crate) async fn execute_call(
         }
     }
 
+    // wrap the command line in a container if needed
     let command_line = if needs_container {
         let container = match container {
             Some(c) => c,
@@ -159,12 +161,17 @@ pub(crate) async fn execute_call(
     }
 
     // finally execute the command line
-    let content = command_line
-        .execute(match can_ssh {
-            true => ssh,
-            false => None,
-        })
-        .await?;
+    let content = if can_ssh {
+        // execute via ssh
+        ssh.as_ref()
+            .unwrap()
+            .execute(command_line.sudo, &command_line.app, &command_line.args)
+            .await?
+    } else {
+        // execute locally
+        command_line.execute().await?
+    };
+
     Ok(openai::CallResultMessage {
         role: "tool".to_string(),
         call_id: call.id.clone(),