Skip to main content

purple_ssh/
snippet.rs

1use std::io;
2use std::path::{Path, PathBuf};
3use std::process::{Command, ExitStatus, Stdio};
4
5use crate::fs_util;
6
7/// A saved command snippet.
8#[derive(Debug, Clone, PartialEq)]
9pub struct Snippet {
10    pub name: String,
11    pub command: String,
12    pub description: String,
13}
14
15/// Result of running a snippet on a host.
16pub struct SnippetResult {
17    pub status: ExitStatus,
18    pub stdout: String,
19    pub stderr: String,
20}
21
22/// Snippet storage backed by ~/.purple/snippets (INI-style).
23#[derive(Debug, Clone, Default)]
24pub struct SnippetStore {
25    pub snippets: Vec<Snippet>,
26    /// Override path for save(). None uses the default ~/.purple/snippets.
27    pub path_override: Option<PathBuf>,
28}
29
30fn config_path() -> Option<PathBuf> {
31    dirs::home_dir().map(|h| h.join(".purple/snippets"))
32}
33
34impl SnippetStore {
35    /// Load snippets from ~/.purple/snippets.
36    /// Returns empty store if file doesn't exist (normal first-use).
37    pub fn load() -> Self {
38        let path = match config_path() {
39            Some(p) => p,
40            None => return Self::default(),
41        };
42        let content = match std::fs::read_to_string(&path) {
43            Ok(c) => c,
44            Err(e) if e.kind() == io::ErrorKind::NotFound => return Self::default(),
45            Err(e) => {
46                eprintln!("! Could not read {}: {}", path.display(), e);
47                return Self::default();
48            }
49        };
50        Self::parse(&content)
51    }
52
53    /// Parse INI-style snippet config.
54    pub fn parse(content: &str) -> Self {
55        let mut snippets = Vec::new();
56        let mut current: Option<Snippet> = None;
57
58        for line in content.lines() {
59            let trimmed = line.trim();
60            if trimmed.is_empty() || trimmed.starts_with('#') {
61                continue;
62            }
63            if trimmed.starts_with('[') && trimmed.ends_with(']') {
64                if let Some(snippet) = current.take() {
65                    if !snippet.command.is_empty()
66                        && !snippets.iter().any(|s: &Snippet| s.name == snippet.name)
67                    {
68                        snippets.push(snippet);
69                    }
70                }
71                let name = trimmed[1..trimmed.len() - 1].trim().to_string();
72                if snippets.iter().any(|s| s.name == name) {
73                    current = None;
74                    continue;
75                }
76                current = Some(Snippet {
77                    name,
78                    command: String::new(),
79                    description: String::new(),
80                });
81            } else if let Some(ref mut snippet) = current {
82                if let Some((key, value)) = trimmed.split_once('=') {
83                    let key = key.trim();
84                    let value = value.trim().to_string();
85                    match key {
86                        "command" => snippet.command = value,
87                        "description" => snippet.description = value,
88                        _ => {}
89                    }
90                }
91            }
92        }
93        if let Some(snippet) = current {
94            if !snippet.command.is_empty()
95                && !snippets.iter().any(|s| s.name == snippet.name)
96            {
97                snippets.push(snippet);
98            }
99        }
100        Self {
101            snippets,
102            path_override: None,
103        }
104    }
105
106    /// Save snippets to ~/.purple/snippets (atomic write, chmod 600).
107    pub fn save(&self) -> io::Result<()> {
108        let path = match &self.path_override {
109            Some(p) => p.clone(),
110            None => match config_path() {
111                Some(p) => p,
112                None => {
113                    return Err(io::Error::new(
114                        io::ErrorKind::NotFound,
115                        "Could not determine home directory",
116                    ))
117                }
118            },
119        };
120
121        let mut content = String::new();
122        for (i, snippet) in self.snippets.iter().enumerate() {
123            if i > 0 {
124                content.push('\n');
125            }
126            content.push_str(&format!("[{}]\n", snippet.name));
127            content.push_str(&format!("command={}\n", snippet.command));
128            if !snippet.description.is_empty() {
129                content.push_str(&format!("description={}\n", snippet.description));
130            }
131        }
132
133        fs_util::atomic_write(&path, content.as_bytes())
134    }
135
136    /// Get a snippet by name.
137    pub fn get(&self, name: &str) -> Option<&Snippet> {
138        self.snippets.iter().find(|s| s.name == name)
139    }
140
141    /// Add or replace a snippet.
142    pub fn set(&mut self, snippet: Snippet) {
143        if let Some(existing) = self.snippets.iter_mut().find(|s| s.name == snippet.name) {
144            *existing = snippet;
145        } else {
146            self.snippets.push(snippet);
147        }
148    }
149
150    /// Remove a snippet by name.
151    pub fn remove(&mut self, name: &str) {
152        self.snippets.retain(|s| s.name != name);
153    }
154}
155
156/// Validate a snippet name: non-empty, no whitespace, no `#`, no `[`, no `]`,
157/// no control characters.
158pub fn validate_name(name: &str) -> Result<(), String> {
159    if name.is_empty() {
160        return Err("Snippet name cannot be empty.".to_string());
161    }
162    if name.contains(char::is_whitespace) {
163        return Err("Snippet name cannot contain whitespace.".to_string());
164    }
165    if name.contains('#') || name.contains('[') || name.contains(']') {
166        return Err("Snippet name cannot contain #, [ or ].".to_string());
167    }
168    if name.contains(|c: char| c.is_control()) {
169        return Err("Snippet name cannot contain control characters.".to_string());
170    }
171    Ok(())
172}
173
174/// Validate a snippet command: non-empty, no control characters (except tab).
175pub fn validate_command(command: &str) -> Result<(), String> {
176    if command.trim().is_empty() {
177        return Err("Command cannot be empty.".to_string());
178    }
179    if command.contains(|c: char| c.is_control() && c != '\t') {
180        return Err("Command cannot contain control characters.".to_string());
181    }
182    Ok(())
183}
184
185/// Run a snippet on a single host via SSH.
186/// When `capture` is true, stdout/stderr are piped and returned in the result.
187/// When `capture` is false, stdout/stderr are inherited (streamed to terminal
188/// in real-time) and the returned strings are empty.
189pub fn run_snippet(
190    alias: &str,
191    config_path: &Path,
192    command: &str,
193    askpass: Option<&str>,
194    bw_session: Option<&str>,
195    capture: bool,
196    has_active_tunnel: bool,
197) -> anyhow::Result<SnippetResult> {
198    let mut cmd = Command::new("ssh");
199    cmd.arg("-F")
200        .arg(config_path)
201        .arg("-o")
202        .arg("ConnectTimeout=10");
203
204    // When a tunnel is already running for this host, disable forwards
205    // to avoid "Address already in use" bind conflicts.
206    if has_active_tunnel {
207        cmd.arg("-o").arg("ClearAllForwardings=yes");
208    }
209
210    cmd.arg("--")
211        .arg(alias)
212        .arg(command)
213        .stdin(Stdio::inherit());
214
215    if capture {
216        cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
217    } else {
218        cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
219    }
220
221    if askpass.is_some() {
222        let exe = std::env::current_exe()
223            .ok()
224            .map(|p| p.to_string_lossy().to_string())
225            .or_else(|| std::env::args().next())
226            .unwrap_or_else(|| "purple".to_string());
227        cmd.env("SSH_ASKPASS", &exe)
228            .env("SSH_ASKPASS_REQUIRE", "prefer")
229            .env("PURPLE_ASKPASS_MODE", "1")
230            .env("PURPLE_HOST_ALIAS", alias)
231            .env("PURPLE_CONFIG_PATH", config_path.as_os_str());
232    }
233
234    if let Some(token) = bw_session {
235        cmd.env("BW_SESSION", token);
236    }
237
238    if capture {
239        let output = cmd
240            .output()
241            .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
242
243        Ok(SnippetResult {
244            status: output.status,
245            stdout: String::from_utf8_lossy(&output.stdout).to_string(),
246            stderr: String::from_utf8_lossy(&output.stderr).to_string(),
247        })
248    } else {
249        let status = cmd
250            .status()
251            .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
252
253        Ok(SnippetResult {
254            status,
255            stdout: String::new(),
256            stderr: String::new(),
257        })
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    // =========================================================================
266    // Parse
267    // =========================================================================
268
269    #[test]
270    fn test_parse_empty() {
271        let store = SnippetStore::parse("");
272        assert!(store.snippets.is_empty());
273    }
274
275    #[test]
276    fn test_parse_single_snippet() {
277        let content = "\
278[check-disk]
279command=df -h
280description=Check disk usage
281";
282        let store = SnippetStore::parse(content);
283        assert_eq!(store.snippets.len(), 1);
284        let s = &store.snippets[0];
285        assert_eq!(s.name, "check-disk");
286        assert_eq!(s.command, "df -h");
287        assert_eq!(s.description, "Check disk usage");
288    }
289
290    #[test]
291    fn test_parse_multiple_snippets() {
292        let content = "\
293[check-disk]
294command=df -h
295
296[uptime]
297command=uptime
298description=Check server uptime
299";
300        let store = SnippetStore::parse(content);
301        assert_eq!(store.snippets.len(), 2);
302        assert_eq!(store.snippets[0].name, "check-disk");
303        assert_eq!(store.snippets[1].name, "uptime");
304    }
305
306    #[test]
307    fn test_parse_comments_and_blanks() {
308        let content = "\
309# Snippet config
310
311[check-disk]
312# Main command
313command=df -h
314";
315        let store = SnippetStore::parse(content);
316        assert_eq!(store.snippets.len(), 1);
317        assert_eq!(store.snippets[0].command, "df -h");
318    }
319
320    #[test]
321    fn test_parse_duplicate_sections_first_wins() {
322        let content = "\
323[check-disk]
324command=df -h
325
326[check-disk]
327command=du -sh *
328";
329        let store = SnippetStore::parse(content);
330        assert_eq!(store.snippets.len(), 1);
331        assert_eq!(store.snippets[0].command, "df -h");
332    }
333
334    #[test]
335    fn test_parse_snippet_without_command_skipped() {
336        let content = "\
337[empty]
338description=No command here
339
340[valid]
341command=ls -la
342";
343        let store = SnippetStore::parse(content);
344        assert_eq!(store.snippets.len(), 1);
345        assert_eq!(store.snippets[0].name, "valid");
346    }
347
348    #[test]
349    fn test_parse_unknown_keys_ignored() {
350        let content = "\
351[check-disk]
352command=df -h
353unknown=value
354foo=bar
355";
356        let store = SnippetStore::parse(content);
357        assert_eq!(store.snippets.len(), 1);
358        assert_eq!(store.snippets[0].command, "df -h");
359    }
360
361    #[test]
362    fn test_parse_whitespace_in_section_name() {
363        let content = "[ check-disk ]\ncommand=df -h\n";
364        let store = SnippetStore::parse(content);
365        assert_eq!(store.snippets[0].name, "check-disk");
366    }
367
368    #[test]
369    fn test_parse_whitespace_around_key_value() {
370        let content = "[check-disk]\n  command  =  df -h  \n";
371        let store = SnippetStore::parse(content);
372        assert_eq!(store.snippets[0].command, "df -h");
373    }
374
375    #[test]
376    fn test_parse_command_with_equals() {
377        let content = "[env-check]\ncommand=env | grep HOME=\n";
378        let store = SnippetStore::parse(content);
379        assert_eq!(store.snippets[0].command, "env | grep HOME=");
380    }
381
382    #[test]
383    fn test_parse_line_without_equals_ignored() {
384        let content = "[check]\ncommand=ls\ngarbage_line\n";
385        let store = SnippetStore::parse(content);
386        assert_eq!(store.snippets[0].command, "ls");
387    }
388
389    // =========================================================================
390    // Get / Set / Remove
391    // =========================================================================
392
393    #[test]
394    fn test_get_found() {
395        let store = SnippetStore::parse("[check]\ncommand=ls\n");
396        assert!(store.get("check").is_some());
397    }
398
399    #[test]
400    fn test_get_not_found() {
401        let store = SnippetStore::parse("");
402        assert!(store.get("nope").is_none());
403    }
404
405    #[test]
406    fn test_set_adds_new() {
407        let mut store = SnippetStore::default();
408        store.set(Snippet {
409            name: "check".to_string(),
410            command: "ls".to_string(),
411            description: String::new(),
412        });
413        assert_eq!(store.snippets.len(), 1);
414    }
415
416    #[test]
417    fn test_set_replaces_existing() {
418        let mut store = SnippetStore::parse("[check]\ncommand=ls\n");
419        store.set(Snippet {
420            name: "check".to_string(),
421            command: "df -h".to_string(),
422            description: String::new(),
423        });
424        assert_eq!(store.snippets.len(), 1);
425        assert_eq!(store.snippets[0].command, "df -h");
426    }
427
428    #[test]
429    fn test_remove() {
430        let mut store = SnippetStore::parse("[check]\ncommand=ls\n[uptime]\ncommand=uptime\n");
431        store.remove("check");
432        assert_eq!(store.snippets.len(), 1);
433        assert_eq!(store.snippets[0].name, "uptime");
434    }
435
436    #[test]
437    fn test_remove_nonexistent_noop() {
438        let mut store = SnippetStore::parse("[check]\ncommand=ls\n");
439        store.remove("nope");
440        assert_eq!(store.snippets.len(), 1);
441    }
442
443    // =========================================================================
444    // Validate name
445    // =========================================================================
446
447    #[test]
448    fn test_validate_name_valid() {
449        assert!(validate_name("check-disk").is_ok());
450        assert!(validate_name("restart_nginx").is_ok());
451        assert!(validate_name("a").is_ok());
452    }
453
454    #[test]
455    fn test_validate_name_empty() {
456        assert!(validate_name("").is_err());
457    }
458
459    #[test]
460    fn test_validate_name_whitespace() {
461        assert!(validate_name("check disk").is_err());
462        assert!(validate_name("check\tdisk").is_err());
463    }
464
465    #[test]
466    fn test_validate_name_special_chars() {
467        assert!(validate_name("check#disk").is_err());
468        assert!(validate_name("[check]").is_err());
469    }
470
471    #[test]
472    fn test_validate_name_control_chars() {
473        assert!(validate_name("check\x00disk").is_err());
474    }
475
476    // =========================================================================
477    // Validate command
478    // =========================================================================
479
480    #[test]
481    fn test_validate_command_valid() {
482        assert!(validate_command("df -h").is_ok());
483        assert!(validate_command("cat /etc/hosts | grep localhost").is_ok());
484        assert!(validate_command("echo 'hello\tworld'").is_ok()); // tab allowed
485    }
486
487    #[test]
488    fn test_validate_command_empty() {
489        assert!(validate_command("").is_err());
490    }
491
492    #[test]
493    fn test_validate_command_whitespace_only() {
494        assert!(validate_command("   ").is_err());
495        assert!(validate_command(" \t ").is_err());
496    }
497
498    #[test]
499    fn test_validate_command_control_chars() {
500        assert!(validate_command("ls\x00-la").is_err());
501    }
502
503    // =========================================================================
504    // Save / roundtrip
505    // =========================================================================
506
507    #[test]
508    fn test_save_roundtrip() {
509        let mut store = SnippetStore::default();
510        store.set(Snippet {
511            name: "check-disk".to_string(),
512            command: "df -h".to_string(),
513            description: "Check disk usage".to_string(),
514        });
515        store.set(Snippet {
516            name: "uptime".to_string(),
517            command: "uptime".to_string(),
518            description: String::new(),
519        });
520
521        // Serialize
522        let mut content = String::new();
523        for (i, snippet) in store.snippets.iter().enumerate() {
524            if i > 0 {
525                content.push('\n');
526            }
527            content.push_str(&format!("[{}]\n", snippet.name));
528            content.push_str(&format!("command={}\n", snippet.command));
529            if !snippet.description.is_empty() {
530                content.push_str(&format!("description={}\n", snippet.description));
531            }
532        }
533
534        // Re-parse
535        let reparsed = SnippetStore::parse(&content);
536        assert_eq!(reparsed.snippets.len(), 2);
537        assert_eq!(reparsed.snippets[0].name, "check-disk");
538        assert_eq!(reparsed.snippets[0].command, "df -h");
539        assert_eq!(reparsed.snippets[0].description, "Check disk usage");
540        assert_eq!(reparsed.snippets[1].name, "uptime");
541        assert_eq!(reparsed.snippets[1].command, "uptime");
542        assert!(reparsed.snippets[1].description.is_empty());
543    }
544
545    #[test]
546    fn test_save_to_temp_file() {
547        let dir = std::env::temp_dir().join(format!("purple_snippet_test_{}", std::process::id()));
548        let _ = std::fs::create_dir_all(&dir);
549        let path = dir.join("snippets");
550
551        let mut store = SnippetStore {
552            path_override: Some(path.clone()),
553            ..Default::default()
554        };
555        store.set(Snippet {
556            name: "test".to_string(),
557            command: "echo hello".to_string(),
558            description: "Test snippet".to_string(),
559        });
560        store.save().unwrap();
561
562        // Read back
563        let content = std::fs::read_to_string(&path).unwrap();
564        let reloaded = SnippetStore::parse(&content);
565        assert_eq!(reloaded.snippets.len(), 1);
566        assert_eq!(reloaded.snippets[0].name, "test");
567        assert_eq!(reloaded.snippets[0].command, "echo hello");
568
569        // Cleanup
570        let _ = std::fs::remove_dir_all(&dir);
571    }
572
573    // =========================================================================
574    // Edge cases
575    // =========================================================================
576
577    #[test]
578    fn test_set_multiple_then_remove_all() {
579        let mut store = SnippetStore::default();
580        for name in ["a", "b", "c"] {
581            store.set(Snippet {
582                name: name.to_string(),
583                command: "cmd".to_string(),
584                description: String::new(),
585            });
586        }
587        assert_eq!(store.snippets.len(), 3);
588        store.remove("a");
589        store.remove("b");
590        store.remove("c");
591        assert!(store.snippets.is_empty());
592    }
593
594    #[test]
595    fn test_snippet_with_complex_command() {
596        let content = "[complex]\ncommand=for i in $(seq 1 5); do echo $i; done\n";
597        let store = SnippetStore::parse(content);
598        assert_eq!(
599            store.snippets[0].command,
600            "for i in $(seq 1 5); do echo $i; done"
601        );
602    }
603
604    #[test]
605    fn test_snippet_command_with_pipes_and_redirects() {
606        let content = "[logs]\ncommand=tail -100 /var/log/syslog | grep error | head -20\n";
607        let store = SnippetStore::parse(content);
608        assert_eq!(
609            store.snippets[0].command,
610            "tail -100 /var/log/syslog | grep error | head -20"
611        );
612    }
613
614    #[test]
615    fn test_description_optional() {
616        let content = "[check]\ncommand=ls\n";
617        let store = SnippetStore::parse(content);
618        assert!(store.snippets[0].description.is_empty());
619    }
620
621    #[test]
622    fn test_description_with_equals() {
623        let content = "[env]\ncommand=env\ndescription=Check HOME= and PATH= vars\n";
624        let store = SnippetStore::parse(content);
625        assert_eq!(store.snippets[0].description, "Check HOME= and PATH= vars");
626    }
627
628    #[test]
629    fn test_name_with_equals_roundtrip() {
630        let mut store = SnippetStore::default();
631        store.set(Snippet {
632            name: "check=disk".to_string(),
633            command: "df -h".to_string(),
634            description: String::new(),
635        });
636
637        let mut content = String::new();
638        for (i, snippet) in store.snippets.iter().enumerate() {
639            if i > 0 {
640                content.push('\n');
641            }
642            content.push_str(&format!("[{}]\n", snippet.name));
643            content.push_str(&format!("command={}\n", snippet.command));
644            if !snippet.description.is_empty() {
645                content.push_str(&format!("description={}\n", snippet.description));
646            }
647        }
648
649        let reparsed = SnippetStore::parse(&content);
650        assert_eq!(reparsed.snippets.len(), 1);
651        assert_eq!(reparsed.snippets[0].name, "check=disk");
652    }
653
654    #[test]
655    fn test_validate_name_with_equals() {
656        assert!(validate_name("check=disk").is_ok());
657    }
658
659    #[test]
660    fn test_parse_only_comments_and_blanks() {
661        let content = "# comment\n\n# another\n";
662        let store = SnippetStore::parse(content);
663        assert!(store.snippets.is_empty());
664    }
665
666    #[test]
667    fn test_parse_section_without_close_bracket() {
668        let content = "[incomplete\ncommand=ls\n";
669        let store = SnippetStore::parse(content);
670        assert!(store.snippets.is_empty());
671    }
672
673    #[test]
674    fn test_parse_trailing_content_after_last_section() {
675        let content = "[check]\ncommand=ls\n";
676        let store = SnippetStore::parse(content);
677        assert_eq!(store.snippets.len(), 1);
678        assert_eq!(store.snippets[0].command, "ls");
679    }
680
681    #[test]
682    fn test_set_overwrite_preserves_order() {
683        let mut store = SnippetStore::default();
684        store.set(Snippet { name: "a".into(), command: "1".into(), description: String::new() });
685        store.set(Snippet { name: "b".into(), command: "2".into(), description: String::new() });
686        store.set(Snippet { name: "c".into(), command: "3".into(), description: String::new() });
687        store.set(Snippet { name: "b".into(), command: "updated".into(), description: String::new() });
688        assert_eq!(store.snippets.len(), 3);
689        assert_eq!(store.snippets[0].name, "a");
690        assert_eq!(store.snippets[1].name, "b");
691        assert_eq!(store.snippets[1].command, "updated");
692        assert_eq!(store.snippets[2].name, "c");
693    }
694
695    #[test]
696    fn test_validate_command_with_tab() {
697        assert!(validate_command("echo\thello").is_ok());
698    }
699
700    #[test]
701    fn test_validate_command_with_newline() {
702        assert!(validate_command("echo\nhello").is_err());
703    }
704
705    #[test]
706    fn test_validate_name_newline() {
707        assert!(validate_name("check\ndisk").is_err());
708    }
709
710}