shuttle/
exec.rs

1use crate::client::SshClient;
2use crate::config::ShuttleConfig;
3use rand_core::RngCore;
4use serde_json::{json, Value};
5use std::sync::Arc;
6use std::time::Duration;
7use tokio::sync::Mutex;
8
9const ALLOWED_INTERPRETERS: &[&str] = &[
10    "/bin/bash",
11    "/bin/sh",
12    "/usr/bin/bash",
13    "/usr/bin/sh",
14    "/usr/bin/python3",
15    "/usr/bin/python",
16    "/usr/bin/perl",
17    "/usr/bin/ruby",
18];
19
20/// Execute a single command on a remote host.
21pub async fn ssh_exec(
22    client: &Arc<Mutex<SshClient>>,
23    config: &ShuttleConfig,
24    params: &Value,
25) -> omegon_extension::Result<Value> {
26    let command = params
27        .get("command")
28        .and_then(|v| v.as_str())
29        .ok_or_else(|| omegon_extension::Error::invalid_params("missing 'command'"))?;
30
31    let timeout_secs = params
32        .get("timeout_secs")
33        .and_then(|v| v.as_u64())
34        .unwrap_or(config.default_timeout_secs)
35        .min(3600);
36
37    let client = client.lock().await;
38    let result = client
39        .exec(
40            command,
41            Duration::from_secs(timeout_secs),
42            config.max_output_bytes,
43        )
44        .await
45        .map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
46
47    let stdout = String::from_utf8_lossy(&result.stdout);
48    let stderr = String::from_utf8_lossy(&result.stderr);
49
50    Ok(json!({
51        "host": client.host_name(),
52        "exit_code": result.exit_code,
53        "stdout": stdout,
54        "stderr": stderr,
55        "truncated": result.truncated,
56    }))
57}
58
59/// Upload and execute a multi-line script on a remote host.
60///
61/// The script is piped through a random-delimited heredoc to the interpreter.
62/// The delimiter is generated per-invocation to prevent injection.
63pub async fn ssh_script(
64    client: &Arc<Mutex<SshClient>>,
65    config: &ShuttleConfig,
66    params: &Value,
67) -> omegon_extension::Result<Value> {
68    let script = params
69        .get("script")
70        .and_then(|v| v.as_str())
71        .ok_or_else(|| omegon_extension::Error::invalid_params("missing 'script'"))?;
72
73    let interpreter = params
74        .get("interpreter")
75        .and_then(|v| v.as_str())
76        .unwrap_or("/bin/bash");
77
78    if !ALLOWED_INTERPRETERS.contains(&interpreter) {
79        return Err(omegon_extension::Error::invalid_params(format!(
80            "interpreter not allowed: {interpreter}. Allowed: {}",
81            ALLOWED_INTERPRETERS.join(", ")
82        )));
83    }
84
85    let timeout_secs = params
86        .get("timeout_secs")
87        .and_then(|v| v.as_u64())
88        .unwrap_or(config.default_timeout_secs)
89        .min(3600);
90
91    // Generate a random delimiter that cannot appear in the script body.
92    let delimiter = format!(
93        "SHUTTLE_EOF_{:016x}",
94        rand_core::OsRng.next_u64()
95    );
96
97    let wrapped = format!(
98        "{interpreter} -s <<'{delimiter}'\n{script}\n{delimiter}"
99    );
100
101    let client = client.lock().await;
102    let result = client
103        .exec(
104            &wrapped,
105            Duration::from_secs(timeout_secs),
106            config.max_output_bytes,
107        )
108        .await
109        .map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
110
111    let stdout = String::from_utf8_lossy(&result.stdout);
112    let stderr = String::from_utf8_lossy(&result.stderr);
113
114    Ok(json!({
115        "host": client.host_name(),
116        "exit_code": result.exit_code,
117        "stdout": stdout,
118        "stderr": stderr,
119        "truncated": result.truncated,
120    }))
121}
shuttle/exec.rs

shuttle/
exec.rs
