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() && !snippets.iter().any(|s| s.name == snippet.name) {
97 snippets.push(snippet);
98 }
99 }
100 Self {
101 snippets,
102 path_override: None,
103 }
104 }
105
106 pub fn save(&self) -> io::Result<()> {
108 if crate::demo_flag::is_demo() {
109 return Ok(());
110 }
111 let path = match &self.path_override {
112 Some(p) => p.clone(),
113 None => match config_path() {
114 Some(p) => p,
115 None => {
116 return Err(io::Error::new(
117 io::ErrorKind::NotFound,
118 "Could not determine home directory",
119 ));
120 }
121 },
122 };
123
124 let mut content = String::new();
125 for (i, snippet) in self.snippets.iter().enumerate() {
126 if i > 0 {
127 content.push('\n');
128 }
129 content.push_str(&format!("[{}]\n", snippet.name));
130 content.push_str(&format!("command={}\n", snippet.command));
131 if !snippet.description.is_empty() {
132 content.push_str(&format!("description={}\n", snippet.description));
133 }
134 }
135
136 fs_util::atomic_write(&path, content.as_bytes())
137 }
138
139 pub fn get(&self, name: &str) -> Option<&Snippet> {
141 self.snippets.iter().find(|s| s.name == name)
142 }
143
144 pub fn set(&mut self, snippet: Snippet) {
146 if let Some(existing) = self.snippets.iter_mut().find(|s| s.name == snippet.name) {
147 *existing = snippet;
148 } else {
149 self.snippets.push(snippet);
150 }
151 }
152
153 pub fn remove(&mut self, name: &str) {
155 self.snippets.retain(|s| s.name != name);
156 }
157}
158
159pub fn validate_name(name: &str) -> Result<(), String> {
162 if name.trim().is_empty() {
163 return Err("Snippet name cannot be empty.".to_string());
164 }
165 if name != name.trim() {
166 return Err("Snippet name cannot have leading or trailing whitespace.".to_string());
167 }
168 if name.contains('#') || name.contains('[') || name.contains(']') {
169 return Err("Snippet name cannot contain #, [ or ].".to_string());
170 }
171 if name.contains(|c: char| c.is_control()) {
172 return Err("Snippet name cannot contain control characters.".to_string());
173 }
174 Ok(())
175}
176
177pub fn validate_command(command: &str) -> Result<(), String> {
179 if command.trim().is_empty() {
180 return Err("Command cannot be empty.".to_string());
181 }
182 if command.contains(|c: char| c.is_control() && c != '\t') {
183 return Err("Command cannot contain control characters.".to_string());
184 }
185 Ok(())
186}
187
188#[derive(Debug, Clone, PartialEq)]
194pub struct SnippetParam {
195 pub name: String,
196 pub default: Option<String>,
197}
198
199pub fn shell_escape(s: &str) -> String {
202 format!("'{}'", s.replace('\'', "'\\''"))
203}
204
205pub fn parse_params(command: &str) -> Vec<SnippetParam> {
208 let mut params = Vec::new();
209 let mut seen = std::collections::HashSet::new();
210 let bytes = command.as_bytes();
211 let len = bytes.len();
212 let mut i = 0;
213 while i + 3 < len {
214 if bytes[i] == b'{' && bytes.get(i + 1) == Some(&b'{') {
215 if let Some(end) = command[i + 2..].find("}}") {
216 let inner = &command[i + 2..i + 2 + end];
217 let (name, default) = if let Some((n, d)) = inner.split_once(':') {
218 (n.to_string(), Some(d.to_string()))
219 } else {
220 (inner.to_string(), None)
221 };
222 if validate_param_name(&name).is_ok() && !seen.contains(&name) && params.len() < 20
223 {
224 seen.insert(name.clone());
225 params.push(SnippetParam { name, default });
226 }
227 i = i + 2 + end + 2;
228 continue;
229 }
230 }
231 i += 1;
232 }
233 params
234}
235
236pub fn validate_param_name(name: &str) -> Result<(), String> {
239 if name.is_empty() {
240 return Err("Parameter name cannot be empty.".to_string());
241 }
242 if !name
243 .chars()
244 .all(|c| c.is_alphanumeric() || c == '_' || c == '-')
245 {
246 return Err(format!(
247 "Parameter name '{}' contains invalid characters.",
248 name
249 ));
250 }
251 Ok(())
252}
253
254pub fn substitute_params(
257 command: &str,
258 values: &std::collections::HashMap<String, String>,
259) -> String {
260 let mut result = String::with_capacity(command.len());
261 let bytes = command.as_bytes();
262 let len = bytes.len();
263 let mut i = 0;
264 while i < len {
265 if i + 3 < len && bytes[i] == b'{' && bytes[i + 1] == b'{' {
266 if let Some(end) = command[i + 2..].find("}}") {
267 let inner = &command[i + 2..i + 2 + end];
268 let (name, default) = if let Some((n, d)) = inner.split_once(':') {
269 (n, Some(d))
270 } else {
271 (inner, None)
272 };
273 let value = values
274 .get(name)
275 .filter(|v| !v.is_empty())
276 .map(|v| v.as_str())
277 .or(default)
278 .unwrap_or("");
279 result.push_str(&shell_escape(value));
280 i = i + 2 + end + 2;
281 continue;
282 }
283 }
284 let ch = command[i..].chars().next().unwrap();
286 result.push(ch);
287 i += ch.len_utf8();
288 }
289 result
290}
291
292pub fn sanitize_output(input: &str) -> String {
299 let mut out = String::with_capacity(input.len());
300 let mut chars = input.chars().peekable();
301 while let Some(c) = chars.next() {
302 match c {
303 '\x1b' => {
304 match chars.peek() {
305 Some('[') => {
306 chars.next();
307 while let Some(&ch) = chars.peek() {
309 chars.next();
310 if ('\x40'..='\x7e').contains(&ch) {
311 break;
312 }
313 }
314 }
315 Some(']') | Some('P') | Some('X') | Some('^') | Some('_') => {
316 chars.next();
317 consume_until_st(&mut chars);
319 }
320 _ => {
321 chars.next();
323 }
324 }
325 }
326 c if ('\u{0080}'..='\u{009F}').contains(&c) => {
327 }
329 c if c.is_control() && c != '\n' && c != '\t' => {
330 }
332 _ => out.push(c),
333 }
334 }
335 out
336}
337
338fn consume_until_st(chars: &mut std::iter::Peekable<std::str::Chars<'_>>) {
340 while let Some(&ch) = chars.peek() {
341 if ch == '\x07' {
342 chars.next();
343 break;
344 }
345 if ch == '\x1b' {
346 chars.next();
347 if chars.peek() == Some(&'\\') {
348 chars.next();
349 }
350 break;
351 }
352 chars.next();
353 }
354}
355
356const MAX_OUTPUT_LINES: usize = 10_000;
363
364pub enum SnippetEvent {
367 HostDone {
368 run_id: u64,
369 alias: String,
370 stdout: String,
371 stderr: String,
372 exit_code: Option<i32>,
373 },
374 Progress {
375 run_id: u64,
376 completed: usize,
377 total: usize,
378 },
379 AllDone {
380 run_id: u64,
381 },
382}
383
384pub struct ChildGuard {
387 inner: std::sync::Mutex<Option<std::process::Child>>,
388 pgid: i32,
389}
390
391impl ChildGuard {
392 fn new(child: std::process::Child) -> Self {
393 let pgid = i32::try_from(child.id()).unwrap_or(-1);
397 Self {
398 inner: std::sync::Mutex::new(Some(child)),
399 pgid,
400 }
401 }
402}
403
404impl Drop for ChildGuard {
405 fn drop(&mut self) {
406 let mut lock = self.inner.lock().unwrap_or_else(|e| e.into_inner());
407 if let Some(ref mut child) = *lock {
408 if let Ok(Some(_)) = child.try_wait() {
410 return;
411 }
412 #[cfg(unix)]
414 unsafe {
415 libc::kill(-self.pgid, libc::SIGTERM);
416 }
417 let deadline = std::time::Instant::now() + std::time::Duration::from_millis(500);
419 loop {
420 if let Ok(Some(_)) = child.try_wait() {
421 return;
422 }
423 if std::time::Instant::now() >= deadline {
424 break;
425 }
426 std::thread::sleep(std::time::Duration::from_millis(50));
427 }
428 #[cfg(unix)]
430 unsafe {
431 libc::kill(-self.pgid, libc::SIGKILL);
432 }
433 let _ = child.kill();
435 let _ = child.wait();
436 }
437 }
438}
439
440fn read_pipe_capped<R: io::Read>(reader: R) -> String {
443 use io::BufRead;
444 let mut reader = io::BufReader::new(reader);
445 let mut output = String::new();
446 let mut line_count = 0;
447 let mut capped = false;
448 let mut buf = Vec::new();
449 loop {
450 buf.clear();
451 match reader.read_until(b'\n', &mut buf) {
452 Ok(0) => break, Ok(_) => {
454 if !capped {
455 if line_count < MAX_OUTPUT_LINES {
456 if line_count > 0 {
457 output.push('\n');
458 }
459 if buf.last() == Some(&b'\n') {
461 buf.pop();
462 if buf.last() == Some(&b'\r') {
463 buf.pop();
464 }
465 }
466 output.push_str(&String::from_utf8_lossy(&buf));
468 line_count += 1;
469 } else {
470 output.push_str("\n[Output truncated at 10,000 lines]");
471 capped = true;
472 }
473 }
474 }
476 Err(_) => break,
477 }
478 }
479 output
480}
481
482fn base_ssh_command(
486 alias: &str,
487 config_path: &Path,
488 command: &str,
489 askpass: Option<&str>,
490 bw_session: Option<&str>,
491 has_active_tunnel: bool,
492) -> Command {
493 let mut cmd = Command::new("ssh");
494 cmd.arg("-F")
495 .arg(config_path)
496 .arg("-o")
497 .arg("ConnectTimeout=10")
498 .arg("-o")
499 .arg("ControlMaster=no")
500 .arg("-o")
501 .arg("ControlPath=none");
502
503 if has_active_tunnel {
504 cmd.arg("-o").arg("ClearAllForwardings=yes");
505 }
506
507 cmd.arg("--").arg(alias).arg(command);
508
509 if askpass.is_some() {
510 crate::askpass_env::configure_ssh_command(&mut cmd, alias, config_path);
511 }
512
513 if let Some(token) = bw_session {
514 cmd.env("BW_SESSION", token);
515 }
516
517 cmd
518}
519
520fn build_snippet_command(
522 alias: &str,
523 config_path: &Path,
524 command: &str,
525 askpass: Option<&str>,
526 bw_session: Option<&str>,
527 has_active_tunnel: bool,
528) -> Command {
529 let mut cmd = base_ssh_command(
530 alias,
531 config_path,
532 command,
533 askpass,
534 bw_session,
535 has_active_tunnel,
536 );
537 cmd.stdin(Stdio::null())
538 .stdout(Stdio::piped())
539 .stderr(Stdio::piped());
540
541 #[cfg(unix)]
544 unsafe {
545 use std::os::unix::process::CommandExt;
546 cmd.pre_exec(|| {
547 libc::setpgid(0, 0);
548 Ok(())
549 });
550 }
551
552 cmd
553}
554
555#[allow(clippy::too_many_arguments)]
557fn execute_host(
558 run_id: u64,
559 alias: &str,
560 config_path: &Path,
561 command: &str,
562 askpass: Option<&str>,
563 bw_session: Option<&str>,
564 has_active_tunnel: bool,
565 tx: &std::sync::mpsc::Sender<SnippetEvent>,
566) -> Option<std::sync::Arc<ChildGuard>> {
567 let mut cmd = build_snippet_command(
568 alias,
569 config_path,
570 command,
571 askpass,
572 bw_session,
573 has_active_tunnel,
574 );
575
576 match cmd.spawn() {
577 Ok(child) => {
578 let guard = std::sync::Arc::new(ChildGuard::new(child));
579
580 let stdout_pipe = {
582 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
583 lock.as_mut().and_then(|c| c.stdout.take())
584 };
585 let stderr_pipe = {
586 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
587 lock.as_mut().and_then(|c| c.stderr.take())
588 };
589
590 let stdout_handle = std::thread::spawn(move || match stdout_pipe {
592 Some(pipe) => read_pipe_capped(pipe),
593 None => String::new(),
594 });
595 let stderr_handle = std::thread::spawn(move || match stderr_pipe {
596 Some(pipe) => read_pipe_capped(pipe),
597 None => String::new(),
598 });
599
600 let stdout_text = stdout_handle.join().unwrap_or_default();
602 let stderr_text = stderr_handle.join().unwrap_or_default();
603
604 let exit_code = {
607 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
608 let status = lock.as_mut().and_then(|c| c.wait().ok());
609 let _ = lock.take(); status.and_then(|s| {
611 #[cfg(unix)]
612 {
613 use std::os::unix::process::ExitStatusExt;
614 s.code().or_else(|| s.signal().map(|sig| 128 + sig))
615 }
616 #[cfg(not(unix))]
617 {
618 s.code()
619 }
620 })
621 };
622
623 let _ = tx.send(SnippetEvent::HostDone {
624 run_id,
625 alias: alias.to_string(),
626 stdout: sanitize_output(&stdout_text),
627 stderr: sanitize_output(&stderr_text),
628 exit_code,
629 });
630
631 Some(guard)
632 }
633 Err(e) => {
634 let _ = tx.send(SnippetEvent::HostDone {
635 run_id,
636 alias: alias.to_string(),
637 stdout: String::new(),
638 stderr: format!("Failed to launch ssh: {}", e),
639 exit_code: None,
640 });
641 None
642 }
643 }
644}
645
646#[allow(clippy::too_many_arguments)]
649pub fn spawn_snippet_execution(
650 run_id: u64,
651 askpass_map: Vec<(String, Option<String>)>,
652 config_path: PathBuf,
653 command: String,
654 bw_session: Option<String>,
655 tunnel_aliases: std::collections::HashSet<String>,
656 cancel: std::sync::Arc<std::sync::atomic::AtomicBool>,
657 tx: std::sync::mpsc::Sender<SnippetEvent>,
658 parallel: bool,
659) {
660 let total = askpass_map.len();
661 let max_concurrent: usize = 20;
662
663 std::thread::Builder::new()
664 .name("snippet-coordinator".into())
665 .spawn(move || {
666 let guards: std::sync::Arc<std::sync::Mutex<Vec<std::sync::Arc<ChildGuard>>>> =
667 std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
668
669 if parallel && total > 1 {
670 let (slot_tx, slot_rx) = std::sync::mpsc::channel::<()>();
672 for _ in 0..max_concurrent.min(total) {
673 let _ = slot_tx.send(());
674 }
675
676 let completed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
677 let mut worker_handles = Vec::new();
678
679 for (alias, askpass) in askpass_map {
680 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
681 break;
682 }
683
684 loop {
686 match slot_rx.recv_timeout(std::time::Duration::from_millis(100)) {
687 Ok(()) => break,
688 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
689 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
690 break;
691 }
692 }
693 Err(_) => break, }
695 }
696
697 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
698 break;
699 }
700
701 let config_path = config_path.clone();
702 let command = command.clone();
703 let bw_session = bw_session.clone();
704 let has_tunnel = tunnel_aliases.contains(&alias);
705 let tx = tx.clone();
706 let slot_tx = slot_tx.clone();
707 let guards = guards.clone();
708 let completed = completed.clone();
709 let total = total;
710
711 let handle = std::thread::spawn(move || {
712 struct SlotRelease(Option<std::sync::mpsc::Sender<()>>);
714 impl Drop for SlotRelease {
715 fn drop(&mut self) {
716 if let Some(tx) = self.0.take() {
717 let _ = tx.send(());
718 }
719 }
720 }
721 let _slot = SlotRelease(Some(slot_tx));
722
723 let guard = execute_host(
724 run_id,
725 &alias,
726 &config_path,
727 &command,
728 askpass.as_deref(),
729 bw_session.as_deref(),
730 has_tunnel,
731 &tx,
732 );
733
734 if let Some(g) = guard {
736 guards.lock().unwrap_or_else(|e| e.into_inner()).push(g);
737 }
738
739 let c = completed.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
740 let _ = tx.send(SnippetEvent::Progress {
741 run_id,
742 completed: c,
743 total,
744 });
745 });
747 worker_handles.push(handle);
748 }
749
750 for handle in worker_handles {
752 let _ = handle.join();
753 }
754 } else {
755 for (i, (alias, askpass)) in askpass_map.into_iter().enumerate() {
757 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
758 break;
759 }
760
761 let has_tunnel = tunnel_aliases.contains(&alias);
762 let guard = execute_host(
763 run_id,
764 &alias,
765 &config_path,
766 &command,
767 askpass.as_deref(),
768 bw_session.as_deref(),
769 has_tunnel,
770 &tx,
771 );
772
773 if let Some(g) = guard {
774 guards.lock().unwrap_or_else(|e| e.into_inner()).push(g);
775 }
776
777 let _ = tx.send(SnippetEvent::Progress {
778 run_id,
779 completed: i + 1,
780 total,
781 });
782 }
783 }
784
785 let _ = tx.send(SnippetEvent::AllDone { run_id });
786 })
788 .expect("failed to spawn snippet coordinator");
789}
790
791pub fn run_snippet(
796 alias: &str,
797 config_path: &Path,
798 command: &str,
799 askpass: Option<&str>,
800 bw_session: Option<&str>,
801 capture: bool,
802 has_active_tunnel: bool,
803) -> anyhow::Result<SnippetResult> {
804 let mut cmd = base_ssh_command(
805 alias,
806 config_path,
807 command,
808 askpass,
809 bw_session,
810 has_active_tunnel,
811 );
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 {
1270 name: "a".into(),
1271 command: "1".into(),
1272 description: String::new(),
1273 });
1274 store.set(Snippet {
1275 name: "b".into(),
1276 command: "2".into(),
1277 description: String::new(),
1278 });
1279 store.set(Snippet {
1280 name: "c".into(),
1281 command: "3".into(),
1282 description: String::new(),
1283 });
1284 store.set(Snippet {
1285 name: "b".into(),
1286 command: "updated".into(),
1287 description: String::new(),
1288 });
1289 assert_eq!(store.snippets.len(), 3);
1290 assert_eq!(store.snippets[0].name, "a");
1291 assert_eq!(store.snippets[1].name, "b");
1292 assert_eq!(store.snippets[1].command, "updated");
1293 assert_eq!(store.snippets[2].name, "c");
1294 }
1295
1296 #[test]
1297 fn test_validate_command_with_tab() {
1298 assert!(validate_command("echo\thello").is_ok());
1299 }
1300
1301 #[test]
1302 fn test_validate_command_with_newline() {
1303 assert!(validate_command("echo\nhello").is_err());
1304 }
1305
1306 #[test]
1307 fn test_validate_name_newline() {
1308 assert!(validate_name("check\ndisk").is_err());
1309 }
1310
1311 #[test]
1316 fn test_shell_escape_simple() {
1317 assert_eq!(shell_escape("hello"), "'hello'");
1318 }
1319
1320 #[test]
1321 fn test_shell_escape_with_single_quote() {
1322 assert_eq!(shell_escape("it's"), "'it'\\''s'");
1323 }
1324
1325 #[test]
1326 fn test_shell_escape_with_spaces() {
1327 assert_eq!(shell_escape("hello world"), "'hello world'");
1328 }
1329
1330 #[test]
1331 fn test_shell_escape_with_semicolon() {
1332 assert_eq!(shell_escape("; rm -rf /"), "'; rm -rf /'");
1333 }
1334
1335 #[test]
1336 fn test_shell_escape_with_dollar() {
1337 assert_eq!(shell_escape("$(whoami)"), "'$(whoami)'");
1338 }
1339
1340 #[test]
1341 fn test_shell_escape_empty() {
1342 assert_eq!(shell_escape(""), "''");
1343 }
1344
1345 #[test]
1350 fn test_parse_params_none() {
1351 assert!(parse_params("df -h").is_empty());
1352 }
1353
1354 #[test]
1355 fn test_parse_params_single() {
1356 let params = parse_params("df -h {{path}}");
1357 assert_eq!(params.len(), 1);
1358 assert_eq!(params[0].name, "path");
1359 assert_eq!(params[0].default, None);
1360 }
1361
1362 #[test]
1363 fn test_parse_params_with_default() {
1364 let params = parse_params("df -h {{path:/var/log}}");
1365 assert_eq!(params.len(), 1);
1366 assert_eq!(params[0].name, "path");
1367 assert_eq!(params[0].default, Some("/var/log".to_string()));
1368 }
1369
1370 #[test]
1371 fn test_parse_params_multiple() {
1372 let params = parse_params("grep {{pattern}} {{file}}");
1373 assert_eq!(params.len(), 2);
1374 assert_eq!(params[0].name, "pattern");
1375 assert_eq!(params[1].name, "file");
1376 }
1377
1378 #[test]
1379 fn test_parse_params_deduplicate() {
1380 let params = parse_params("echo {{name}} {{name}}");
1381 assert_eq!(params.len(), 1);
1382 }
1383
1384 #[test]
1385 fn test_parse_params_invalid_name_skipped() {
1386 let params = parse_params("echo {{valid}} {{bad name}} {{ok}}");
1387 assert_eq!(params.len(), 2);
1388 assert_eq!(params[0].name, "valid");
1389 assert_eq!(params[1].name, "ok");
1390 }
1391
1392 #[test]
1393 fn test_parse_params_unclosed_brace() {
1394 let params = parse_params("echo {{unclosed");
1395 assert!(params.is_empty());
1396 }
1397
1398 #[test]
1399 fn test_parse_params_max_20() {
1400 let cmd: String = (0..25)
1401 .map(|i| format!("{{{{p{}}}}}", i))
1402 .collect::<Vec<_>>()
1403 .join(" ");
1404 let params = parse_params(&cmd);
1405 assert_eq!(params.len(), 20);
1406 }
1407
1408 #[test]
1413 fn test_validate_param_name_valid() {
1414 assert!(validate_param_name("path").is_ok());
1415 assert!(validate_param_name("my-param").is_ok());
1416 assert!(validate_param_name("my_param").is_ok());
1417 assert!(validate_param_name("param1").is_ok());
1418 }
1419
1420 #[test]
1421 fn test_validate_param_name_empty() {
1422 assert!(validate_param_name("").is_err());
1423 }
1424
1425 #[test]
1426 fn test_validate_param_name_rejects_braces() {
1427 assert!(validate_param_name("a{b").is_err());
1428 assert!(validate_param_name("a}b").is_err());
1429 }
1430
1431 #[test]
1432 fn test_validate_param_name_rejects_quote() {
1433 assert!(validate_param_name("it's").is_err());
1434 }
1435
1436 #[test]
1437 fn test_validate_param_name_rejects_whitespace() {
1438 assert!(validate_param_name("a b").is_err());
1439 }
1440
1441 #[test]
1446 fn test_substitute_simple() {
1447 let mut values = std::collections::HashMap::new();
1448 values.insert("path".to_string(), "/var/log".to_string());
1449 let result = substitute_params("df -h {{path}}", &values);
1450 assert_eq!(result, "df -h '/var/log'");
1451 }
1452
1453 #[test]
1454 fn test_substitute_with_default() {
1455 let values = std::collections::HashMap::new();
1456 let result = substitute_params("df -h {{path:/tmp}}", &values);
1457 assert_eq!(result, "df -h '/tmp'");
1458 }
1459
1460 #[test]
1461 fn test_substitute_overrides_default() {
1462 let mut values = std::collections::HashMap::new();
1463 values.insert("path".to_string(), "/home".to_string());
1464 let result = substitute_params("df -h {{path:/tmp}}", &values);
1465 assert_eq!(result, "df -h '/home'");
1466 }
1467
1468 #[test]
1469 fn test_substitute_escapes_injection() {
1470 let mut values = std::collections::HashMap::new();
1471 values.insert("name".to_string(), "; rm -rf /".to_string());
1472 let result = substitute_params("echo {{name}}", &values);
1473 assert_eq!(result, "echo '; rm -rf /'");
1474 }
1475
1476 #[test]
1477 fn test_substitute_no_recursive_expansion() {
1478 let mut values = std::collections::HashMap::new();
1479 values.insert("a".to_string(), "{{b}}".to_string());
1480 values.insert("b".to_string(), "gotcha".to_string());
1481 let result = substitute_params("echo {{a}}", &values);
1482 assert_eq!(result, "echo '{{b}}'");
1483 }
1484
1485 #[test]
1486 fn test_substitute_default_also_escaped() {
1487 let values = std::collections::HashMap::new();
1488 let result = substitute_params("echo {{x:$(whoami)}}", &values);
1489 assert_eq!(result, "echo '$(whoami)'");
1490 }
1491
1492 #[test]
1497 fn test_sanitize_plain_text() {
1498 assert_eq!(sanitize_output("hello world"), "hello world");
1499 }
1500
1501 #[test]
1502 fn test_sanitize_preserves_newlines_tabs() {
1503 assert_eq!(sanitize_output("line1\nline2\tok"), "line1\nline2\tok");
1504 }
1505
1506 #[test]
1507 fn test_sanitize_strips_csi() {
1508 assert_eq!(sanitize_output("\x1b[31mred\x1b[0m"), "red");
1509 }
1510
1511 #[test]
1512 fn test_sanitize_strips_osc_bel() {
1513 assert_eq!(sanitize_output("\x1b]0;title\x07text"), "text");
1514 }
1515
1516 #[test]
1517 fn test_sanitize_strips_osc_st() {
1518 assert_eq!(sanitize_output("\x1b]52;c;dGVzdA==\x1b\\text"), "text");
1519 }
1520
1521 #[test]
1522 fn test_sanitize_strips_c1_range() {
1523 assert_eq!(sanitize_output("a\u{0090}b\u{009C}c"), "abc");
1524 }
1525
1526 #[test]
1527 fn test_sanitize_strips_control_chars() {
1528 assert_eq!(sanitize_output("a\x01b\x07c"), "abc");
1529 }
1530
1531 #[test]
1532 fn test_sanitize_strips_dcs() {
1533 assert_eq!(sanitize_output("\x1bPdata\x1b\\text"), "text");
1534 }
1535
1536 #[test]
1541 fn test_shell_escape_only_single_quotes() {
1542 assert_eq!(shell_escape("'''"), "''\\'''\\'''\\'''");
1543 }
1544
1545 #[test]
1546 fn test_shell_escape_consecutive_single_quotes() {
1547 assert_eq!(shell_escape("a''b"), "'a'\\'''\\''b'");
1548 }
1549
1550 #[test]
1555 fn test_parse_params_adjacent() {
1556 let params = parse_params("{{a}}{{b}}");
1557 assert_eq!(params.len(), 2);
1558 assert_eq!(params[0].name, "a");
1559 assert_eq!(params[1].name, "b");
1560 }
1561
1562 #[test]
1563 fn test_parse_params_command_is_only_param() {
1564 let params = parse_params("{{cmd}}");
1565 assert_eq!(params.len(), 1);
1566 assert_eq!(params[0].name, "cmd");
1567 }
1568
1569 #[test]
1570 fn test_parse_params_nested_braces_rejected() {
1571 let params = parse_params("{{{a}}}");
1573 assert!(params.is_empty());
1574 }
1575
1576 #[test]
1577 fn test_parse_params_colon_empty_default() {
1578 let params = parse_params("echo {{name:}}");
1579 assert_eq!(params.len(), 1);
1580 assert_eq!(params[0].name, "name");
1581 assert_eq!(params[0].default, Some("".to_string()));
1582 }
1583
1584 #[test]
1585 fn test_parse_params_empty_inner() {
1586 let params = parse_params("echo {{}}");
1587 assert!(params.is_empty());
1588 }
1589
1590 #[test]
1591 fn test_parse_params_single_braces_ignored() {
1592 let params = parse_params("echo {notaparam}");
1593 assert!(params.is_empty());
1594 }
1595
1596 #[test]
1597 fn test_parse_params_default_with_colons() {
1598 let params = parse_params("{{url:http://localhost:8080}}");
1599 assert_eq!(params.len(), 1);
1600 assert_eq!(params[0].name, "url");
1601 assert_eq!(params[0].default, Some("http://localhost:8080".to_string()));
1602 }
1603
1604 #[test]
1609 fn test_validate_param_name_unicode() {
1610 assert!(validate_param_name("caf\u{00e9}").is_ok());
1611 }
1612
1613 #[test]
1614 fn test_validate_param_name_hyphen_only() {
1615 assert!(validate_param_name("-").is_ok());
1616 }
1617
1618 #[test]
1619 fn test_validate_param_name_underscore_only() {
1620 assert!(validate_param_name("_").is_ok());
1621 }
1622
1623 #[test]
1624 fn test_validate_param_name_rejects_dot() {
1625 assert!(validate_param_name("a.b").is_err());
1626 }
1627
1628 #[test]
1633 fn test_substitute_no_params_passthrough() {
1634 let values = std::collections::HashMap::new();
1635 let result = substitute_params("df -h /tmp", &values);
1636 assert_eq!(result, "df -h /tmp");
1637 }
1638
1639 #[test]
1640 fn test_substitute_missing_param_no_default() {
1641 let values = std::collections::HashMap::new();
1642 let result = substitute_params("echo {{name}}", &values);
1643 assert_eq!(result, "echo ''");
1644 }
1645
1646 #[test]
1647 fn test_substitute_empty_value_falls_to_default() {
1648 let mut values = std::collections::HashMap::new();
1649 values.insert("name".to_string(), "".to_string());
1650 let result = substitute_params("echo {{name:fallback}}", &values);
1651 assert_eq!(result, "echo 'fallback'");
1652 }
1653
1654 #[test]
1655 fn test_substitute_non_ascii_around_params() {
1656 let mut values = std::collections::HashMap::new();
1657 values.insert("x".to_string(), "val".to_string());
1658 let result = substitute_params("\u{00e9}cho {{x}} \u{2603}", &values);
1659 assert_eq!(result, "\u{00e9}cho 'val' \u{2603}");
1660 }
1661
1662 #[test]
1663 fn test_substitute_adjacent_params() {
1664 let mut values = std::collections::HashMap::new();
1665 values.insert("a".to_string(), "x".to_string());
1666 values.insert("b".to_string(), "y".to_string());
1667 let result = substitute_params("{{a}}{{b}}", &values);
1668 assert_eq!(result, "'x''y'");
1669 }
1670
1671 #[test]
1676 fn test_sanitize_empty() {
1677 assert_eq!(sanitize_output(""), "");
1678 }
1679
1680 #[test]
1681 fn test_sanitize_only_escapes() {
1682 assert_eq!(sanitize_output("\x1b[31m\x1b[0m\x1b[1m"), "");
1683 }
1684
1685 #[test]
1686 fn test_sanitize_lone_esc_at_end() {
1687 assert_eq!(sanitize_output("hello\x1b"), "hello");
1688 }
1689
1690 #[test]
1691 fn test_sanitize_truncated_csi_no_terminator() {
1692 assert_eq!(sanitize_output("hello\x1b[123"), "hello");
1693 }
1694
1695 #[test]
1696 fn test_sanitize_apc_sequence() {
1697 assert_eq!(sanitize_output("\x1b_payload\x1b\\visible"), "visible");
1698 }
1699
1700 #[test]
1701 fn test_sanitize_pm_sequence() {
1702 assert_eq!(sanitize_output("\x1b^payload\x1b\\visible"), "visible");
1703 }
1704
1705 #[test]
1706 fn test_sanitize_dcs_terminated_by_bel() {
1707 assert_eq!(sanitize_output("\x1bPdata\x07text"), "text");
1708 }
1709
1710 #[test]
1711 fn test_sanitize_lone_esc_plus_letter() {
1712 assert_eq!(sanitize_output("a\x1bMb"), "ab");
1713 }
1714
1715 #[test]
1716 fn test_sanitize_multiple_mixed_sequences() {
1717 let input = "\x1b[1mbold\x1b[0m \x1b]0;title\x07normal \x01gone";
1719 assert_eq!(sanitize_output(input), "bold normal gone");
1720 }
1721}