1use std::io;
2use std::path::{Path, PathBuf};
3use std::process::{Command, ExitStatus, Stdio};
4
5use crate::fs_util;
6
7#[derive(Debug, Clone, PartialEq)]
9pub struct Snippet {
10 pub name: String,
11 pub command: String,
12 pub description: String,
13}
14
15pub struct SnippetResult {
17 pub status: ExitStatus,
18 pub stdout: String,
19 pub stderr: String,
20}
21
22#[derive(Debug, Clone, Default)]
24pub struct SnippetStore {
25 pub snippets: Vec<Snippet>,
26 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 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 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_start().to_string();
87 match key {
88 "command" => snippet.command = value,
89 "description" => snippet.description = value,
90 _ => {}
91 }
92 }
93 }
94 }
95 if let Some(snippet) = current {
96 if !snippet.command.is_empty()
97 && !snippets.iter().any(|s| s.name == snippet.name)
98 {
99 snippets.push(snippet);
100 }
101 }
102 Self {
103 snippets,
104 path_override: None,
105 }
106 }
107
108 pub fn save(&self) -> io::Result<()> {
110 let path = match &self.path_override {
111 Some(p) => p.clone(),
112 None => match config_path() {
113 Some(p) => p,
114 None => {
115 return Err(io::Error::new(
116 io::ErrorKind::NotFound,
117 "Could not determine home directory",
118 ))
119 }
120 },
121 };
122
123 let mut content = String::new();
124 for (i, snippet) in self.snippets.iter().enumerate() {
125 if i > 0 {
126 content.push('\n');
127 }
128 content.push_str(&format!("[{}]\n", snippet.name));
129 content.push_str(&format!("command={}\n", snippet.command));
130 if !snippet.description.is_empty() {
131 content.push_str(&format!("description={}\n", snippet.description));
132 }
133 }
134
135 fs_util::atomic_write(&path, content.as_bytes())
136 }
137
138 pub fn get(&self, name: &str) -> Option<&Snippet> {
140 self.snippets.iter().find(|s| s.name == name)
141 }
142
143 pub fn set(&mut self, snippet: Snippet) {
145 if let Some(existing) = self.snippets.iter_mut().find(|s| s.name == snippet.name) {
146 *existing = snippet;
147 } else {
148 self.snippets.push(snippet);
149 }
150 }
151
152 pub fn remove(&mut self, name: &str) {
154 self.snippets.retain(|s| s.name != name);
155 }
156}
157
158pub fn validate_name(name: &str) -> Result<(), String> {
161 if name.is_empty() {
162 return Err("Snippet name cannot be empty.".to_string());
163 }
164 if name.contains(char::is_whitespace) {
165 return Err("Snippet name cannot contain whitespace.".to_string());
166 }
167 if name.contains('#') || name.contains('[') || name.contains(']') {
168 return Err("Snippet name cannot contain #, [ or ].".to_string());
169 }
170 if name.contains(|c: char| c.is_control()) {
171 return Err("Snippet name cannot contain control characters.".to_string());
172 }
173 Ok(())
174}
175
176pub fn validate_command(command: &str) -> Result<(), String> {
178 if command.trim().is_empty() {
179 return Err("Command cannot be empty.".to_string());
180 }
181 if command.contains(|c: char| c.is_control() && c != '\t') {
182 return Err("Command cannot contain control characters.".to_string());
183 }
184 Ok(())
185}
186
187#[derive(Debug, Clone, PartialEq)]
193pub struct SnippetParam {
194 pub name: String,
195 pub default: Option<String>,
196}
197
198pub fn shell_escape(s: &str) -> String {
201 format!("'{}'", s.replace('\'', "'\\''"))
202}
203
204pub fn parse_params(command: &str) -> Vec<SnippetParam> {
207 let mut params = Vec::new();
208 let mut seen = std::collections::HashSet::new();
209 let bytes = command.as_bytes();
210 let len = bytes.len();
211 let mut i = 0;
212 while i + 3 < len {
213 if bytes[i] == b'{' && bytes.get(i + 1) == Some(&b'{') {
214 if let Some(end) = command[i + 2..].find("}}") {
215 let inner = &command[i + 2..i + 2 + end];
216 let (name, default) = if let Some((n, d)) = inner.split_once(':') {
217 (n.to_string(), Some(d.to_string()))
218 } else {
219 (inner.to_string(), None)
220 };
221 if validate_param_name(&name).is_ok()
222 && !seen.contains(&name)
223 && params.len() < 20
224 {
225 seen.insert(name.clone());
226 params.push(SnippetParam { name, default });
227 }
228 i = i + 2 + end + 2;
229 continue;
230 }
231 }
232 i += 1;
233 }
234 params
235}
236
237pub fn validate_param_name(name: &str) -> Result<(), String> {
240 if name.is_empty() {
241 return Err("Parameter name cannot be empty.".to_string());
242 }
243 if !name
244 .chars()
245 .all(|c| c.is_alphanumeric() || c == '_' || c == '-')
246 {
247 return Err(format!(
248 "Parameter name '{}' contains invalid characters.",
249 name
250 ));
251 }
252 Ok(())
253}
254
255pub fn substitute_params(
258 command: &str,
259 values: &std::collections::HashMap<String, String>,
260) -> String {
261 let mut result = String::with_capacity(command.len());
262 let bytes = command.as_bytes();
263 let len = bytes.len();
264 let mut i = 0;
265 while i < len {
266 if i + 3 < len && bytes[i] == b'{' && bytes[i + 1] == b'{' {
267 if let Some(end) = command[i + 2..].find("}}") {
268 let inner = &command[i + 2..i + 2 + end];
269 let (name, default) = if let Some((n, d)) = inner.split_once(':') {
270 (n, Some(d))
271 } else {
272 (inner, None)
273 };
274 let value = values
275 .get(name)
276 .filter(|v| !v.is_empty())
277 .map(|v| v.as_str())
278 .or(default)
279 .unwrap_or("");
280 result.push_str(&shell_escape(value));
281 i = i + 2 + end + 2;
282 continue;
283 }
284 }
285 let ch = command[i..].chars().next().unwrap();
287 result.push(ch);
288 i += ch.len_utf8();
289 }
290 result
291}
292
293pub fn sanitize_output(input: &str) -> String {
300 let mut out = String::with_capacity(input.len());
301 let mut chars = input.chars().peekable();
302 while let Some(c) = chars.next() {
303 match c {
304 '\x1b' => {
305 match chars.peek() {
306 Some('[') => {
307 chars.next();
308 while let Some(&ch) = chars.peek() {
310 chars.next();
311 if ('\x40'..='\x7e').contains(&ch) {
312 break;
313 }
314 }
315 }
316 Some(']') | Some('P') | Some('X') | Some('^') | Some('_') => {
317 chars.next();
318 consume_until_st(&mut chars);
320 }
321 _ => {
322 chars.next();
324 }
325 }
326 }
327 c if ('\u{0080}'..='\u{009F}').contains(&c) => {
328 }
330 c if c.is_control() && c != '\n' && c != '\t' => {
331 }
333 _ => out.push(c),
334 }
335 }
336 out
337}
338
339fn consume_until_st(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
341 while let Some(&ch) = chars.peek() {
342 if ch == '\x07' {
343 chars.next();
344 break;
345 }
346 if ch == '\x1b' {
347 chars.next();
348 if chars.peek() == Some(&'\\') {
349 chars.next();
350 }
351 break;
352 }
353 chars.next();
354 }
355}
356
357const MAX_OUTPUT_LINES: usize = 10_000;
364
365pub enum SnippetEvent {
368 HostDone {
369 run_id: u64,
370 alias: String,
371 stdout: String,
372 stderr: String,
373 exit_code: Option<i32>,
374 },
375 Progress {
376 run_id: u64,
377 completed: usize,
378 total: usize,
379 },
380 AllDone {
381 run_id: u64,
382 },
383}
384
385pub struct ChildGuard {
388 inner: std::sync::Mutex<Option<std::process::Child>>,
389 pgid: i32,
390}
391
392impl ChildGuard {
393 fn new(child: std::process::Child) -> Self {
394 let pgid = child.id() as i32;
395 Self {
396 inner: std::sync::Mutex::new(Some(child)),
397 pgid,
398 }
399 }
400}
401
402impl Drop for ChildGuard {
403 fn drop(&mut self) {
404 let mut lock = self.inner.lock().unwrap_or_else(|e| e.into_inner());
405 if let Some(ref mut child) = *lock {
406 if let Ok(Some(_)) = child.try_wait() {
408 return;
409 }
410 #[cfg(unix)]
412 unsafe {
413 libc::kill(-self.pgid, libc::SIGTERM);
414 }
415 let deadline =
417 std::time::Instant::now() + std::time::Duration::from_millis(500);
418 loop {
419 if let Ok(Some(_)) = child.try_wait() {
420 return;
421 }
422 if std::time::Instant::now() >= deadline {
423 break;
424 }
425 std::thread::sleep(std::time::Duration::from_millis(50));
426 }
427 #[cfg(unix)]
429 unsafe {
430 libc::kill(-self.pgid, libc::SIGKILL);
431 }
432 let _ = child.kill();
434 let _ = child.wait();
435 }
436 }
437}
438
439fn read_pipe_capped<R: io::Read>(reader: R) -> String {
442 use io::BufRead;
443 let mut reader = io::BufReader::new(reader);
444 let mut output = String::new();
445 let mut line_count = 0;
446 let mut capped = false;
447 let mut buf = Vec::new();
448 loop {
449 buf.clear();
450 match reader.read_until(b'\n', &mut buf) {
451 Ok(0) => break, Ok(_) => {
453 if !capped {
454 if line_count < MAX_OUTPUT_LINES {
455 if line_count > 0 {
456 output.push('\n');
457 }
458 if buf.last() == Some(&b'\n') {
460 buf.pop();
461 if buf.last() == Some(&b'\r') {
462 buf.pop();
463 }
464 }
465 output.push_str(&String::from_utf8_lossy(&buf));
467 line_count += 1;
468 } else {
469 output.push_str("\n[Output truncated at 10,000 lines]");
470 capped = true;
471 }
472 }
473 }
475 Err(_) => break,
476 }
477 }
478 output
479}
480
481fn base_ssh_command(
485 alias: &str,
486 config_path: &Path,
487 command: &str,
488 askpass: Option<&str>,
489 bw_session: Option<&str>,
490 has_active_tunnel: bool,
491) -> Command {
492 let mut cmd = Command::new("ssh");
493 cmd.arg("-F")
494 .arg(config_path)
495 .arg("-o")
496 .arg("ConnectTimeout=10")
497 .arg("-o")
498 .arg("ControlMaster=no")
499 .arg("-o")
500 .arg("ControlPath=none");
501
502 if has_active_tunnel {
503 cmd.arg("-o").arg("ClearAllForwardings=yes");
504 }
505
506 cmd.arg("--").arg(alias).arg(command);
507
508 if askpass.is_some() {
509 let exe = std::env::current_exe()
510 .ok()
511 .map(|p| p.to_string_lossy().to_string())
512 .or_else(|| std::env::args().next())
513 .unwrap_or_else(|| "purple".to_string());
514 cmd.env("SSH_ASKPASS", &exe)
515 .env("SSH_ASKPASS_REQUIRE", "prefer")
516 .env("PURPLE_ASKPASS_MODE", "1")
517 .env("PURPLE_HOST_ALIAS", alias)
518 .env("PURPLE_CONFIG_PATH", config_path.as_os_str());
519 }
520
521 if let Some(token) = bw_session {
522 cmd.env("BW_SESSION", token);
523 }
524
525 cmd
526}
527
528fn build_snippet_command(
530 alias: &str,
531 config_path: &Path,
532 command: &str,
533 askpass: Option<&str>,
534 bw_session: Option<&str>,
535 has_active_tunnel: bool,
536) -> Command {
537 let mut cmd = base_ssh_command(alias, config_path, command, askpass, bw_session, has_active_tunnel);
538 cmd.stdin(Stdio::null())
539 .stdout(Stdio::piped())
540 .stderr(Stdio::piped());
541
542 #[cfg(unix)]
545 unsafe {
546 use std::os::unix::process::CommandExt;
547 cmd.pre_exec(|| {
548 libc::setpgid(0, 0);
549 Ok(())
550 });
551 }
552
553 cmd
554}
555
556#[allow(clippy::too_many_arguments)]
558fn execute_host(
559 run_id: u64,
560 alias: &str,
561 config_path: &Path,
562 command: &str,
563 askpass: Option<&str>,
564 bw_session: Option<&str>,
565 has_active_tunnel: bool,
566 tx: &std::sync::mpsc::Sender<SnippetEvent>,
567) -> Option<std::sync::Arc<ChildGuard>> {
568 let mut cmd =
569 build_snippet_command(alias, config_path, command, askpass, bw_session, has_active_tunnel);
570
571 match cmd.spawn() {
572 Ok(child) => {
573 let guard = std::sync::Arc::new(ChildGuard::new(child));
574
575 let stdout_pipe = {
577 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
578 lock.as_mut().and_then(|c| c.stdout.take())
579 };
580 let stderr_pipe = {
581 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
582 lock.as_mut().and_then(|c| c.stderr.take())
583 };
584
585 let stdout_handle = std::thread::spawn(move || match stdout_pipe {
587 Some(pipe) => read_pipe_capped(pipe),
588 None => String::new(),
589 });
590 let stderr_handle = std::thread::spawn(move || match stderr_pipe {
591 Some(pipe) => read_pipe_capped(pipe),
592 None => String::new(),
593 });
594
595 let stdout_text = stdout_handle.join().unwrap_or_default();
597 let stderr_text = stderr_handle.join().unwrap_or_default();
598
599 let exit_code = {
602 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
603 let status = lock.as_mut().and_then(|c| c.wait().ok());
604 let _ = lock.take(); status.and_then(|s| {
606 #[cfg(unix)]
607 {
608 use std::os::unix::process::ExitStatusExt;
609 s.code().or_else(|| s.signal().map(|sig| 128 + sig))
610 }
611 #[cfg(not(unix))]
612 {
613 s.code()
614 }
615 })
616 };
617
618 let _ = tx.send(SnippetEvent::HostDone {
619 run_id,
620 alias: alias.to_string(),
621 stdout: sanitize_output(&stdout_text),
622 stderr: sanitize_output(&stderr_text),
623 exit_code,
624 });
625
626 Some(guard)
627 }
628 Err(e) => {
629 let _ = tx.send(SnippetEvent::HostDone {
630 run_id,
631 alias: alias.to_string(),
632 stdout: String::new(),
633 stderr: format!("Failed to launch ssh: {}", e),
634 exit_code: None,
635 });
636 None
637 }
638 }
639}
640
641#[allow(clippy::too_many_arguments)]
644pub fn spawn_snippet_execution(
645 run_id: u64,
646 askpass_map: Vec<(String, Option<String>)>,
647 config_path: PathBuf,
648 command: String,
649 bw_session: Option<String>,
650 tunnel_aliases: std::collections::HashSet<String>,
651 cancel: std::sync::Arc<std::sync::atomic::AtomicBool>,
652 tx: std::sync::mpsc::Sender<SnippetEvent>,
653 parallel: bool,
654) {
655 let total = askpass_map.len();
656 let max_concurrent: usize = 20;
657
658 std::thread::Builder::new()
659 .name("snippet-coordinator".into())
660 .spawn(move || {
661 let guards: std::sync::Arc<
662 std::sync::Mutex<Vec<std::sync::Arc<ChildGuard>>>,
663 > = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
664
665 if parallel && total > 1 {
666 let (slot_tx, slot_rx) = std::sync::mpsc::channel::<()>();
668 for _ in 0..max_concurrent.min(total) {
669 let _ = slot_tx.send(());
670 }
671
672 let completed =
673 std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
674 let mut worker_handles = Vec::new();
675
676 for (alias, askpass) in askpass_map {
677 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
678 break;
679 }
680
681 loop {
683 match slot_rx.recv_timeout(std::time::Duration::from_millis(100))
684 {
685 Ok(()) => break,
686 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
687 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
688 break;
689 }
690 }
691 Err(_) => break, }
693 }
694
695 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
696 break;
697 }
698
699 let config_path = config_path.clone();
700 let command = command.clone();
701 let bw_session = bw_session.clone();
702 let has_tunnel = tunnel_aliases.contains(&alias);
703 let tx = tx.clone();
704 let slot_tx = slot_tx.clone();
705 let guards = guards.clone();
706 let completed = completed.clone();
707 let total = total;
708
709 let handle = std::thread::spawn(move || {
710 struct SlotRelease(Option<std::sync::mpsc::Sender<()>>);
712 impl Drop for SlotRelease {
713 fn drop(&mut self) {
714 if let Some(tx) = self.0.take() {
715 let _ = tx.send(());
716 }
717 }
718 }
719 let _slot = SlotRelease(Some(slot_tx));
720
721 let guard = execute_host(
722 run_id,
723 &alias,
724 &config_path,
725 &command,
726 askpass.as_deref(),
727 bw_session.as_deref(),
728 has_tunnel,
729 &tx,
730 );
731
732 if let Some(g) = guard {
734 guards
735 .lock()
736 .unwrap_or_else(|e| e.into_inner())
737 .push(g);
738 }
739
740 let c = completed.fetch_add(
741 1,
742 std::sync::atomic::Ordering::Relaxed,
743 ) + 1;
744 let _ = tx.send(SnippetEvent::Progress {
745 run_id,
746 completed: c,
747 total,
748 });
749 });
751 worker_handles.push(handle);
752 }
753
754 for handle in worker_handles {
756 let _ = handle.join();
757 }
758 } else {
759 for (i, (alias, askpass)) in askpass_map.into_iter().enumerate() {
761 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
762 break;
763 }
764
765 let has_tunnel = tunnel_aliases.contains(&alias);
766 let guard = execute_host(
767 run_id,
768 &alias,
769 &config_path,
770 &command,
771 askpass.as_deref(),
772 bw_session.as_deref(),
773 has_tunnel,
774 &tx,
775 );
776
777 if let Some(g) = guard {
778 guards
779 .lock()
780 .unwrap_or_else(|e| e.into_inner())
781 .push(g);
782 }
783
784 let _ = tx.send(SnippetEvent::Progress {
785 run_id,
786 completed: i + 1,
787 total,
788 });
789 }
790 }
791
792 let _ = tx.send(SnippetEvent::AllDone { run_id });
793 })
795 .expect("failed to spawn snippet coordinator");
796}
797
798pub fn run_snippet(
803 alias: &str,
804 config_path: &Path,
805 command: &str,
806 askpass: Option<&str>,
807 bw_session: Option<&str>,
808 capture: bool,
809 has_active_tunnel: bool,
810) -> anyhow::Result<SnippetResult> {
811 let mut cmd = base_ssh_command(alias, config_path, command, askpass, bw_session, has_active_tunnel);
812 cmd.stdin(Stdio::inherit());
813
814 if capture {
815 cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
816 } else {
817 cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
818 }
819
820 if capture {
821 let output = cmd
822 .output()
823 .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
824
825 Ok(SnippetResult {
826 status: output.status,
827 stdout: String::from_utf8_lossy(&output.stdout).to_string(),
828 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
829 })
830 } else {
831 let status = cmd
832 .status()
833 .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
834
835 Ok(SnippetResult {
836 status,
837 stdout: String::new(),
838 stderr: String::new(),
839 })
840 }
841}
842
843#[cfg(test)]
844mod tests {
845 use super::*;
846
847 #[test]
852 fn test_parse_empty() {
853 let store = SnippetStore::parse("");
854 assert!(store.snippets.is_empty());
855 }
856
857 #[test]
858 fn test_parse_single_snippet() {
859 let content = "\
860[check-disk]
861command=df -h
862description=Check disk usage
863";
864 let store = SnippetStore::parse(content);
865 assert_eq!(store.snippets.len(), 1);
866 let s = &store.snippets[0];
867 assert_eq!(s.name, "check-disk");
868 assert_eq!(s.command, "df -h");
869 assert_eq!(s.description, "Check disk usage");
870 }
871
872 #[test]
873 fn test_parse_multiple_snippets() {
874 let content = "\
875[check-disk]
876command=df -h
877
878[uptime]
879command=uptime
880description=Check server uptime
881";
882 let store = SnippetStore::parse(content);
883 assert_eq!(store.snippets.len(), 2);
884 assert_eq!(store.snippets[0].name, "check-disk");
885 assert_eq!(store.snippets[1].name, "uptime");
886 }
887
888 #[test]
889 fn test_parse_comments_and_blanks() {
890 let content = "\
891# Snippet config
892
893[check-disk]
894# Main command
895command=df -h
896";
897 let store = SnippetStore::parse(content);
898 assert_eq!(store.snippets.len(), 1);
899 assert_eq!(store.snippets[0].command, "df -h");
900 }
901
902 #[test]
903 fn test_parse_duplicate_sections_first_wins() {
904 let content = "\
905[check-disk]
906command=df -h
907
908[check-disk]
909command=du -sh *
910";
911 let store = SnippetStore::parse(content);
912 assert_eq!(store.snippets.len(), 1);
913 assert_eq!(store.snippets[0].command, "df -h");
914 }
915
916 #[test]
917 fn test_parse_snippet_without_command_skipped() {
918 let content = "\
919[empty]
920description=No command here
921
922[valid]
923command=ls -la
924";
925 let store = SnippetStore::parse(content);
926 assert_eq!(store.snippets.len(), 1);
927 assert_eq!(store.snippets[0].name, "valid");
928 }
929
930 #[test]
931 fn test_parse_unknown_keys_ignored() {
932 let content = "\
933[check-disk]
934command=df -h
935unknown=value
936foo=bar
937";
938 let store = SnippetStore::parse(content);
939 assert_eq!(store.snippets.len(), 1);
940 assert_eq!(store.snippets[0].command, "df -h");
941 }
942
943 #[test]
944 fn test_parse_whitespace_in_section_name() {
945 let content = "[ check-disk ]\ncommand=df -h\n";
946 let store = SnippetStore::parse(content);
947 assert_eq!(store.snippets[0].name, "check-disk");
948 }
949
950 #[test]
951 fn test_parse_whitespace_around_key_value() {
952 let content = "[check-disk]\n command = df -h \n";
953 let store = SnippetStore::parse(content);
954 assert_eq!(store.snippets[0].command, "df -h");
955 }
956
957 #[test]
958 fn test_parse_command_with_equals() {
959 let content = "[env-check]\ncommand=env | grep HOME=\n";
960 let store = SnippetStore::parse(content);
961 assert_eq!(store.snippets[0].command, "env | grep HOME=");
962 }
963
964 #[test]
965 fn test_parse_line_without_equals_ignored() {
966 let content = "[check]\ncommand=ls\ngarbage_line\n";
967 let store = SnippetStore::parse(content);
968 assert_eq!(store.snippets[0].command, "ls");
969 }
970
971 #[test]
976 fn test_get_found() {
977 let store = SnippetStore::parse("[check]\ncommand=ls\n");
978 assert!(store.get("check").is_some());
979 }
980
981 #[test]
982 fn test_get_not_found() {
983 let store = SnippetStore::parse("");
984 assert!(store.get("nope").is_none());
985 }
986
987 #[test]
988 fn test_set_adds_new() {
989 let mut store = SnippetStore::default();
990 store.set(Snippet {
991 name: "check".to_string(),
992 command: "ls".to_string(),
993 description: String::new(),
994 });
995 assert_eq!(store.snippets.len(), 1);
996 }
997
998 #[test]
999 fn test_set_replaces_existing() {
1000 let mut store = SnippetStore::parse("[check]\ncommand=ls\n");
1001 store.set(Snippet {
1002 name: "check".to_string(),
1003 command: "df -h".to_string(),
1004 description: String::new(),
1005 });
1006 assert_eq!(store.snippets.len(), 1);
1007 assert_eq!(store.snippets[0].command, "df -h");
1008 }
1009
1010 #[test]
1011 fn test_remove() {
1012 let mut store = SnippetStore::parse("[check]\ncommand=ls\n[uptime]\ncommand=uptime\n");
1013 store.remove("check");
1014 assert_eq!(store.snippets.len(), 1);
1015 assert_eq!(store.snippets[0].name, "uptime");
1016 }
1017
1018 #[test]
1019 fn test_remove_nonexistent_noop() {
1020 let mut store = SnippetStore::parse("[check]\ncommand=ls\n");
1021 store.remove("nope");
1022 assert_eq!(store.snippets.len(), 1);
1023 }
1024
1025 #[test]
1030 fn test_validate_name_valid() {
1031 assert!(validate_name("check-disk").is_ok());
1032 assert!(validate_name("restart_nginx").is_ok());
1033 assert!(validate_name("a").is_ok());
1034 }
1035
1036 #[test]
1037 fn test_validate_name_empty() {
1038 assert!(validate_name("").is_err());
1039 }
1040
1041 #[test]
1042 fn test_validate_name_whitespace() {
1043 assert!(validate_name("check disk").is_err());
1044 assert!(validate_name("check\tdisk").is_err());
1045 }
1046
1047 #[test]
1048 fn test_validate_name_special_chars() {
1049 assert!(validate_name("check#disk").is_err());
1050 assert!(validate_name("[check]").is_err());
1051 }
1052
1053 #[test]
1054 fn test_validate_name_control_chars() {
1055 assert!(validate_name("check\x00disk").is_err());
1056 }
1057
1058 #[test]
1063 fn test_validate_command_valid() {
1064 assert!(validate_command("df -h").is_ok());
1065 assert!(validate_command("cat /etc/hosts | grep localhost").is_ok());
1066 assert!(validate_command("echo 'hello\tworld'").is_ok()); }
1068
1069 #[test]
1070 fn test_validate_command_empty() {
1071 assert!(validate_command("").is_err());
1072 }
1073
1074 #[test]
1075 fn test_validate_command_whitespace_only() {
1076 assert!(validate_command(" ").is_err());
1077 assert!(validate_command(" \t ").is_err());
1078 }
1079
1080 #[test]
1081 fn test_validate_command_control_chars() {
1082 assert!(validate_command("ls\x00-la").is_err());
1083 }
1084
1085 #[test]
1090 fn test_save_roundtrip() {
1091 let mut store = SnippetStore::default();
1092 store.set(Snippet {
1093 name: "check-disk".to_string(),
1094 command: "df -h".to_string(),
1095 description: "Check disk usage".to_string(),
1096 });
1097 store.set(Snippet {
1098 name: "uptime".to_string(),
1099 command: "uptime".to_string(),
1100 description: String::new(),
1101 });
1102
1103 let mut content = String::new();
1105 for (i, snippet) in store.snippets.iter().enumerate() {
1106 if i > 0 {
1107 content.push('\n');
1108 }
1109 content.push_str(&format!("[{}]\n", snippet.name));
1110 content.push_str(&format!("command={}\n", snippet.command));
1111 if !snippet.description.is_empty() {
1112 content.push_str(&format!("description={}\n", snippet.description));
1113 }
1114 }
1115
1116 let reparsed = SnippetStore::parse(&content);
1118 assert_eq!(reparsed.snippets.len(), 2);
1119 assert_eq!(reparsed.snippets[0].name, "check-disk");
1120 assert_eq!(reparsed.snippets[0].command, "df -h");
1121 assert_eq!(reparsed.snippets[0].description, "Check disk usage");
1122 assert_eq!(reparsed.snippets[1].name, "uptime");
1123 assert_eq!(reparsed.snippets[1].command, "uptime");
1124 assert!(reparsed.snippets[1].description.is_empty());
1125 }
1126
1127 #[test]
1128 fn test_save_to_temp_file() {
1129 let dir = std::env::temp_dir().join(format!("purple_snippet_test_{}", std::process::id()));
1130 let _ = std::fs::create_dir_all(&dir);
1131 let path = dir.join("snippets");
1132
1133 let mut store = SnippetStore {
1134 path_override: Some(path.clone()),
1135 ..Default::default()
1136 };
1137 store.set(Snippet {
1138 name: "test".to_string(),
1139 command: "echo hello".to_string(),
1140 description: "Test snippet".to_string(),
1141 });
1142 store.save().unwrap();
1143
1144 let content = std::fs::read_to_string(&path).unwrap();
1146 let reloaded = SnippetStore::parse(&content);
1147 assert_eq!(reloaded.snippets.len(), 1);
1148 assert_eq!(reloaded.snippets[0].name, "test");
1149 assert_eq!(reloaded.snippets[0].command, "echo hello");
1150
1151 let _ = std::fs::remove_dir_all(&dir);
1153 }
1154
1155 #[test]
1160 fn test_set_multiple_then_remove_all() {
1161 let mut store = SnippetStore::default();
1162 for name in ["a", "b", "c"] {
1163 store.set(Snippet {
1164 name: name.to_string(),
1165 command: "cmd".to_string(),
1166 description: String::new(),
1167 });
1168 }
1169 assert_eq!(store.snippets.len(), 3);
1170 store.remove("a");
1171 store.remove("b");
1172 store.remove("c");
1173 assert!(store.snippets.is_empty());
1174 }
1175
1176 #[test]
1177 fn test_snippet_with_complex_command() {
1178 let content = "[complex]\ncommand=for i in $(seq 1 5); do echo $i; done\n";
1179 let store = SnippetStore::parse(content);
1180 assert_eq!(
1181 store.snippets[0].command,
1182 "for i in $(seq 1 5); do echo $i; done"
1183 );
1184 }
1185
1186 #[test]
1187 fn test_snippet_command_with_pipes_and_redirects() {
1188 let content = "[logs]\ncommand=tail -100 /var/log/syslog | grep error | head -20\n";
1189 let store = SnippetStore::parse(content);
1190 assert_eq!(
1191 store.snippets[0].command,
1192 "tail -100 /var/log/syslog | grep error | head -20"
1193 );
1194 }
1195
1196 #[test]
1197 fn test_description_optional() {
1198 let content = "[check]\ncommand=ls\n";
1199 let store = SnippetStore::parse(content);
1200 assert!(store.snippets[0].description.is_empty());
1201 }
1202
1203 #[test]
1204 fn test_description_with_equals() {
1205 let content = "[env]\ncommand=env\ndescription=Check HOME= and PATH= vars\n";
1206 let store = SnippetStore::parse(content);
1207 assert_eq!(store.snippets[0].description, "Check HOME= and PATH= vars");
1208 }
1209
1210 #[test]
1211 fn test_name_with_equals_roundtrip() {
1212 let mut store = SnippetStore::default();
1213 store.set(Snippet {
1214 name: "check=disk".to_string(),
1215 command: "df -h".to_string(),
1216 description: String::new(),
1217 });
1218
1219 let mut content = String::new();
1220 for (i, snippet) in store.snippets.iter().enumerate() {
1221 if i > 0 {
1222 content.push('\n');
1223 }
1224 content.push_str(&format!("[{}]\n", snippet.name));
1225 content.push_str(&format!("command={}\n", snippet.command));
1226 if !snippet.description.is_empty() {
1227 content.push_str(&format!("description={}\n", snippet.description));
1228 }
1229 }
1230
1231 let reparsed = SnippetStore::parse(&content);
1232 assert_eq!(reparsed.snippets.len(), 1);
1233 assert_eq!(reparsed.snippets[0].name, "check=disk");
1234 }
1235
1236 #[test]
1237 fn test_validate_name_with_equals() {
1238 assert!(validate_name("check=disk").is_ok());
1239 }
1240
1241 #[test]
1242 fn test_parse_only_comments_and_blanks() {
1243 let content = "# comment\n\n# another\n";
1244 let store = SnippetStore::parse(content);
1245 assert!(store.snippets.is_empty());
1246 }
1247
1248 #[test]
1249 fn test_parse_section_without_close_bracket() {
1250 let content = "[incomplete\ncommand=ls\n";
1251 let store = SnippetStore::parse(content);
1252 assert!(store.snippets.is_empty());
1253 }
1254
1255 #[test]
1256 fn test_parse_trailing_content_after_last_section() {
1257 let content = "[check]\ncommand=ls\n";
1258 let store = SnippetStore::parse(content);
1259 assert_eq!(store.snippets.len(), 1);
1260 assert_eq!(store.snippets[0].command, "ls");
1261 }
1262
1263 #[test]
1264 fn test_set_overwrite_preserves_order() {
1265 let mut store = SnippetStore::default();
1266 store.set(Snippet { name: "a".into(), command: "1".into(), description: String::new() });
1267 store.set(Snippet { name: "b".into(), command: "2".into(), description: String::new() });
1268 store.set(Snippet { name: "c".into(), command: "3".into(), description: String::new() });
1269 store.set(Snippet { name: "b".into(), command: "updated".into(), description: String::new() });
1270 assert_eq!(store.snippets.len(), 3);
1271 assert_eq!(store.snippets[0].name, "a");
1272 assert_eq!(store.snippets[1].name, "b");
1273 assert_eq!(store.snippets[1].command, "updated");
1274 assert_eq!(store.snippets[2].name, "c");
1275 }
1276
1277 #[test]
1278 fn test_validate_command_with_tab() {
1279 assert!(validate_command("echo\thello").is_ok());
1280 }
1281
1282 #[test]
1283 fn test_validate_command_with_newline() {
1284 assert!(validate_command("echo\nhello").is_err());
1285 }
1286
1287 #[test]
1288 fn test_validate_name_newline() {
1289 assert!(validate_name("check\ndisk").is_err());
1290 }
1291
1292 #[test]
1297 fn test_shell_escape_simple() {
1298 assert_eq!(shell_escape("hello"), "'hello'");
1299 }
1300
1301 #[test]
1302 fn test_shell_escape_with_single_quote() {
1303 assert_eq!(shell_escape("it's"), "'it'\\''s'");
1304 }
1305
1306 #[test]
1307 fn test_shell_escape_with_spaces() {
1308 assert_eq!(shell_escape("hello world"), "'hello world'");
1309 }
1310
1311 #[test]
1312 fn test_shell_escape_with_semicolon() {
1313 assert_eq!(shell_escape("; rm -rf /"), "'; rm -rf /'");
1314 }
1315
1316 #[test]
1317 fn test_shell_escape_with_dollar() {
1318 assert_eq!(shell_escape("$(whoami)"), "'$(whoami)'");
1319 }
1320
1321 #[test]
1322 fn test_shell_escape_empty() {
1323 assert_eq!(shell_escape(""), "''");
1324 }
1325
1326 #[test]
1331 fn test_parse_params_none() {
1332 assert!(parse_params("df -h").is_empty());
1333 }
1334
1335 #[test]
1336 fn test_parse_params_single() {
1337 let params = parse_params("df -h {{path}}");
1338 assert_eq!(params.len(), 1);
1339 assert_eq!(params[0].name, "path");
1340 assert_eq!(params[0].default, None);
1341 }
1342
1343 #[test]
1344 fn test_parse_params_with_default() {
1345 let params = parse_params("df -h {{path:/var/log}}");
1346 assert_eq!(params.len(), 1);
1347 assert_eq!(params[0].name, "path");
1348 assert_eq!(params[0].default, Some("/var/log".to_string()));
1349 }
1350
1351 #[test]
1352 fn test_parse_params_multiple() {
1353 let params = parse_params("grep {{pattern}} {{file}}");
1354 assert_eq!(params.len(), 2);
1355 assert_eq!(params[0].name, "pattern");
1356 assert_eq!(params[1].name, "file");
1357 }
1358
1359 #[test]
1360 fn test_parse_params_deduplicate() {
1361 let params = parse_params("echo {{name}} {{name}}");
1362 assert_eq!(params.len(), 1);
1363 }
1364
1365 #[test]
1366 fn test_parse_params_invalid_name_skipped() {
1367 let params = parse_params("echo {{valid}} {{bad name}} {{ok}}");
1368 assert_eq!(params.len(), 2);
1369 assert_eq!(params[0].name, "valid");
1370 assert_eq!(params[1].name, "ok");
1371 }
1372
1373 #[test]
1374 fn test_parse_params_unclosed_brace() {
1375 let params = parse_params("echo {{unclosed");
1376 assert!(params.is_empty());
1377 }
1378
1379 #[test]
1380 fn test_parse_params_max_20() {
1381 let cmd: String = (0..25)
1382 .map(|i| format!("{{{{p{}}}}}", i))
1383 .collect::<Vec<_>>()
1384 .join(" ");
1385 let params = parse_params(&cmd);
1386 assert_eq!(params.len(), 20);
1387 }
1388
1389 #[test]
1394 fn test_validate_param_name_valid() {
1395 assert!(validate_param_name("path").is_ok());
1396 assert!(validate_param_name("my-param").is_ok());
1397 assert!(validate_param_name("my_param").is_ok());
1398 assert!(validate_param_name("param1").is_ok());
1399 }
1400
1401 #[test]
1402 fn test_validate_param_name_empty() {
1403 assert!(validate_param_name("").is_err());
1404 }
1405
1406 #[test]
1407 fn test_validate_param_name_rejects_braces() {
1408 assert!(validate_param_name("a{b").is_err());
1409 assert!(validate_param_name("a}b").is_err());
1410 }
1411
1412 #[test]
1413 fn test_validate_param_name_rejects_quote() {
1414 assert!(validate_param_name("it's").is_err());
1415 }
1416
1417 #[test]
1418 fn test_validate_param_name_rejects_whitespace() {
1419 assert!(validate_param_name("a b").is_err());
1420 }
1421
1422 #[test]
1427 fn test_substitute_simple() {
1428 let mut values = std::collections::HashMap::new();
1429 values.insert("path".to_string(), "/var/log".to_string());
1430 let result = substitute_params("df -h {{path}}", &values);
1431 assert_eq!(result, "df -h '/var/log'");
1432 }
1433
1434 #[test]
1435 fn test_substitute_with_default() {
1436 let values = std::collections::HashMap::new();
1437 let result = substitute_params("df -h {{path:/tmp}}", &values);
1438 assert_eq!(result, "df -h '/tmp'");
1439 }
1440
1441 #[test]
1442 fn test_substitute_overrides_default() {
1443 let mut values = std::collections::HashMap::new();
1444 values.insert("path".to_string(), "/home".to_string());
1445 let result = substitute_params("df -h {{path:/tmp}}", &values);
1446 assert_eq!(result, "df -h '/home'");
1447 }
1448
1449 #[test]
1450 fn test_substitute_escapes_injection() {
1451 let mut values = std::collections::HashMap::new();
1452 values.insert("name".to_string(), "; rm -rf /".to_string());
1453 let result = substitute_params("echo {{name}}", &values);
1454 assert_eq!(result, "echo '; rm -rf /'");
1455 }
1456
1457 #[test]
1458 fn test_substitute_no_recursive_expansion() {
1459 let mut values = std::collections::HashMap::new();
1460 values.insert("a".to_string(), "{{b}}".to_string());
1461 values.insert("b".to_string(), "gotcha".to_string());
1462 let result = substitute_params("echo {{a}}", &values);
1463 assert_eq!(result, "echo '{{b}}'");
1464 }
1465
1466 #[test]
1467 fn test_substitute_default_also_escaped() {
1468 let values = std::collections::HashMap::new();
1469 let result = substitute_params("echo {{x:$(whoami)}}", &values);
1470 assert_eq!(result, "echo '$(whoami)'");
1471 }
1472
1473 #[test]
1478 fn test_sanitize_plain_text() {
1479 assert_eq!(sanitize_output("hello world"), "hello world");
1480 }
1481
1482 #[test]
1483 fn test_sanitize_preserves_newlines_tabs() {
1484 assert_eq!(sanitize_output("line1\nline2\tok"), "line1\nline2\tok");
1485 }
1486
1487 #[test]
1488 fn test_sanitize_strips_csi() {
1489 assert_eq!(sanitize_output("\x1b[31mred\x1b[0m"), "red");
1490 }
1491
1492 #[test]
1493 fn test_sanitize_strips_osc_bel() {
1494 assert_eq!(sanitize_output("\x1b]0;title\x07text"), "text");
1495 }
1496
1497 #[test]
1498 fn test_sanitize_strips_osc_st() {
1499 assert_eq!(sanitize_output("\x1b]52;c;dGVzdA==\x1b\\text"), "text");
1500 }
1501
1502 #[test]
1503 fn test_sanitize_strips_c1_range() {
1504 assert_eq!(sanitize_output("a\u{0090}b\u{009C}c"), "abc");
1505 }
1506
1507 #[test]
1508 fn test_sanitize_strips_control_chars() {
1509 assert_eq!(sanitize_output("a\x01b\x07c"), "abc");
1510 }
1511
1512 #[test]
1513 fn test_sanitize_strips_dcs() {
1514 assert_eq!(sanitize_output("\x1bPdata\x1b\\text"), "text");
1515 }
1516
1517 #[test]
1522 fn test_shell_escape_only_single_quotes() {
1523 assert_eq!(shell_escape("'''"), "''\\'''\\'''\\'''");
1524 }
1525
1526 #[test]
1527 fn test_shell_escape_consecutive_single_quotes() {
1528 assert_eq!(shell_escape("a''b"), "'a'\\'''\\''b'");
1529 }
1530
1531 #[test]
1536 fn test_parse_params_adjacent() {
1537 let params = parse_params("{{a}}{{b}}");
1538 assert_eq!(params.len(), 2);
1539 assert_eq!(params[0].name, "a");
1540 assert_eq!(params[1].name, "b");
1541 }
1542
1543 #[test]
1544 fn test_parse_params_command_is_only_param() {
1545 let params = parse_params("{{cmd}}");
1546 assert_eq!(params.len(), 1);
1547 assert_eq!(params[0].name, "cmd");
1548 }
1549
1550 #[test]
1551 fn test_parse_params_nested_braces_rejected() {
1552 let params = parse_params("{{{a}}}");
1554 assert!(params.is_empty());
1555 }
1556
1557 #[test]
1558 fn test_parse_params_colon_empty_default() {
1559 let params = parse_params("echo {{name:}}");
1560 assert_eq!(params.len(), 1);
1561 assert_eq!(params[0].name, "name");
1562 assert_eq!(params[0].default, Some("".to_string()));
1563 }
1564
1565 #[test]
1566 fn test_parse_params_empty_inner() {
1567 let params = parse_params("echo {{}}");
1568 assert!(params.is_empty());
1569 }
1570
1571 #[test]
1572 fn test_parse_params_single_braces_ignored() {
1573 let params = parse_params("echo {notaparam}");
1574 assert!(params.is_empty());
1575 }
1576
1577 #[test]
1578 fn test_parse_params_default_with_colons() {
1579 let params = parse_params("{{url:http://localhost:8080}}");
1580 assert_eq!(params.len(), 1);
1581 assert_eq!(params[0].name, "url");
1582 assert_eq!(params[0].default, Some("http://localhost:8080".to_string()));
1583 }
1584
1585 #[test]
1590 fn test_validate_param_name_unicode() {
1591 assert!(validate_param_name("caf\u{00e9}").is_ok());
1592 }
1593
1594 #[test]
1595 fn test_validate_param_name_hyphen_only() {
1596 assert!(validate_param_name("-").is_ok());
1597 }
1598
1599 #[test]
1600 fn test_validate_param_name_underscore_only() {
1601 assert!(validate_param_name("_").is_ok());
1602 }
1603
1604 #[test]
1605 fn test_validate_param_name_rejects_dot() {
1606 assert!(validate_param_name("a.b").is_err());
1607 }
1608
1609 #[test]
1614 fn test_substitute_no_params_passthrough() {
1615 let values = std::collections::HashMap::new();
1616 let result = substitute_params("df -h /tmp", &values);
1617 assert_eq!(result, "df -h /tmp");
1618 }
1619
1620 #[test]
1621 fn test_substitute_missing_param_no_default() {
1622 let values = std::collections::HashMap::new();
1623 let result = substitute_params("echo {{name}}", &values);
1624 assert_eq!(result, "echo ''");
1625 }
1626
1627 #[test]
1628 fn test_substitute_empty_value_falls_to_default() {
1629 let mut values = std::collections::HashMap::new();
1630 values.insert("name".to_string(), "".to_string());
1631 let result = substitute_params("echo {{name:fallback}}", &values);
1632 assert_eq!(result, "echo 'fallback'");
1633 }
1634
1635 #[test]
1636 fn test_substitute_non_ascii_around_params() {
1637 let mut values = std::collections::HashMap::new();
1638 values.insert("x".to_string(), "val".to_string());
1639 let result = substitute_params("\u{00e9}cho {{x}} \u{2603}", &values);
1640 assert_eq!(result, "\u{00e9}cho 'val' \u{2603}");
1641 }
1642
1643 #[test]
1644 fn test_substitute_adjacent_params() {
1645 let mut values = std::collections::HashMap::new();
1646 values.insert("a".to_string(), "x".to_string());
1647 values.insert("b".to_string(), "y".to_string());
1648 let result = substitute_params("{{a}}{{b}}", &values);
1649 assert_eq!(result, "'x''y'");
1650 }
1651
1652 #[test]
1657 fn test_sanitize_empty() {
1658 assert_eq!(sanitize_output(""), "");
1659 }
1660
1661 #[test]
1662 fn test_sanitize_only_escapes() {
1663 assert_eq!(sanitize_output("\x1b[31m\x1b[0m\x1b[1m"), "");
1664 }
1665
1666 #[test]
1667 fn test_sanitize_lone_esc_at_end() {
1668 assert_eq!(sanitize_output("hello\x1b"), "hello");
1669 }
1670
1671 #[test]
1672 fn test_sanitize_truncated_csi_no_terminator() {
1673 assert_eq!(sanitize_output("hello\x1b[123"), "hello");
1674 }
1675
1676 #[test]
1677 fn test_sanitize_apc_sequence() {
1678 assert_eq!(sanitize_output("\x1b_payload\x1b\\visible"), "visible");
1679 }
1680
1681 #[test]
1682 fn test_sanitize_pm_sequence() {
1683 assert_eq!(sanitize_output("\x1b^payload\x1b\\visible"), "visible");
1684 }
1685
1686 #[test]
1687 fn test_sanitize_dcs_terminated_by_bel() {
1688 assert_eq!(sanitize_output("\x1bPdata\x07text"), "text");
1689 }
1690
1691 #[test]
1692 fn test_sanitize_lone_esc_plus_letter() {
1693 assert_eq!(sanitize_output("a\x1bMb"), "ab");
1694 }
1695
1696 #[test]
1697 fn test_sanitize_multiple_mixed_sequences() {
1698 let input = "\x1b[1mbold\x1b[0m \x1b]0;title\x07normal \x01gone";
1700 assert_eq!(sanitize_output(input), "bold normal gone");
1701 }
1702
1703}