From e11e8bbf14be8f84b57013e4ace1e61071853c12 Mon Sep 17 00:00:00 2001 From: Toby Vincent Date: Mon, 20 Mar 2023 01:58:49 -0500 Subject: feat: impl tmux session switching writing to history file --- src/config.rs | 19 +++++++++---- src/history.rs | 46 ++++++++++++++++++++++++++++--- src/lib.rs | 2 +- src/main.rs | 32 ++++++++++++++-------- src/session.rs | 13 +++++++++ src/tmux.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++++++- src/tmux/error.rs | 3 +++ 7 files changed, 175 insertions(+), 21 deletions(-) (limited to 'src') diff --git a/src/config.rs b/src/config.rs index 7745e0b..30eac4b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,14 +1,12 @@ use std::path::PathBuf; -use clap::{Args, Parser}; +use clap::{Args, Parser, Subcommand}; use tracing::{metadata::LevelFilter, Level}; #[derive(Debug, Clone, Parser)] pub struct Config { - pub target: Option, - - #[arg(short, long)] - pub list: bool, + #[command(subcommand)] + pub command: Commands, /// tmux socket-name, equivelent to `tmux -L ` #[arg(short = 'L', long, default_value = "ssh")] @@ -22,6 +20,17 @@ pub struct Config { pub verbosity: Verbosity, } +#[derive(Debug, Clone, Subcommand)] +pub enum Commands { + /// List available session targets + List, + /// Switch to a specific session, or create it if it does not exist + Switch { + /// Target session name or ssh host + name: String, + }, +} + #[derive(Debug, Default, Clone, Args)] pub struct Verbosity { /// Print additional information per occurrence. diff --git a/src/history.rs b/src/history.rs index 9d29468..d7e7063 100644 --- a/src/history.rs +++ b/src/history.rs @@ -1,14 +1,25 @@ use std::{ fs::File, - io::{BufRead, BufReader}, + io::{BufRead, BufReader, BufWriter, Write}, path::PathBuf, }; use directories::ProjectDirs; -use crate::Session; +use crate::{Session, SessionSource}; -use super::SessionSource; +pub use error::Error; + +mod error { + #[derive(Debug, thiserror::Error)] + pub enum Error { + #[error("IO error: {0}")] + IO(#[from] std::io::Error), + + #[error(transparent)] + Format(#[from] ron::Error), + } +} #[derive(Debug)] pub struct History(PathBuf); @@ -24,6 +35,27 @@ impl History { .join("history") .into() } + + pub fn update_session(&self, session: Session) -> Result<(), Error> { + std::fs::create_dir_all(&self.0)?; + + let mut sessions: Vec = self + .sessions()? + .into_iter() + .filter(|s| s.name == session.name) + .collect(); + + sessions.push(session); + + let mut writer = BufWriter::new(File::create(&self.0)?); + + Ok(sessions + .into_iter() + .try_for_each(|s| match ron::to_string(&s) { + Ok(ser) => writeln!(writer, "{ser}"), + Err(err) => Ok(tracing::warn!(?err, "Failed to re-serialize session")), + })?) + } } impl SessionSource for History { @@ -53,3 +85,11 @@ impl SessionSource for History { Ok(sessions) } } + +impl std::ops::Deref for History { + type Target = PathBuf; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} diff --git a/src/lib.rs b/src/lib.rs index 0dcdf7b..e58e676 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,4 @@ -pub use config::Config; +pub use config::{Commands, Config}; pub use history::History; pub use session::{Session, SessionSource}; pub use ssh::KnownHosts; diff --git a/src/main.rs b/src/main.rs index 763194c..b546377 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,20 +2,21 @@ use std::collections::HashMap; use clap::Parser; -use sshr::{Config, History, KnownHosts, Session, SessionSource, Tmux}; +use sshr::{Commands, Config, History, KnownHosts, Session, SessionSource, Tmux}; fn main() -> anyhow::Result<()> { - let config = Config::parse(); + let mut config = Config::parse(); tracing_subscriber::fmt::fmt() .with_max_level(&config.verbosity) .without_time() .init(); - if config.list { - list_sessions(config) - } else { - switch(config.target) + config.history_file = config.history_file.or_else(History::default_path); + + match config.command.to_owned() { + Commands::List => list_sessions(config), + Commands::Switch { name } => switch(config, name), } } @@ -25,9 +26,9 @@ fn list_sessions(config: Config) -> anyhow::Result<()> { sessions = KnownHosts::default().update(sessions)?; sessions = Tmux::new(config.socket_name).update(sessions)?; - if let Some(history_file) = config.history_file.or_else(History::default_path) { - if history_file.exists() { - sessions = History::new(history_file).update(sessions)?; + if let Some(history) = config.history_file.map(History::new) { + if history.exists() { + sessions = history.update(sessions)?; } } @@ -42,6 +43,15 @@ fn list_sessions(config: Config) -> anyhow::Result<()> { Ok(()) } -fn switch(_target: Option) -> anyhow::Result<()> { - todo!() +fn switch(config: Config, name: String) -> anyhow::Result<()> { + let tmux = Tmux::new(config.socket_name); + + if let Some(history) = config.history_file.map(History::new) { + if tmux.switch(&name)?.success() { + let session = tmux.show(&name)?; + history.update_session(session)?; + } + } + + Ok(()) } diff --git a/src/session.rs b/src/session.rs index ea3cc75..fd5137e 100644 --- a/src/session.rs +++ b/src/session.rs @@ -90,3 +90,16 @@ impl From for Session { } } +impl From<(String, State)> for Session { + fn from((name, state): (String, State)) -> Self { + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Current time is pre-epoch. (Time traveler?)"); + + Self { + state, + timestamp, + name, + } + } +} diff --git a/src/tmux.rs b/src/tmux.rs index 6c53831..59420e7 100644 --- a/src/tmux.rs +++ b/src/tmux.rs @@ -1,4 +1,4 @@ -use std::process::Command; +use std::process::{Command, ExitStatus, Output}; use crate::{Session, SessionSource}; @@ -17,6 +17,85 @@ impl Tmux { pub fn new(socket_name: String) -> Self { Self { socket_name } } + + pub fn show(&self, name: &String) -> Result { + let output = Command::new("tmux") + .arg("-L") + .arg(&self.socket_name) + .arg("display") + .arg("-p") + .arg("-t") + .arg(name) + .arg(Self::SESSION_FORMAT) + .output()? + .stdout; + + let output_str = std::str::from_utf8(&output)?; + + Ok(ron::from_str(output_str)?) + } + + pub fn list(&self) -> Result { + Ok(Command::new("tmux") + .arg("-L") + .arg(&self.socket_name) + .arg("list-sessions") + .arg("-F") + .arg(Self::SESSION_FORMAT) + .output()?) + } + + pub fn switch(&self, host: &String) -> Result { + if Command::new("tmux") + .arg("-L") + .arg(&self.socket_name) + .arg("has-session") + .arg("-t") + .arg(host) + .status()? + .success() + { + self.create(host)?; + } + + if std::env::var("TMUX").is_ok() { + self.attach(host) + } else { + self.detach_then_attach(host) + } + } + + fn create(&self, host: &String) -> Result { + Ok(Command::new("tmux") + .arg("-L") + .arg(&self.socket_name) + .arg("new-session") + .arg("-ds") + .arg(host) + .arg("ssh") + .arg("-t") + .arg(host) + .arg("zsh -l -c 'tmux new -A'") + .status()?) + } + + fn attach(&self, host: &String) -> Result { + Ok(Command::new("tmux") + .arg("-L") + .arg(&self.socket_name) + .arg("attach-session") + .arg("-t") + .arg(host) + .status()?) + } + + fn detach_then_attach(&self, host: &String) -> Result { + Ok(Command::new("tmux") + .arg("detach") + .arg("-E") + .arg(format!("tmux -L ssh attach -t {host}")) + .status()?) + } } impl SessionSource for Tmux { diff --git a/src/tmux/error.rs b/src/tmux/error.rs index 6da5830..e6e12dc 100644 --- a/src/tmux/error.rs +++ b/src/tmux/error.rs @@ -8,4 +8,7 @@ pub enum Error { #[error("Parsing error: {0}")] Parse(#[from] std::str::Utf8Error), + + #[error("History file error: {0}")] + History(#[from] crate::history::Error), } -- cgit v1.2.3-70-g09d2