use crate::client::SshClient;
use crate::config::ShuttleConfig;
use serde_json::{json, Value};
use std::path::Path;
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;
const BLOCKED_HOME_PREFIXES: &[&str] = &[
".config/styrene",
".ssh",
".aws",
".gnupg",
".gpg",
".omegon",
".kube",
".docker",
".npmrc",
".pypirc",
".netrc",
".config/gcloud",
".azure",
];
const BLOCKED_ABSOLUTE_PREFIXES: &[&str] = &[
"/etc",
"/private/etc",
"/var/run/secrets",
"/private/var/run/secrets",
"/proc",
"/sys",
"/dev",
"/Library/Keychains",
"/System",
];
fn validate_local_path(path: &str) -> omegon_extension::Result<std::path::PathBuf> {
let p = Path::new(path);
let canonical = p
.canonicalize()
.or_else(|_| {
p.parent()
.and_then(|parent| parent.canonicalize().ok())
.map(|parent| parent.join(p.file_name().unwrap_or_default()))
.ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::NotFound, "parent dir not found")
})
})
.map_err(|e| {
omegon_extension::Error::invalid_params(format!("invalid path: {e}"))
})?;
let canonical_str = canonical.to_string_lossy();
if canonical_str.contains("..") {
return Err(omegon_extension::Error::invalid_params(
"path traversal not allowed",
));
}
for prefix in BLOCKED_ABSOLUTE_PREFIXES {
if canonical.starts_with(prefix) {
return Err(omegon_extension::Error::invalid_params(
"access to system directories is blocked",
));
}
}
if let Some(home) = dirs::home_dir() {
for prefix in BLOCKED_HOME_PREFIXES {
let sensitive = home.join(prefix);
if canonical.starts_with(&sensitive) {
return Err(omegon_extension::Error::invalid_params(
"access to credential/config directories is blocked",
));
}
}
}
Ok(canonical)
}
pub async fn scp_push(
client: &Arc<Mutex<SshClient>>,
_config: &ShuttleConfig,
params: &Value,
) -> omegon_extension::Result<Value> {
let local_path = params
.get("local_path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'local_path'"))?;
let remote_path = params
.get("remote_path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'remote_path'"))?;
let local = validate_local_path(local_path)?;
if !local.exists() {
return Err(omegon_extension::Error::invalid_params(format!(
"local file not found: {local_path}"
)));
}
if !local.is_file() {
return Err(omegon_extension::Error::invalid_params(format!(
"local path is not a file: {local_path}"
)));
}
let data = tokio::fs::read(&local)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("read {local_path}: {e}")))?;
let bytes_written = data.len();
let client_guard = client.lock().await;
let sftp = client_guard
.sftp()
.await
.map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
sftp.create(remote_path)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("sftp create: {e}")))?
.write_all(&data)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("sftp write: {e}")))?;
Ok(json!({
"host": client_guard.host_name(),
"local_path": local_path,
"remote_path": remote_path,
"bytes_written": bytes_written,
}))
}
pub async fn scp_pull(
client: &Arc<Mutex<SshClient>>,
_config: &ShuttleConfig,
params: &Value,
) -> omegon_extension::Result<Value> {
let remote_path = params
.get("remote_path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'remote_path'"))?;
let local_path = params
.get("local_path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'local_path'"))?;
let local = validate_local_path(local_path)?;
let client_guard = client.lock().await;
let sftp = client_guard
.sftp()
.await
.map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
let data = sftp
.read(remote_path)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("sftp read: {e}")))?;
tokio::fs::write(&local, &data)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("write {local_path}: {e}")))?;
Ok(json!({
"host": client_guard.host_name(),
"remote_path": remote_path,
"local_path": local_path,
"bytes_written": data.len(),
}))
}
pub async fn sftp_ls(
client: &Arc<Mutex<SshClient>>,
_config: &ShuttleConfig,
params: &Value,
) -> omegon_extension::Result<Value> {
let path = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'path'"))?;
let client_guard = client.lock().await;
let sftp = client_guard
.sftp()
.await
.map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
let entries = sftp
.read_dir(path)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("sftp readdir: {e}")))?;
let items: Vec<Value> = entries
.into_iter()
.map(|entry| {
json!({
"name": entry.file_name(),
"size": entry.metadata().size.unwrap_or(0),
"is_dir": entry.metadata().is_dir(),
})
})
.collect();
Ok(json!({
"host": client_guard.host_name(),
"path": path,
"entries": items,
}))
}
pub async fn sftp_read(
client: &Arc<Mutex<SshClient>>,
config: &ShuttleConfig,
params: &Value,
) -> omegon_extension::Result<Value> {
let path = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| omegon_extension::Error::invalid_params("missing 'path'"))?;
let max_bytes = params
.get("max_bytes")
.and_then(|v| v.as_u64())
.map(|v| v as usize)
.unwrap_or(config.max_output_bytes)
.min(config.max_output_bytes);
let client_guard = client.lock().await;
let sftp = client_guard
.sftp()
.await
.map_err(|e| omegon_extension::Error::internal_error(e.to_string()))?;
let data = sftp
.read(path)
.await
.map_err(|e| omegon_extension::Error::internal_error(format!("sftp read: {e}")))?;
let truncated = data.len() > max_bytes;
let content = if truncated {
String::from_utf8_lossy(&data[..max_bytes]).into_owned()
} else {
String::from_utf8_lossy(&data).into_owned()
};
Ok(json!({
"host": client_guard.host_name(),
"path": path,
"content": content,
"size": data.len(),
"truncated": truncated,
}))
}