feat: add exclude flag and localhost session

4 files changed, 68 insertions, 21 deletions

diff --git a/src/config.rs b/src/config.rs
index 03e962d..63dfbc7 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -13,6 +13,10 @@ pub struct Config {
     #[arg(short = 'L', long, default_value = "ssh")]
     pub socket_name: String,
 
+    /// Name of host to exclude
+    #[arg(short, long)]
+    pub exclude: Vec<String>,
+
     /// path to history file [default: $XDG_DATA_HOME/sshr/history]
     #[arg(short = 'f', long)]
     pub history_file: Option<PathBuf>,
diff --git a/src/main.rs b/src/main.rs
index 28b7bea..eb62d52 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,9 +19,16 @@ fn main() -> anyhow::Result<()> {
         None => Err(std::io::Error::from(ErrorKind::NotFound)),
     };
 
-    let mut sessions = Sessions::new();
+    let tmux = Tmux::new(config.socket_name);
 
-    match Tmux::new(config.socket_name).list(None) {
+    let mut sessions = Sessions::new(config.exclude);
+
+    match tmux.host() {
+        Ok(p) => sessions.add(p),
+        Err(err) => tracing::warn!(?err, "Failed to get tmux host"),
+    };
+
+    match tmux.list(None) {
         Ok(p) => sessions.extend(p),
         Err(err) => tracing::warn!(?err, "Failed to list tmux sessions"),
     };
diff --git a/src/session.rs b/src/session.rs
index 360e2e1..a7ed261 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -11,21 +11,13 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Default)]
 pub struct Sessions {
     inner: HashMap<String, State>,
-    hostname: Option<String>,
+    exclude: Vec<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
-            }
-        };
-
+    pub fn new(exclude: Vec<String>) -> Self {
         Self {
-            hostname,
+            exclude,
             ..Default::default()
         }
     }
@@ -49,16 +41,12 @@ impl Sessions {
         Ok(())
     }
 
-    fn add(&mut self, item: Session) {
+    pub 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()
-        {
+        if self.exclude.contains(&item.name) {
+            tracing::debug!(item.name, "Skipping excluded item");
             return;
         }
 
@@ -107,6 +95,7 @@ pub enum State {
     Created(Duration),
     #[serde(with = "epoch_timestamp")]
     Attached(Duration),
+    LocalHost,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
diff --git a/src/tmux.rs b/src/tmux.rs
index fee0f49..88f521f 100644
--- a/src/tmux.rs
+++ b/src/tmux.rs
@@ -33,7 +33,7 @@ impl Tmux {
             .output()?
             .stdout;
 
-        let sessions = std::str::from_utf8(&stdout)?
+        let sessions: Vec<Session> = std::str::from_utf8(&stdout)?
             .lines()
             .filter_map(|s| match ron::from_str(s) {
                 Ok(session) => Some(session),
@@ -46,6 +46,24 @@ impl Tmux {
 
         Ok(sessions)
     }
+
+    pub fn host(&self) -> Result<Session, Error> {
+        let stdout = Command::new("tmux")
+            .arg("-L")
+            .arg(&self.socket_name)
+            .arg("display")
+            .arg("-p")
+            .arg("#h")
+            .output()?
+            .stdout;
+
+        let name = std::str::from_utf8(&stdout)?.trim().into();
+
+        Ok(Session {
+            state: crate::State::LocalHost,
+            name,
+        })
+    }
 }
 
 #[cfg(test)]
@@ -81,4 +99,33 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_tmux_host() -> Result<(), Error> {
+        Command::new("tmux")
+            .arg("-L")
+            .arg(SOCKET_NAME)
+            .arg("new-session")
+            .arg("-d")
+            .status()?;
+
+        let name = hostname::get()?.to_string_lossy().into();
+        let expected_session = Session {
+            state: crate::State::LocalHost,
+            name,
+        };
+
+        let tmux = Tmux::new(SOCKET_NAME.to_owned());
+        let session = tmux.host()?;
+
+        Command::new("tmux")
+            .arg("-L")
+            .arg(SOCKET_NAME)
+            .arg("kill-server")
+            .status()?;
+
+        assert_eq!(expected_session, session);
+
+        Ok(())
+    }
 }