feat: rewrite to use Extend trait and remove tmux control

10 files changed, 281 insertions, 281 deletions

diff --git a/src/config.rs b/src/config.rs
index 30eac4b..03e962d 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,12 +1,13 @@
 use std::path::PathBuf;
 
-use clap::{Args, Parser, Subcommand};
+use clap::{Args, Parser};
 use tracing::{metadata::LevelFilter, Level};
 
 #[derive(Debug, Clone, Parser)]
 pub struct Config {
-    #[command(subcommand)]
-    pub command: Commands,
+    /// Update the history file from the current sessions
+    #[arg(short, long)]
+    pub update: bool,
 
     /// tmux socket-name, equivelent to `tmux -L <socket-name>`
     #[arg(short = 'L', long, default_value = "ssh")]
@@ -20,17 +21,6 @@ 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 d7e7063..863be8e 100644
--- a/src/history.rs
+++ b/src/history.rs
@@ -1,32 +1,34 @@
 use std::{
     fs::File,
-    io::{BufRead, BufReader, BufWriter, Write},
+    io::{BufRead, BufReader, Write},
     path::PathBuf,
 };
 
 use directories::ProjectDirs;
 
-use crate::{Session, SessionSource};
+use crate::Session;
 
 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),
-    }
-}
+mod error;
 
 #[derive(Debug)]
-pub struct History(PathBuf);
+pub struct History {
+    pub file: File,
+    pub entries: Vec<Session>,
+}
 
 impl History {
-    pub fn new(history_file: PathBuf) -> Self {
-        Self(history_file)
+    pub fn open(history_file: PathBuf) -> Result<Self, std::io::Error> {
+        let file = File::options().write(true).open(&history_file)?;
+
+        let entries = BufReader::new(File::open(history_file)?)
+            .lines()
+            .flatten()
+            .flat_map(|item| ron::from_str(&item))
+            .collect();
+
+        Ok(Self { file, entries })
     }
 
     pub fn default_path() -> Option<PathBuf> {
@@ -35,61 +37,24 @@ 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<Session> = 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 {
-    type Error = std::io::Error;
+impl IntoIterator for History {
+    type Item = Session;
 
-    type Iter = Vec<Session>;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
 
-    fn sessions(&self) -> Result<Self::Iter, Self::Error> {
-        let mut sessions = Vec::new();
-
-        let file = File::open(&self.0)?;
-        for res in BufReader::new(file).lines() {
-            let entry = match res {
-                Ok(entry) => entry,
-                Err(err) => {
-                    tracing::warn!(?err, "Failed to read line from history file");
-                    continue;
-                }
-            };
-
-            match ron::from_str(&entry) {
-                Ok(session) => sessions.push(session),
-                Err(err) => tracing::warn!(?err, "Invalid history entry"),
-            }
-        }
-
-        Ok(sessions)
+    fn into_iter(self) -> Self::IntoIter {
+        self.entries.into_iter()
     }
 }
 
-impl std::ops::Deref for History {
-    type Target = PathBuf;
+impl Write for History {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.file.write(buf)
+    }
 
-    fn deref(&self) -> &Self::Target {
-        &self.0
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.file.flush()
     }
 }
diff --git a/src/history/error.rs b/src/history/error.rs
new file mode 100644
index 0000000..8bdd2d6
--- /dev/null
+++ b/src/history/error.rs
@@ -0,0 +1,14 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("IO error: {0}")]
+    IO(#[from] std::io::Error),
+
+    #[error(transparent)]
+    Serialize(#[from] ron::Error),
+
+    #[error(transparent)]
+    Format(#[from] ron::de::SpannedError),
+
+    #[error(transparent)]
+    Tmux(#[from] crate::tmux::Error),
+}
diff --git a/src/lib.rs b/src/lib.rs
index e58e676..c749b20 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,9 @@
-pub use config::{Commands, Config};
+pub use config::Config;
 pub use history::History;
-pub use session::{Session, SessionSource};
+pub use session::{Session, Sessions, State};
 pub use ssh::KnownHosts;
 pub use tmux::Tmux;
+pub use unix::Hosts;
 
 mod config;
 mod history;
diff --git a/src/main.rs b/src/main.rs
index b546377..28b7bea 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,8 @@
-use std::collections::HashMap;
+use std::io::ErrorKind;
 
 use clap::Parser;
 
-use sshr::{Commands, Config, History, KnownHosts, Session, SessionSource, Tmux};
+use sshr::{Config, History, Hosts, KnownHosts, Sessions, Tmux};
 
 fn main() -> anyhow::Result<()> {
     let mut config = Config::parse();
@@ -14,44 +14,42 @@ fn main() -> anyhow::Result<()> {
 
     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),
-    }
-}
+    let history = match config.history_file {
+        Some(path) => History::open(path),
+        None => Err(std::io::Error::from(ErrorKind::NotFound)),
+    };
+
+    let mut sessions = Sessions::new();
 
-fn list_sessions(config: Config) -> anyhow::Result<()> {
-    let mut sessions = HashMap::new();
+    match Tmux::new(config.socket_name).list(None) {
+        Ok(p) => sessions.extend(p),
+        Err(err) => tracing::warn!(?err, "Failed to list tmux sessions"),
+    };
 
-    sessions = KnownHosts::default().update(sessions)?;
-    sessions = Tmux::new(config.socket_name).update(sessions)?;
+    match history {
+        Ok(p) => {
+            sessions.extend(p.entries);
 
-    if let Some(history) = config.history_file.map(History::new) {
-        if history.exists() {
-            sessions = history.update(sessions)?;
+            if config.update {
+                sessions.write_into(p.file)?;
+            }
         }
-    }
+        Err(err) => tracing::warn!(?err, "Failed to open history file"),
+    };
 
-    let mut sessions: Vec<Session> = sessions.into_values().collect();
+    match KnownHosts::open() {
+        Ok(p) => sessions.extend(p),
+        Err(err) => tracing::warn!(?err, "Failed to read KnownHost file"),
+    };
 
-    sessions.sort();
+    match Hosts::open() {
+        Ok(p) => sessions.extend(p),
+        Err(err) => tracing::warn!(?err, "Failed to read /etc/hosts"),
+    };
 
-    for session in sessions {
+    for session in sessions.sorted() {
         println!("{session}");
     }
 
     Ok(())
 }
-
-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 fd5137e..360e2e1 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -1,52 +1,117 @@
 use std::{
-    collections::{hash_map::Entry, HashMap},
+    collections::{hash_map::Entry, HashMap, HashSet},
     fmt::Display,
+    io::{BufWriter, Write},
     iter::IntoIterator,
-    time::{Duration, SystemTime, UNIX_EPOCH},
+    time::Duration,
 };
 
 use serde::{Deserialize, Serialize};
 
-pub trait SessionSource {
-    type Error;
-
-    type Iter: IntoIterator<Item = Session>;
-
-    fn sessions(&self) -> Result<Self::Iter, Self::Error>;
-
-    fn update(
-        &self,
-        mut hash_map: HashMap<String, Session>,
-    ) -> Result<HashMap<String, Session>, Self::Error> {
-        for session in self.sessions()? {
-            match hash_map.entry(session.name.to_owned()) {
-                Entry::Occupied(mut o) if &session > o.get() => {
-                    o.insert(session);
-                }
-                Entry::Vacant(v) => {
-                    v.insert(session);
-                }
-                _ => {}
+#[derive(Debug, Default)]
+pub struct Sessions {
+    inner: HashMap<String, State>,
+    hostname: Option<String>,
+}
+
+impl Sessions {
+    pub fn new() -> Self {
+        let hostname = match hostname::get() {
+            Ok(h) => Some(h.to_string_lossy().into()),
+            Err(err) => {
+                tracing::error!(%err, "Failed to get hostname to filter output.");
+                None
+            }
+        };
+
+        Self {
+            hostname,
+            ..Default::default()
+        }
+    }
+
+    pub fn sorted(self) -> Vec<Session> {
+        let mut sessions: Vec<Session> = self.into_iter().map(Session::from).collect();
+        sessions.sort();
+        sessions
+    }
+
+    pub fn write_into<W: Write>(&self, writer: W) -> std::io::Result<()> {
+        let mut buf_writer = BufWriter::new(writer);
+
+        for session in self.inner.iter().map(Session::from) {
+            match ron::to_string(&session) {
+                Ok(ser) => writeln!(buf_writer, "{ser}")?,
+                Err(err) => tracing::warn!(%err, "Failed to serialize session"),
             }
         }
-        Ok(hash_map)
+
+        Ok(())
+    }
+
+    fn add(&mut self, item: Session) {
+        let span = tracing::trace_span!("Entry", ?item);
+        let _guard = span.enter();
+
+        if self
+            .hostname
+            .as_ref()
+            .map(|h| h == &item.name)
+            .unwrap_or_default()
+        {
+            return;
+        }
+
+        match self.inner.entry(item.name) {
+            Entry::Occupied(mut occupied) if &item.state > occupied.get() => {
+                tracing::trace!(?occupied, new_value=?item.state, "New entry is more recent, replacing");
+                occupied.insert(item.state);
+            }
+            Entry::Occupied(occupied) => {
+                tracing::trace!(?occupied, new_value=?item.state, "Previous entry is more recent, skipping");
+            }
+            Entry::Vacant(v) => {
+                tracing::trace!(?item.state, "No previous entry exists, inserting");
+                v.insert(item.state);
+            }
+        }
+    }
+}
+
+impl IntoIterator for Sessions {
+    type Item = Session;
+
+    type IntoIter = std::collections::hash_set::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.inner
+            .into_iter()
+            .map(Into::into)
+            .collect::<HashSet<Session>>()
+            .into_iter()
+    }
+}
+
+impl Extend<Session> for Sessions {
+    fn extend<T: IntoIterator<Item = Session>>(&mut self, iter: T) {
+        for item in iter {
+            self.add(item)
+        }
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub enum State {
     Discovered,
-    Connected,
-    Updated,
+    #[serde(with = "epoch_timestamp")]
+    Created(Duration),
+    #[serde(with = "epoch_timestamp")]
+    Attached(Duration),
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 pub struct Session {
     pub state: State,
-
-    #[serde(with = "epoch_timestamp")]
-    pub timestamp: Duration,
-
     pub name: String,
 }
 
@@ -76,15 +141,19 @@ impl Display for Session {
     }
 }
 
+impl From<&str> for Session {
+    fn from(value: &str) -> Self {
+        Self {
+            state: State::Discovered,
+            name: value.to_owned(),
+        }
+    }
+}
+
 impl From<String> for Session {
     fn from(name: String) -> Self {
-        let timestamp = SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .expect("Current time is pre-epoch. (Time traveler?)");
-
         Self {
             state: State::Discovered,
-            timestamp,
             name,
         }
     }
@@ -92,14 +161,15 @@ impl From<String> 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, name }
+    }
+}
 
+impl From<(&String, &State)> for Session {
+    fn from((name, &state): (&String, &State)) -> Self {
         Self {
             state,
-            timestamp,
-            name,
+            name: name.to_owned(),
         }
     }
 }
diff --git a/src/ssh.rs b/src/ssh.rs
index cd8c61a..3508a36 100644
--- a/src/ssh.rs
+++ b/src/ssh.rs
@@ -1,49 +1,41 @@
 use std::{
-    collections::HashSet,
     fs::File,
-    io::{BufRead, BufReader},
+    io::{BufRead, BufReader, Error, ErrorKind},
 };
 
 use directories::UserDirs;
 
-use crate::{Session, SessionSource};
+use crate::Session;
 
-pub struct KnownHosts;
+pub struct KnownHosts(Vec<Session>);
 
 impl KnownHosts {
-    pub fn new() -> Self {
-        Self
-    }
-}
-
-impl SessionSource for KnownHosts {
-    type Error = std::io::Error;
-
-    type Iter = Vec<Session>;
-
-    fn sessions(&self) -> Result<Self::Iter, Self::Error> {
+    pub fn open() -> Result<Self, Error> {
         let path = UserDirs::new()
-            .ok_or(std::io::ErrorKind::NotFound)?
+            .ok_or(ErrorKind::NotFound)?
             .home_dir()
             .join(".ssh/known_hosts");
 
-        let reader = BufReader::new(File::open(path)?);
+        let file = File::open(path)?;
 
-        let sessions: Vec<Session> = reader
+        let inner = BufReader::new(file)
             .lines()
             .flatten()
-            .filter_map(|s| s.split_whitespace().next().map(str::to_owned))
-            .collect::<HashSet<String>>()
-            .into_iter()
+            .take_while(|s| !s.starts_with('#'))
+            .filter_map(|l| l.split_whitespace().next().map(str::to_owned))
             .map(Session::from)
             .collect();
 
-        Ok(sessions)
+        Ok(Self(inner))
     }
 }
 
-impl Default for KnownHosts {
-    fn default() -> Self {
-        Self::new()
+impl IntoIterator for KnownHosts {
+    type Item = Session;
+
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
     }
 }
diff --git a/src/tmux.rs b/src/tmux.rs
index 59420e7..fee0f49 100644
--- a/src/tmux.rs
+++ b/src/tmux.rs
@@ -1,6 +1,6 @@
-use std::process::{Command, ExitStatus, Output};
+use std::process::Command;
 
-use crate::{Session, SessionSource};
+use crate::Session;
 
 pub use error::Error;
 
@@ -12,108 +12,28 @@ pub struct Tmux {
 }
 
 impl Tmux {
-    const SESSION_FORMAT: &str = r##"Session(name: "#S", #{?session_last_attached,timestamp: #{session_last_attached}#, state: Updated,timestamp: #{session_created}#, state: Connected})"##;
+    const SESSION_FORMAT: &str = r##"Session(name: "#S", state: #{?session_last_attached,Attached(#{session_last_attached}),Created(#{session_created})})"##;
 
     pub fn new(socket_name: String) -> Self {
         Self { socket_name }
     }
 
-    pub fn show(&self, name: &String) -> Result<Session, Error> {
-        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)?;
+    pub fn list(&self, name: Option<String>) -> Result<Vec<Session>, Error> {
+        let filter = name
+            .map(|s| vec!["-f".into(), format!("#{{==:#S,{s}}}")])
+            .unwrap_or_default();
 
-        Ok(ron::from_str(output_str)?)
-    }
-
-    pub fn list(&self) -> Result<Output, Error> {
-        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<ExitStatus, Error> {
-        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<ExitStatus, Error> {
-        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<ExitStatus, Error> {
-        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<ExitStatus, Error> {
-        Ok(Command::new("tmux")
-            .arg("detach")
-            .arg("-E")
-            .arg(format!("tmux -L ssh attach -t {host}"))
-            .status()?)
-    }
-}
-
-impl SessionSource for Tmux {
-    type Error = Error;
-
-    type Iter = Vec<Session>;
-
-    fn sessions(&self) -> Result<Self::Iter, Self::Error> {
-        let output = Command::new("tmux")
+        let stdout = Command::new("tmux")
             .arg("-L")
             .arg(&self.socket_name)
             .arg("list-sessions")
+            .args(filter)
             .arg("-F")
             .arg(Self::SESSION_FORMAT)
             .output()?
             .stdout;
 
-        let sessions = std::str::from_utf8(&output)?
+        let sessions = std::str::from_utf8(&stdout)?
             .lines()
             .filter_map(|s| match ron::from_str(s) {
                 Ok(session) => Some(session),
@@ -127,3 +47,38 @@ impl SessionSource for Tmux {
         Ok(sessions)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const SOCKET_NAME: &str = "test";
+
+    #[test]
+    fn test_tmux_list() -> Result<(), Error> {
+        let names = Vec::from(["test_1", "test_2", "test_3", "test_4"]);
+
+        for name in names.iter().cloned() {
+            Command::new("tmux")
+                .arg("-L")
+                .arg(SOCKET_NAME)
+                .arg("new-session")
+                .arg("-ds")
+                .arg(name)
+                .status()?;
+        }
+
+        let tmux = Tmux::new(SOCKET_NAME.to_owned());
+        let sessions: Vec<_> = tmux.list(None)?.into_iter().map(|s| s.name).collect();
+
+        Command::new("tmux")
+            .arg("-L")
+            .arg(SOCKET_NAME)
+            .arg("kill-server")
+            .status()?;
+
+        assert_eq!(names, sessions);
+
+        Ok(())
+    }
+}
diff --git a/src/tmux/error.rs b/src/tmux/error.rs
index e6e12dc..6da5830 100644
--- a/src/tmux/error.rs
+++ b/src/tmux/error.rs
@@ -8,7 +8,4 @@ pub enum Error {
 
     #[error("Parsing error: {0}")]
     Parse(#[from] std::str::Utf8Error),
-
-    #[error("History file error: {0}")]
-    History(#[from] crate::history::Error),
 }
diff --git a/src/unix.rs b/src/unix.rs
index 5414125..402e476 100644
--- a/src/unix.rs
+++ b/src/unix.rs
@@ -1,19 +1,37 @@
-use crate::{Session, SessionSource};
+use std::{
+    fs::File,
+    io::{BufRead, BufReader, Error},
+};
 
-pub struct HostFile;
+use crate::Session;
 
-impl SessionSource for HostFile {
-    type Error = &'static str;
+pub struct Hosts(Vec<Session>);
 
-    type Iter = Vec<Session>;
-
-    fn sessions(&self) -> Result<Self::Iter, Self::Error> {
-        let sessions = hostfile::parse_hostfile()?
-            .into_iter()
-            .flat_map(|h| h.names.into_iter())
-            .map(Into::into)
+impl Hosts {
+    pub fn open() -> Result<Self, Error> {
+        let buf_reader = BufReader::new(File::open("/etc/hosts")?);
+        let inner: Vec<Session> = buf_reader
+            .lines()
+            .flatten()
+            .take_while(|s| !s.starts_with('#'))
+            .flat_map(|l| {
+                l.split_whitespace()
+                    .skip(1)
+                    .take_while(|s| !s.starts_with('#'))
+                    .map(Session::from)
+                    .collect::<Vec<_>>()
+            })
             .collect();
+        Ok(Self(inner))
+    }
+}
+
+impl IntoIterator for Hosts {
+    type Item = Session;
+
+    type IntoIter = std::vec::IntoIter<Self::Item>;
 
-        Ok(sessions)
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
     }
 }