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.trim().is_empty() {
162 return Err("Snippet name cannot be empty.".to_string());
163 }
164 if name != name.trim() {
165 return Err("Snippet name cannot have leading or trailing 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_ok());
1044 assert!(validate_name("check\tdisk").is_err()); assert!(validate_name(" ").is_err()); assert!(validate_name(" leading").is_err()); assert!(validate_name("trailing ").is_err()); }
1049
1050 #[test]
1051 fn test_validate_name_special_chars() {
1052 assert!(validate_name("check#disk").is_err());
1053 assert!(validate_name("[check]").is_err());
1054 }
1055
1056 #[test]
1057 fn test_validate_name_control_chars() {
1058 assert!(validate_name("check\x00disk").is_err());
1059 }
1060
1061 #[test]
1066 fn test_validate_command_valid() {
1067 assert!(validate_command("df -h").is_ok());
1068 assert!(validate_command("cat /etc/hosts | grep localhost").is_ok());
1069 assert!(validate_command("echo 'hello\tworld'").is_ok()); }
1071
1072 #[test]
1073 fn test_validate_command_empty() {
1074 assert!(validate_command("").is_err());
1075 }
1076
1077 #[test]
1078 fn test_validate_command_whitespace_only() {
1079 assert!(validate_command(" ").is_err());
1080 assert!(validate_command(" \t ").is_err());
1081 }
1082
1083 #[test]
1084 fn test_validate_command_control_chars() {
1085 assert!(validate_command("ls\x00-la").is_err());
1086 }
1087
1088 #[test]
1093 fn test_save_roundtrip() {
1094 let mut store = SnippetStore::default();
1095 store.set(Snippet {
1096 name: "check-disk".to_string(),
1097 command: "df -h".to_string(),
1098 description: "Check disk usage".to_string(),
1099 });
1100 store.set(Snippet {
1101 name: "uptime".to_string(),
1102 command: "uptime".to_string(),
1103 description: String::new(),
1104 });
1105
1106 let mut content = String::new();
1108 for (i, snippet) in store.snippets.iter().enumerate() {
1109 if i > 0 {
1110 content.push('\n');
1111 }
1112 content.push_str(&format!("[{}]\n", snippet.name));
1113 content.push_str(&format!("command={}\n", snippet.command));
1114 if !snippet.description.is_empty() {
1115 content.push_str(&format!("description={}\n", snippet.description));
1116 }
1117 }
1118
1119 let reparsed = SnippetStore::parse(&content);
1121 assert_eq!(reparsed.snippets.len(), 2);
1122 assert_eq!(reparsed.snippets[0].name, "check-disk");
1123 assert_eq!(reparsed.snippets[0].command, "df -h");
1124 assert_eq!(reparsed.snippets[0].description, "Check disk usage");
1125 assert_eq!(reparsed.snippets[1].name, "uptime");
1126 assert_eq!(reparsed.snippets[1].command, "uptime");
1127 assert!(reparsed.snippets[1].description.is_empty());
1128 }
1129
1130 #[test]
1131 fn test_save_to_temp_file() {
1132 let dir = std::env::temp_dir().join(format!("purple_snippet_test_{}", std::process::id()));
1133 let _ = std::fs::create_dir_all(&dir);
1134 let path = dir.join("snippets");
1135
1136 let mut store = SnippetStore {
1137 path_override: Some(path.clone()),
1138 ..Default::default()
1139 };
1140 store.set(Snippet {
1141 name: "test".to_string(),
1142 command: "echo hello".to_string(),
1143 description: "Test snippet".to_string(),
1144 });
1145 store.save().unwrap();
1146
1147 let content = std::fs::read_to_string(&path).unwrap();
1149 let reloaded = SnippetStore::parse(&content);
1150 assert_eq!(reloaded.snippets.len(), 1);
1151 assert_eq!(reloaded.snippets[0].name, "test");
1152 assert_eq!(reloaded.snippets[0].command, "echo hello");
1153
1154 let _ = std::fs::remove_dir_all(&dir);
1156 }
1157
1158 #[test]
1163 fn test_set_multiple_then_remove_all() {
1164 let mut store = SnippetStore::default();
1165 for name in ["a", "b", "c"] {
1166 store.set(Snippet {
1167 name: name.to_string(),
1168 command: "cmd".to_string(),
1169 description: String::new(),
1170 });
1171 }
1172 assert_eq!(store.snippets.len(), 3);
1173 store.remove("a");
1174 store.remove("b");
1175 store.remove("c");
1176 assert!(store.snippets.is_empty());
1177 }
1178
1179 #[test]
1180 fn test_snippet_with_complex_command() {
1181 let content = "[complex]\ncommand=for i in $(seq 1 5); do echo $i; done\n";
1182 let store = SnippetStore::parse(content);
1183 assert_eq!(
1184 store.snippets[0].command,
1185 "for i in $(seq 1 5); do echo $i; done"
1186 );
1187 }
1188
1189 #[test]
1190 fn test_snippet_command_with_pipes_and_redirects() {
1191 let content = "[logs]\ncommand=tail -100 /var/log/syslog | grep error | head -20\n";
1192 let store = SnippetStore::parse(content);
1193 assert_eq!(
1194 store.snippets[0].command,
1195 "tail -100 /var/log/syslog | grep error | head -20"
1196 );
1197 }
1198
1199 #[test]
1200 fn test_description_optional() {
1201 let content = "[check]\ncommand=ls\n";
1202 let store = SnippetStore::parse(content);
1203 assert!(store.snippets[0].description.is_empty());
1204 }
1205
1206 #[test]
1207 fn test_description_with_equals() {
1208 let content = "[env]\ncommand=env\ndescription=Check HOME= and PATH= vars\n";
1209 let store = SnippetStore::parse(content);
1210 assert_eq!(store.snippets[0].description, "Check HOME= and PATH= vars");
1211 }
1212
1213 #[test]
1214 fn test_name_with_equals_roundtrip() {
1215 let mut store = SnippetStore::default();
1216 store.set(Snippet {
1217 name: "check=disk".to_string(),
1218 command: "df -h".to_string(),
1219 description: String::new(),
1220 });
1221
1222 let mut content = String::new();
1223 for (i, snippet) in store.snippets.iter().enumerate() {
1224 if i > 0 {
1225 content.push('\n');
1226 }
1227 content.push_str(&format!("[{}]\n", snippet.name));
1228 content.push_str(&format!("command={}\n", snippet.command));
1229 if !snippet.description.is_empty() {
1230 content.push_str(&format!("description={}\n", snippet.description));
1231 }
1232 }
1233
1234 let reparsed = SnippetStore::parse(&content);
1235 assert_eq!(reparsed.snippets.len(), 1);
1236 assert_eq!(reparsed.snippets[0].name, "check=disk");
1237 }
1238
1239 #[test]
1240 fn test_validate_name_with_equals() {
1241 assert!(validate_name("check=disk").is_ok());
1242 }
1243
1244 #[test]
1245 fn test_parse_only_comments_and_blanks() {
1246 let content = "# comment\n\n# another\n";
1247 let store = SnippetStore::parse(content);
1248 assert!(store.snippets.is_empty());
1249 }
1250
1251 #[test]
1252 fn test_parse_section_without_close_bracket() {
1253 let content = "[incomplete\ncommand=ls\n";
1254 let store = SnippetStore::parse(content);
1255 assert!(store.snippets.is_empty());
1256 }
1257
1258 #[test]
1259 fn test_parse_trailing_content_after_last_section() {
1260 let content = "[check]\ncommand=ls\n";
1261 let store = SnippetStore::parse(content);
1262 assert_eq!(store.snippets.len(), 1);
1263 assert_eq!(store.snippets[0].command, "ls");
1264 }
1265
1266 #[test]
1267 fn test_set_overwrite_preserves_order() {
1268 let mut store = SnippetStore::default();
1269 store.set(Snippet { name: "a".into(), command: "1".into(), description: String::new() });
1270 store.set(Snippet { name: "b".into(), command: "2".into(), description: String::new() });
1271 store.set(Snippet { name: "c".into(), command: "3".into(), description: String::new() });
1272 store.set(Snippet { name: "b".into(), command: "updated".into(), description: String::new() });
1273 assert_eq!(store.snippets.len(), 3);
1274 assert_eq!(store.snippets[0].name, "a");
1275 assert_eq!(store.snippets[1].name, "b");
1276 assert_eq!(store.snippets[1].command, "updated");
1277 assert_eq!(store.snippets[2].name, "c");
1278 }
1279
1280 #[test]
1281 fn test_validate_command_with_tab() {
1282 assert!(validate_command("echo\thello").is_ok());
1283 }
1284
1285 #[test]
1286 fn test_validate_command_with_newline() {
1287 assert!(validate_command("echo\nhello").is_err());
1288 }
1289
1290 #[test]
1291 fn test_validate_name_newline() {
1292 assert!(validate_name("check\ndisk").is_err());
1293 }
1294
1295 #[test]
1300 fn test_shell_escape_simple() {
1301 assert_eq!(shell_escape("hello"), "'hello'");
1302 }
1303
1304 #[test]
1305 fn test_shell_escape_with_single_quote() {
1306 assert_eq!(shell_escape("it's"), "'it'\\''s'");
1307 }
1308
1309 #[test]
1310 fn test_shell_escape_with_spaces() {
1311 assert_eq!(shell_escape("hello world"), "'hello world'");
1312 }
1313
1314 #[test]
1315 fn test_shell_escape_with_semicolon() {
1316 assert_eq!(shell_escape("; rm -rf /"), "'; rm -rf /'");
1317 }
1318
1319 #[test]
1320 fn test_shell_escape_with_dollar() {
1321 assert_eq!(shell_escape("$(whoami)"), "'$(whoami)'");
1322 }
1323
1324 #[test]
1325 fn test_shell_escape_empty() {
1326 assert_eq!(shell_escape(""), "''");
1327 }
1328
1329 #[test]
1334 fn test_parse_params_none() {
1335 assert!(parse_params("df -h").is_empty());
1336 }
1337
1338 #[test]
1339 fn test_parse_params_single() {
1340 let params = parse_params("df -h {{path}}");
1341 assert_eq!(params.len(), 1);
1342 assert_eq!(params[0].name, "path");
1343 assert_eq!(params[0].default, None);
1344 }
1345
1346 #[test]
1347 fn test_parse_params_with_default() {
1348 let params = parse_params("df -h {{path:/var/log}}");
1349 assert_eq!(params.len(), 1);
1350 assert_eq!(params[0].name, "path");
1351 assert_eq!(params[0].default, Some("/var/log".to_string()));
1352 }
1353
1354 #[test]
1355 fn test_parse_params_multiple() {
1356 let params = parse_params("grep {{pattern}} {{file}}");
1357 assert_eq!(params.len(), 2);
1358 assert_eq!(params[0].name, "pattern");
1359 assert_eq!(params[1].name, "file");
1360 }
1361
1362 #[test]
1363 fn test_parse_params_deduplicate() {
1364 let params = parse_params("echo {{name}} {{name}}");
1365 assert_eq!(params.len(), 1);
1366 }
1367
1368 #[test]
1369 fn test_parse_params_invalid_name_skipped() {
1370 let params = parse_params("echo {{valid}} {{bad name}} {{ok}}");
1371 assert_eq!(params.len(), 2);
1372 assert_eq!(params[0].name, "valid");
1373 assert_eq!(params[1].name, "ok");
1374 }
1375
1376 #[test]
1377 fn test_parse_params_unclosed_brace() {
1378 let params = parse_params("echo {{unclosed");
1379 assert!(params.is_empty());
1380 }
1381
1382 #[test]
1383 fn test_parse_params_max_20() {
1384 let cmd: String = (0..25)
1385 .map(|i| format!("{{{{p{}}}}}", i))
1386 .collect::<Vec<_>>()
1387 .join(" ");
1388 let params = parse_params(&cmd);
1389 assert_eq!(params.len(), 20);
1390 }
1391
1392 #[test]
1397 fn test_validate_param_name_valid() {
1398 assert!(validate_param_name("path").is_ok());
1399 assert!(validate_param_name("my-param").is_ok());
1400 assert!(validate_param_name("my_param").is_ok());
1401 assert!(validate_param_name("param1").is_ok());
1402 }
1403
1404 #[test]
1405 fn test_validate_param_name_empty() {
1406 assert!(validate_param_name("").is_err());
1407 }
1408
1409 #[test]
1410 fn test_validate_param_name_rejects_braces() {
1411 assert!(validate_param_name("a{b").is_err());
1412 assert!(validate_param_name("a}b").is_err());
1413 }
1414
1415 #[test]
1416 fn test_validate_param_name_rejects_quote() {
1417 assert!(validate_param_name("it's").is_err());
1418 }
1419
1420 #[test]
1421 fn test_validate_param_name_rejects_whitespace() {
1422 assert!(validate_param_name("a b").is_err());
1423 }
1424
1425 #[test]
1430 fn test_substitute_simple() {
1431 let mut values = std::collections::HashMap::new();
1432 values.insert("path".to_string(), "/var/log".to_string());
1433 let result = substitute_params("df -h {{path}}", &values);
1434 assert_eq!(result, "df -h '/var/log'");
1435 }
1436
1437 #[test]
1438 fn test_substitute_with_default() {
1439 let values = std::collections::HashMap::new();
1440 let result = substitute_params("df -h {{path:/tmp}}", &values);
1441 assert_eq!(result, "df -h '/tmp'");
1442 }
1443
1444 #[test]
1445 fn test_substitute_overrides_default() {
1446 let mut values = std::collections::HashMap::new();
1447 values.insert("path".to_string(), "/home".to_string());
1448 let result = substitute_params("df -h {{path:/tmp}}", &values);
1449 assert_eq!(result, "df -h '/home'");
1450 }
1451
1452 #[test]
1453 fn test_substitute_escapes_injection() {
1454 let mut values = std::collections::HashMap::new();
1455 values.insert("name".to_string(), "; rm -rf /".to_string());
1456 let result = substitute_params("echo {{name}}", &values);
1457 assert_eq!(result, "echo '; rm -rf /'");
1458 }
1459
1460 #[test]
1461 fn test_substitute_no_recursive_expansion() {
1462 let mut values = std::collections::HashMap::new();
1463 values.insert("a".to_string(), "{{b}}".to_string());
1464 values.insert("b".to_string(), "gotcha".to_string());
1465 let result = substitute_params("echo {{a}}", &values);
1466 assert_eq!(result, "echo '{{b}}'");
1467 }
1468
1469 #[test]
1470 fn test_substitute_default_also_escaped() {
1471 let values = std::collections::HashMap::new();
1472 let result = substitute_params("echo {{x:$(whoami)}}", &values);
1473 assert_eq!(result, "echo '$(whoami)'");
1474 }
1475
1476 #[test]
1481 fn test_sanitize_plain_text() {
1482 assert_eq!(sanitize_output("hello world"), "hello world");
1483 }
1484
1485 #[test]
1486 fn test_sanitize_preserves_newlines_tabs() {
1487 assert_eq!(sanitize_output("line1\nline2\tok"), "line1\nline2\tok");
1488 }
1489
1490 #[test]
1491 fn test_sanitize_strips_csi() {
1492 assert_eq!(sanitize_output("\x1b[31mred\x1b[0m"), "red");
1493 }
1494
1495 #[test]
1496 fn test_sanitize_strips_osc_bel() {
1497 assert_eq!(sanitize_output("\x1b]0;title\x07text"), "text");
1498 }
1499
1500 #[test]
1501 fn test_sanitize_strips_osc_st() {
1502 assert_eq!(sanitize_output("\x1b]52;c;dGVzdA==\x1b\\text"), "text");
1503 }
1504
1505 #[test]
1506 fn test_sanitize_strips_c1_range() {
1507 assert_eq!(sanitize_output("a\u{0090}b\u{009C}c"), "abc");
1508 }
1509
1510 #[test]
1511 fn test_sanitize_strips_control_chars() {
1512 assert_eq!(sanitize_output("a\x01b\x07c"), "abc");
1513 }
1514
1515 #[test]
1516 fn test_sanitize_strips_dcs() {
1517 assert_eq!(sanitize_output("\x1bPdata\x1b\\text"), "text");
1518 }
1519
1520 #[test]
1525 fn test_shell_escape_only_single_quotes() {
1526 assert_eq!(shell_escape("'''"), "''\\'''\\'''\\'''");
1527 }
1528
1529 #[test]
1530 fn test_shell_escape_consecutive_single_quotes() {
1531 assert_eq!(shell_escape("a''b"), "'a'\\'''\\''b'");
1532 }
1533
1534 #[test]
1539 fn test_parse_params_adjacent() {
1540 let params = parse_params("{{a}}{{b}}");
1541 assert_eq!(params.len(), 2);
1542 assert_eq!(params[0].name, "a");
1543 assert_eq!(params[1].name, "b");
1544 }
1545
1546 #[test]
1547 fn test_parse_params_command_is_only_param() {
1548 let params = parse_params("{{cmd}}");
1549 assert_eq!(params.len(), 1);
1550 assert_eq!(params[0].name, "cmd");
1551 }
1552
1553 #[test]
1554 fn test_parse_params_nested_braces_rejected() {
1555 let params = parse_params("{{{a}}}");
1557 assert!(params.is_empty());
1558 }
1559
1560 #[test]
1561 fn test_parse_params_colon_empty_default() {
1562 let params = parse_params("echo {{name:}}");
1563 assert_eq!(params.len(), 1);
1564 assert_eq!(params[0].name, "name");
1565 assert_eq!(params[0].default, Some("".to_string()));
1566 }
1567
1568 #[test]
1569 fn test_parse_params_empty_inner() {
1570 let params = parse_params("echo {{}}");
1571 assert!(params.is_empty());
1572 }
1573
1574 #[test]
1575 fn test_parse_params_single_braces_ignored() {
1576 let params = parse_params("echo {notaparam}");
1577 assert!(params.is_empty());
1578 }
1579
1580 #[test]
1581 fn test_parse_params_default_with_colons() {
1582 let params = parse_params("{{url:http://localhost:8080}}");
1583 assert_eq!(params.len(), 1);
1584 assert_eq!(params[0].name, "url");
1585 assert_eq!(params[0].default, Some("http://localhost:8080".to_string()));
1586 }
1587
1588 #[test]
1593 fn test_validate_param_name_unicode() {
1594 assert!(validate_param_name("caf\u{00e9}").is_ok());
1595 }
1596
1597 #[test]
1598 fn test_validate_param_name_hyphen_only() {
1599 assert!(validate_param_name("-").is_ok());
1600 }
1601
1602 #[test]
1603 fn test_validate_param_name_underscore_only() {
1604 assert!(validate_param_name("_").is_ok());
1605 }
1606
1607 #[test]
1608 fn test_validate_param_name_rejects_dot() {
1609 assert!(validate_param_name("a.b").is_err());
1610 }
1611
1612 #[test]
1617 fn test_substitute_no_params_passthrough() {
1618 let values = std::collections::HashMap::new();
1619 let result = substitute_params("df -h /tmp", &values);
1620 assert_eq!(result, "df -h /tmp");
1621 }
1622
1623 #[test]
1624 fn test_substitute_missing_param_no_default() {
1625 let values = std::collections::HashMap::new();
1626 let result = substitute_params("echo {{name}}", &values);
1627 assert_eq!(result, "echo ''");
1628 }
1629
1630 #[test]
1631 fn test_substitute_empty_value_falls_to_default() {
1632 let mut values = std::collections::HashMap::new();
1633 values.insert("name".to_string(), "".to_string());
1634 let result = substitute_params("echo {{name:fallback}}", &values);
1635 assert_eq!(result, "echo 'fallback'");
1636 }
1637
1638 #[test]
1639 fn test_substitute_non_ascii_around_params() {
1640 let mut values = std::collections::HashMap::new();
1641 values.insert("x".to_string(), "val".to_string());
1642 let result = substitute_params("\u{00e9}cho {{x}} \u{2603}", &values);
1643 assert_eq!(result, "\u{00e9}cho 'val' \u{2603}");
1644 }
1645
1646 #[test]
1647 fn test_substitute_adjacent_params() {
1648 let mut values = std::collections::HashMap::new();
1649 values.insert("a".to_string(), "x".to_string());
1650 values.insert("b".to_string(), "y".to_string());
1651 let result = substitute_params("{{a}}{{b}}", &values);
1652 assert_eq!(result, "'x''y'");
1653 }
1654
1655 #[test]
1660 fn test_sanitize_empty() {
1661 assert_eq!(sanitize_output(""), "");
1662 }
1663
1664 #[test]
1665 fn test_sanitize_only_escapes() {
1666 assert_eq!(sanitize_output("\x1b[31m\x1b[0m\x1b[1m"), "");
1667 }
1668
1669 #[test]
1670 fn test_sanitize_lone_esc_at_end() {
1671 assert_eq!(sanitize_output("hello\x1b"), "hello");
1672 }
1673
1674 #[test]
1675 fn test_sanitize_truncated_csi_no_terminator() {
1676 assert_eq!(sanitize_output("hello\x1b[123"), "hello");
1677 }
1678
1679 #[test]
1680 fn test_sanitize_apc_sequence() {
1681 assert_eq!(sanitize_output("\x1b_payload\x1b\\visible"), "visible");
1682 }
1683
1684 #[test]
1685 fn test_sanitize_pm_sequence() {
1686 assert_eq!(sanitize_output("\x1b^payload\x1b\\visible"), "visible");
1687 }
1688
1689 #[test]
1690 fn test_sanitize_dcs_terminated_by_bel() {
1691 assert_eq!(sanitize_output("\x1bPdata\x07text"), "text");
1692 }
1693
1694 #[test]
1695 fn test_sanitize_lone_esc_plus_letter() {
1696 assert_eq!(sanitize_output("a\x1bMb"), "ab");
1697 }
1698
1699 #[test]
1700 fn test_sanitize_multiple_mixed_sequences() {
1701 let input = "\x1b[1mbold\x1b[0m \x1b]0;title\x07normal \x01gone";
1703 assert_eq!(sanitize_output(input), "bold normal gone");
1704 }
1705
1706}