feat: add --on-session-start-template and `--on-session-end-templat…

…e` options Closes #98 Signed-off-by: Nathanael DEMACON <[email protected]>
quantumsheep · Aug 14, 2024 · c9c7e5f · c9c7e5f
1 parent 59d70cb
commit c9c7e5f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -38,6 +38,14 @@ struct Args {
     #[arg(short, long, default_value = "ssh \"{{{name}}}\"")]
     template: String,
 
+    /// Handlebars template of the command to execute when an SSH session starts
+    #[arg(long, value_name = "TEMPLATE")]
+    on_session_start_template: Option<String>,
+
+    /// Handlebars template of the command to execute when an SSH session ends
+    #[arg(long, value_name = "TEMPLATE")]
+    on_session_end_template: Option<String>,
+
     /// Exit after ending the SSH session
     #[arg(short, long, default_value_t = false)]
     exit: bool,
@@ -52,7 +60,9 @@ fn main() -> Result<()> {
         sort_by_name: args.sort,
         show_proxy_command: args.show_proxy_command,
         command_template: args.template,
-        exit_after_ssh: args.exit,
+        command_template_on_session_start: args.on_session_start_template,
+        command_template_on_session_end: args.on_session_end_template,
+        exit_after_ssh_session_ends: args.exit,
     })?;
     app.start()?;
 

diff --git a/src/ui.rs b/src/ui.rs
@@ -35,7 +35,9 @@ pub struct AppConfig {
     pub show_proxy_command: bool,
 
     pub command_template: String,
-    pub exit_after_ssh: bool,
+    pub command_template_on_session_start: Option<String>,
+    pub command_template_on_session_end: Option<String>,
+    pub exit_after_ssh_session_ends: bool,
 }
 
 pub struct App {
@@ -223,11 +225,19 @@ impl App {
 
                 restore_terminal(terminal).expect("Failed to restore terminal");
 
+                if let Some(template) = &self.config.command_template_on_session_start {
+                    host.run_command_template(template)?;
+                }
+
                 host.run_command_template(&self.config.command_template)?;
 
+                if let Some(template) = &self.config.command_template_on_session_end {
+                    host.run_command_template(template)?;
+                }
+
                 setup_terminal(terminal).expect("Failed to setup terminal");
 
-                if self.config.exit_after_ssh {
+                if self.config.exit_after_ssh_session_ends {
                     return Ok(AppKeyAction::Stop);
                 }
             }