forked from Eeveid/lightOps
123 lines
3.7 KiB
Rust
123 lines
3.7 KiB
Rust
use anyhow::Result;
|
|
use base64::Engine;
|
|
use lightops_common::protocol::AgentMessage;
|
|
use portable_pty::{native_pty_system, CommandBuilder, PtySize};
|
|
use serde_json::Value;
|
|
use std::{
|
|
io::{Read, Write},
|
|
sync::{Arc, Mutex},
|
|
thread,
|
|
};
|
|
use tokio::sync::mpsc;
|
|
|
|
pub struct TerminalHandle {
|
|
writer: Arc<Mutex<Box<dyn Write + Send>>>,
|
|
}
|
|
|
|
impl TerminalHandle {
|
|
pub fn write(&self, data: String, binary: bool) -> Result<()> {
|
|
let bytes = if binary {
|
|
base64::engine::general_purpose::STANDARD.decode(data)?
|
|
} else {
|
|
data.into_bytes()
|
|
};
|
|
self.writer.lock().expect("pty writer").write_all(&bytes)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub fn open(
|
|
stream_id: String,
|
|
tx: mpsc::UnboundedSender<AgentMessage>,
|
|
meta: Value,
|
|
) -> Result<TerminalHandle> {
|
|
let cols = meta.get("cols").and_then(Value::as_u64).unwrap_or(100) as u16;
|
|
let rows = meta.get("rows").and_then(Value::as_u64).unwrap_or(30) as u16;
|
|
let pty_system = native_pty_system();
|
|
let pair = pty_system.openpty(PtySize {
|
|
rows,
|
|
cols,
|
|
pixel_width: 0,
|
|
pixel_height: 0,
|
|
})?;
|
|
let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string());
|
|
let mut cmd = CommandBuilder::new(shell);
|
|
open_with_command(stream_id, tx, pair, &mut cmd)
|
|
}
|
|
|
|
pub fn open_docker_exec(
|
|
stream_id: String,
|
|
tx: mpsc::UnboundedSender<AgentMessage>,
|
|
meta: Value,
|
|
) -> Result<TerminalHandle> {
|
|
let cols = meta.get("cols").and_then(Value::as_u64).unwrap_or(100) as u16;
|
|
let rows = meta.get("rows").and_then(Value::as_u64).unwrap_or(30) as u16;
|
|
let container_id = meta
|
|
.get("container_id")
|
|
.and_then(Value::as_str)
|
|
.ok_or_else(|| anyhow::anyhow!("缺少容器 ID"))?;
|
|
validate_docker_id(container_id)?;
|
|
let shell = meta.get("shell").and_then(Value::as_str).unwrap_or("sh");
|
|
let shell = if shell == "bash" { "bash" } else { "sh" };
|
|
let pty_system = native_pty_system();
|
|
let pair = pty_system.openpty(PtySize {
|
|
rows,
|
|
cols,
|
|
pixel_width: 0,
|
|
pixel_height: 0,
|
|
})?;
|
|
let mut cmd = CommandBuilder::new("docker");
|
|
cmd.arg("exec");
|
|
cmd.arg("-it");
|
|
cmd.arg(container_id);
|
|
cmd.arg(shell);
|
|
open_with_command(stream_id, tx, pair, &mut cmd)
|
|
}
|
|
|
|
fn open_with_command(
|
|
stream_id: String,
|
|
tx: mpsc::UnboundedSender<AgentMessage>,
|
|
pair: portable_pty::PtyPair,
|
|
cmd: &mut CommandBuilder,
|
|
) -> Result<TerminalHandle> {
|
|
cmd.env("TERM", "xterm-256color");
|
|
let mut child = pair.slave.spawn_command(cmd)?;
|
|
let mut reader = pair.master.try_clone_reader()?;
|
|
let writer = Arc::new(Mutex::new(pair.master.take_writer()?));
|
|
let close_id = stream_id.clone();
|
|
thread::spawn(move || {
|
|
let mut buf = [0u8; 8192];
|
|
loop {
|
|
match reader.read(&mut buf) {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
let data = String::from_utf8_lossy(&buf[..n]).to_string();
|
|
let _ = tx.send(AgentMessage::StreamData {
|
|
stream_id: stream_id.clone(),
|
|
data,
|
|
binary: false,
|
|
});
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
let _ = child.kill();
|
|
let _ = tx.send(AgentMessage::StreamClose {
|
|
stream_id: close_id,
|
|
reason: Some("终端已关闭".into()),
|
|
});
|
|
});
|
|
Ok(TerminalHandle { writer })
|
|
}
|
|
|
|
fn validate_docker_id(id: &str) -> Result<()> {
|
|
if id.len() > 200
|
|
|| !id
|
|
.chars()
|
|
.all(|c| c.is_ascii_alphanumeric() || ".:/@_-".contains(c))
|
|
{
|
|
anyhow::bail!("Docker 容器标识无效");
|
|
}
|
|
Ok(())
|
|
}
|