Skip to content

Instantly share code, notes, and snippets.

@matthewjberger

matthewjberger/sync-main.rs

Last active August 4, 2024 09:59
Show Gist options
  • Save matthewjberger/62b364c8b5e9c8512f092e3e28744437 to your computer and use it in GitHub Desktop.
Save matthewjberger/62b364c8b5e9c8512f092e3e28744437 to your computer and use it in GitHub Desktop.
Download ZIP
Teleport interactively in rust
Raw
sync-main.rs
// thiserror = "1.0.63"
use std::{
io::{BufRead, BufReader, Write},
process::{Child, Command, Stdio},
sync::{
mpsc::{self, Sender},
Arc, Mutex,
},
thread,
time::Duration,
};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum TeleportError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("UTF-8 conversion error: {0}")]
Utf8(#[from] std::string::FromUtf8Error),
#[error("No hostnames found")]
NoHostnames,
#[error("Invalid hostname")]
InvalidHostname,
#[error("Tunnel already open")]
TunnelAlreadyOpen,
#[error("No tunnel is currently open")]
NoOpenTunnel,
#[error("Failed to capture stdout")]
StdoutCaptureFailed,
#[error("Failed to capture stderr")]
StderrCaptureFailed,
#[error("Failed to capture stdin")]
StdinCaptureFailed,
#[error("Invalid input")]
InvalidInput,
}
pub struct Teleport {
hostnames: Vec<String>,
current_tunnel: Arc<Mutex<Option<(Child, Sender<()>)>>>,
}
impl Default for Teleport {
fn default() -> Self {
Self {
hostnames: Vec::new(),
current_tunnel: Arc::new(Mutex::new(None)),
}
}
}
impl Teleport {
pub fn update_hostnames(&mut self) -> Result<(), TeleportError> {
let output = Command::new("tsh")
.arg("ls")
.stdout(Stdio::piped())
.output()?;
let output = String::from_utf8(output.stdout)?;
let hostnames = Self::parse_hostnames(&output);
if hostnames.is_empty() {
return Err(TeleportError::NoHostnames);
}
self.hostnames = hostnames;
Ok(())
}
fn parse_hostnames(output: &str) -> Vec<String> {
output
.lines()
.skip(2)
.filter(|line| !line.is_empty())
.filter_map(|line| line.split_whitespace().next())
.map(String::from)
.collect()
}
pub fn get_hostnames(&self) -> &[String] {
&self.hostnames
}
pub fn open_tunnel(&self, hostname: &str) -> Result<(), TeleportError> {
let mut current_tunnel = self.current_tunnel.lock().unwrap();
if current_tunnel.is_some() {
return Err(TeleportError::TunnelAlreadyOpen);
}
if !self.hostnames.contains(&hostname.to_string()) {
return Err(TeleportError::InvalidHostname);
}
println!("Opening tunnel to {}", hostname);
let mut child = Command::new("tsh")
.args(&[
"ssh",
"-L",
"9000:localhost:9000",
&format!("root@{}", hostname),
])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()?;
let stdin = child
.stdin
.take()
.ok_or(TeleportError::StdinCaptureFailed)?;
let stdout = child
.stdout
.take()
.ok_or(TeleportError::StdoutCaptureFailed)?;
let stderr = child
.stderr
.take()
.ok_or(TeleportError::StderrCaptureFailed)?;
let (tx, rx) = mpsc::channel();
let handle = thread::spawn(move || {
let stdout_reader = BufReader::new(stdout);
let stderr_reader = BufReader::new(stderr);
let mut stdin = stdin;
let stdout_handle = thread::spawn(move || {
for line in stdout_reader.lines() {
match line {
Ok(line) => println!("Stdout: {}", line),
Err(e) => eprintln!("Error reading stdout: {}", e),
}
}
});
let stderr_handle = thread::spawn(move || {
for line in stderr_reader.lines() {
match line {
Ok(line) => eprintln!("Stderr: {}", line),
Err(e) => eprintln!("Error reading stderr: {}", e),
}
}
});
rx.recv().ok();
println!("Received notification to close tunnel.");
if let Err(e) = stdin.write_all(b"exit\n") {
eprintln!("Error sending exit command: {}", e);
}
stdout_handle.join().ok();
stderr_handle.join().ok();
});
*current_tunnel = Some((child, tx));
println!("Tunnel opened successfully");
Ok(())
}
pub fn close_tunnel(&self) -> Result<(), TeleportError> {
let mut current_tunnel = self.current_tunnel.lock().unwrap();
if let Some((mut child, tx)) = current_tunnel.take() {
println!("Closing tunnel...");
tx.send(()).ok();
let status = child.wait()?;
println!("Tunnel closed with status: {}", status);
Ok(())
} else {
Err(TeleportError::NoOpenTunnel)
}
}
pub fn is_tunnel_open(&self) -> bool {
let current_tunnel = self.current_tunnel.lock().unwrap();
current_tunnel.is_some()
}
pub fn get_tunnel_status(&self) -> String {
if self.is_tunnel_open() {
"Open".to_string()
} else {
"Closed".to_string()
}
}
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let teleport = Arc::new(Mutex::new(Teleport::default()));
teleport.lock().unwrap().update_hostnames()?;
println!("Available hostnames:");
for (i, hostname) in teleport.lock().unwrap().get_hostnames().iter().enumerate() {
println!("{}. {}", i + 1, hostname);
}
print!("Enter the number of the hostname you want to tunnel to: ");
std::io::stdout().flush()?;
let mut choice = String::new();
std::io::stdin().read_line(&mut choice)?;
let choice: usize = choice
.trim()
.parse()
.map_err(|_| TeleportError::InvalidInput)?;
if choice == 0 || choice > teleport.lock().unwrap().get_hostnames().len() {
return Err(TeleportError::InvalidInput.into());
}
let hostname = teleport.lock().unwrap().get_hostnames()[choice - 1].clone();
teleport.lock().unwrap().open_tunnel(&hostname)?;
println!("Tunnel is now open. Press Enter to close the tunnel.");
println!("Monitoring tunnel status every 5 seconds...");
let teleport_clone = Arc::clone(&teleport);
let status_thread = thread::spawn(move || loop {
thread::sleep(Duration::from_secs(5));
let status = teleport_clone.lock().unwrap().get_tunnel_status();
println!("Tunnel status: {}", status);
if status == "Closed" {
break;
}
});
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
teleport.lock().unwrap().close_tunnel()?;
// Wait for the status thread to finish
status_thread.join().unwrap();
Ok(())
}
Raw
teleport-async.rs
// [dependencies]
// thiserror = "1.0.63"
// tokio = { version = "1.39.2", features = ["full"] }
use std::{process::Stdio, sync::Arc};
use thiserror::Error;
use tokio::{
io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
process::{Child, Command},
sync::{Mutex, Notify},
task::JoinHandle,
};
#[derive(Error, Debug)]
pub enum TeleportError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("UTF-8 conversion error: {0}")]
Utf8(#[from] std::string::FromUtf8Error),
#[error("No hostnames found")]
NoHostnames,
#[error("Invalid hostname")]
InvalidHostname,
#[error("Tunnel already open")]
TunnelAlreadyOpen,
#[error("No tunnel is currently open")]
NoOpenTunnel,
#[error("Task join error: {0}")]
TaskJoin(#[from] tokio::task::JoinError),
#[error("Failed to capture stdout")]
StdoutCaptureFailed,
#[error("Failed to capture stderr")]
StderrCaptureFailed,
#[error("Failed to capture stdin")]
StdinCaptureFailed,
#[error("Invalid input")]
InvalidInput,
}
pub struct Teleport {
hostnames: Vec<String>,
current_tunnel: Arc<Mutex<Option<(Child, JoinHandle<()>, Arc<Notify>)>>>,
}
impl Default for Teleport {
fn default() -> Self {
Self {
hostnames: Vec::new(),
current_tunnel: Arc::new(Mutex::new(None)),
}
}
}
impl Teleport {
pub async fn get_hostnames(&mut self) -> Result<(), TeleportError> {
let output = Command::new("tsh")
.arg("ls")
.stdout(Stdio::piped())
.output()
.await?;
let output = String::from_utf8(output.stdout)?;
let hostnames = Self::parse_hostnames(&output);
if hostnames.is_empty() {
return Err(TeleportError::NoHostnames);
}
self.hostnames = hostnames;
Ok(())
}
fn parse_hostnames(output: &str) -> Vec<String> {
output
.lines()
.skip(2)
.filter(|line| !line.is_empty())
.filter_map(|line| line.split_whitespace().next())
.map(String::from)
.collect()
}
pub fn get_hostnames_list(&self) -> &[String] {
&self.hostnames
}
pub async fn open_tunnel(&self, hostname: &str) -> Result<(), TeleportError> {
let mut current_tunnel = self.current_tunnel.lock().await;
if current_tunnel.is_some() {
return Err(TeleportError::TunnelAlreadyOpen);
}
if !self.hostnames.contains(&hostname.to_string()) {
return Err(TeleportError::InvalidHostname);
}
println!("Opening tunnel to {}", hostname);
let mut child = Command::new("tsh")
.args(&[
"ssh",
"-L",
"9000:localhost:9000",
&format!("root@{}", hostname),
])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()?;
let stdin = child
.stdin
.take()
.ok_or(TeleportError::StdinCaptureFailed)?;
let stdout = child
.stdout
.take()
.ok_or(TeleportError::StdoutCaptureFailed)?;
let stderr = child
.stderr
.take()
.ok_or(TeleportError::StderrCaptureFailed)?;
let notify = Arc::new(Notify::new());
let notify_clone = Arc::clone(&notify);
let handle = tokio::spawn(async move {
let mut stdout_reader = BufReader::new(stdout).lines();
let mut stderr_reader = BufReader::new(stderr).lines();
let mut stdin = stdin;
loop {
tokio::select! {
result = stdout_reader.next_line() => {
match result {
Ok(Some(line)) => println!("Stdout: {}", line),
Ok(None) => break,
Err(e) => eprintln!("Error reading stdout: {}", e),
}
}
result = stderr_reader.next_line() => {
match result {
Ok(Some(line)) => eprintln!("Stderr: {}", line),
Ok(None) => break,
Err(e) => eprintln!("Error reading stderr: {}", e),
}
}
_ = notify_clone.notified() => {
println!("Received notification to close tunnel.");
if let Err(e) = stdin.write_all(b"exit\n").await {
eprintln!("Error sending exit command: {}", e);
}
break;
}
}
}
});
*current_tunnel = Some((child, handle, notify));
println!("Tunnel opened successfully");
Ok(())
}
pub async fn close_tunnel(&self) -> Result<(), TeleportError> {
let mut current_tunnel = self.current_tunnel.lock().await;
if let Some((mut child, handle, notify)) = current_tunnel.take() {
println!("Closing tunnel...");
notify.notify_one();
handle.await?;
let status = child.wait().await?;
println!("Tunnel closed with status: {}", status);
Ok(())
} else {
Err(TeleportError::NoOpenTunnel)
}
}
pub async fn is_tunnel_open(&self) -> bool {
let current_tunnel = self.current_tunnel.lock().await;
current_tunnel.is_some()
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
use std::io::{self, Write};
let mut teleport = Teleport::default();
teleport.get_hostnames().await?;
println!("Available hostnames:");
for (i, hostname) in teleport.get_hostnames_list().iter().enumerate() {
println!("{}. {}", i + 1, hostname);
}
print!("Enter the number of the hostname you want to tunnel to: ");
io::stdout().flush()?;
let mut choice = String::new();
io::stdin().read_line(&mut choice)?;
let choice: usize = choice
.trim()
.parse()
.map_err(|_| TeleportError::InvalidInput)?;
if choice == 0 || choice > teleport.get_hostnames_list().len() {
return Err(TeleportError::InvalidInput.into());
}
let hostname = &teleport.get_hostnames_list()[choice - 1];
teleport.open_tunnel(hostname).await?;
println!("Tunnel is now open. Press Enter to close the tunnel.");
let mut input = String::new();
io::stdin().read_line(&mut input)?;
teleport.close_tunnel().await?;
Ok(())
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment