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 log::warn!("[config] 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 struct ChildGuard {
367 inner: std::sync::Mutex<Option<std::process::Child>>,
368 pgid: i32,
369}
370
371impl ChildGuard {
372 fn new(child: std::process::Child) -> Self {
373 let pgid = i32::try_from(child.id()).unwrap_or(-1);
377 Self {
378 inner: std::sync::Mutex::new(Some(child)),
379 pgid,
380 }
381 }
382}
383
384impl Drop for ChildGuard {
385 fn drop(&mut self) {
386 let mut lock = self.inner.lock().unwrap_or_else(|e| e.into_inner());
387 if let Some(ref mut child) = *lock {
388 if let Ok(Some(_)) = child.try_wait() {
390 return;
391 }
392 #[cfg(unix)]
398 unsafe {
399 libc::kill(-self.pgid, libc::SIGTERM);
400 }
401 let deadline = std::time::Instant::now() + std::time::Duration::from_millis(500);
403 loop {
404 if let Ok(Some(_)) = child.try_wait() {
405 return;
406 }
407 if std::time::Instant::now() >= deadline {
408 break;
409 }
410 std::thread::sleep(std::time::Duration::from_millis(50));
411 }
412 #[cfg(unix)]
414 unsafe {
415 libc::kill(-self.pgid, libc::SIGKILL);
416 }
417 let _ = child.kill();
419 let _ = child.wait();
420 }
421 }
422}
423
424fn read_pipe_capped<R: io::Read>(reader: R) -> String {
427 use io::BufRead;
428 let mut reader = io::BufReader::new(reader);
429 let mut output = String::new();
430 let mut line_count = 0;
431 let mut capped = false;
432 let mut buf = Vec::new();
433 loop {
434 buf.clear();
435 match reader.read_until(b'\n', &mut buf) {
436 Ok(0) => break, Ok(_) => {
438 if !capped {
439 if line_count < MAX_OUTPUT_LINES {
440 if line_count > 0 {
441 output.push('\n');
442 }
443 if buf.last() == Some(&b'\n') {
445 buf.pop();
446 if buf.last() == Some(&b'\r') {
447 buf.pop();
448 }
449 }
450 output.push_str(&String::from_utf8_lossy(&buf));
452 line_count += 1;
453 } else {
454 output.push_str("\n[Output truncated at 10,000 lines]");
455 capped = true;
456 }
457 }
458 }
460 Err(_) => break,
461 }
462 }
463 output
464}
465
466fn base_ssh_command(
470 alias: &str,
471 config_path: &Path,
472 command: &str,
473 askpass: Option<&str>,
474 bw_session: Option<&str>,
475 has_active_tunnel: bool,
476) -> Command {
477 let mut cmd = Command::new("ssh");
478 cmd.arg("-F")
479 .arg(config_path)
480 .arg("-o")
481 .arg("ConnectTimeout=10")
482 .arg("-o")
483 .arg("ControlMaster=no")
484 .arg("-o")
485 .arg("ControlPath=none");
486
487 if has_active_tunnel {
488 cmd.arg("-o").arg("ClearAllForwardings=yes");
489 }
490
491 cmd.arg("--").arg(alias).arg(command);
492
493 if askpass.is_some() {
494 crate::askpass_env::configure_ssh_command(&mut cmd, alias, config_path);
495 }
496
497 if let Some(token) = bw_session {
498 cmd.env("BW_SESSION", token);
499 }
500
501 cmd
502}
503
504fn build_snippet_command(
506 alias: &str,
507 config_path: &Path,
508 command: &str,
509 askpass: Option<&str>,
510 bw_session: Option<&str>,
511 has_active_tunnel: bool,
512) -> Command {
513 let mut cmd = base_ssh_command(
514 alias,
515 config_path,
516 command,
517 askpass,
518 bw_session,
519 has_active_tunnel,
520 );
521 cmd.stdin(Stdio::null())
522 .stdout(Stdio::piped())
523 .stderr(Stdio::piped());
524
525 #[cfg(unix)]
528 unsafe {
529 use std::os::unix::process::CommandExt;
530 cmd.pre_exec(|| {
531 libc::setpgid(0, 0);
532 Ok(())
533 });
534 }
535
536 cmd
537}
538
539fn execute_host(
541 run_id: u64,
542 ctx: &crate::ssh_context::SshContext<'_>,
543 command: &str,
544 tx: &std::sync::mpsc::Sender<crate::event::AppEvent>,
545) -> Option<std::sync::Arc<ChildGuard>> {
546 let alias = ctx.alias;
547 let mut cmd = build_snippet_command(
548 alias,
549 ctx.config_path,
550 command,
551 ctx.askpass,
552 ctx.bw_session,
553 ctx.has_tunnel,
554 );
555
556 match cmd.spawn() {
557 Ok(child) => {
558 let guard = std::sync::Arc::new(ChildGuard::new(child));
559
560 let stdout_pipe = {
562 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
563 lock.as_mut().and_then(|c| c.stdout.take())
564 };
565 let stderr_pipe = {
566 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
567 lock.as_mut().and_then(|c| c.stderr.take())
568 };
569
570 let stdout_handle = std::thread::spawn(move || match stdout_pipe {
572 Some(pipe) => read_pipe_capped(pipe),
573 None => String::new(),
574 });
575 let stderr_handle = std::thread::spawn(move || match stderr_pipe {
576 Some(pipe) => read_pipe_capped(pipe),
577 None => String::new(),
578 });
579
580 let stdout_text = stdout_handle.join().unwrap_or_else(|_| {
582 log::warn!("[purple] Snippet stdout reader thread panicked");
583 String::new()
584 });
585 let stderr_text = stderr_handle.join().unwrap_or_else(|_| {
586 log::warn!("[purple] Snippet stderr reader thread panicked");
587 String::new()
588 });
589
590 let exit_code = {
593 let mut lock = guard.inner.lock().unwrap_or_else(|e| e.into_inner());
594 let status = lock.as_mut().and_then(|c| c.wait().ok());
595 let _ = lock.take(); status.and_then(|s| {
597 #[cfg(unix)]
598 {
599 use std::os::unix::process::ExitStatusExt;
600 s.code().or_else(|| s.signal().map(|sig| 128 + sig))
601 }
602 #[cfg(not(unix))]
603 {
604 s.code()
605 }
606 })
607 };
608
609 let _ = tx.send(crate::event::AppEvent::SnippetHostDone {
610 run_id,
611 alias: alias.to_string(),
612 stdout: sanitize_output(&stdout_text),
613 stderr: sanitize_output(&stderr_text),
614 exit_code,
615 });
616
617 Some(guard)
618 }
619 Err(e) => {
620 let _ = tx.send(crate::event::AppEvent::SnippetHostDone {
621 run_id,
622 alias: alias.to_string(),
623 stdout: String::new(),
624 stderr: format!("Failed to launch ssh: {}", e),
625 exit_code: None,
626 });
627 None
628 }
629 }
630}
631
632#[allow(clippy::too_many_arguments)]
635pub fn spawn_snippet_execution(
636 run_id: u64,
637 askpass_map: Vec<(String, Option<String>)>,
638 config_path: PathBuf,
639 command: String,
640 bw_session: Option<String>,
641 tunnel_aliases: std::collections::HashSet<String>,
642 cancel: std::sync::Arc<std::sync::atomic::AtomicBool>,
643 tx: std::sync::mpsc::Sender<crate::event::AppEvent>,
644 parallel: bool,
645) {
646 let total = askpass_map.len();
647 let max_concurrent: usize = 20;
648
649 std::thread::Builder::new()
650 .name("snippet-coordinator".into())
651 .spawn(move || {
652 let guards: std::sync::Arc<std::sync::Mutex<Vec<std::sync::Arc<ChildGuard>>>> =
653 std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
654
655 if parallel && total > 1 {
656 let (slot_tx, slot_rx) = std::sync::mpsc::channel::<()>();
658 for _ in 0..max_concurrent.min(total) {
659 let _ = slot_tx.send(());
660 }
661
662 let completed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
663 let mut worker_handles = Vec::new();
664
665 for (alias, askpass) in askpass_map {
666 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
667 break;
668 }
669
670 loop {
672 match slot_rx.recv_timeout(std::time::Duration::from_millis(100)) {
673 Ok(()) => break,
674 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
675 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
676 break;
677 }
678 }
679 Err(_) => break, }
681 }
682
683 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
684 break;
685 }
686
687 let config_path = config_path.clone();
688 let command = command.clone();
689 let bw_session = bw_session.clone();
690 let has_tunnel = tunnel_aliases.contains(&alias);
691 let tx = tx.clone();
692 let slot_tx = slot_tx.clone();
693 let guards = guards.clone();
694 let completed = completed.clone();
695 let total = total;
696
697 let handle = std::thread::spawn(move || {
698 struct SlotRelease(Option<std::sync::mpsc::Sender<()>>);
700 impl Drop for SlotRelease {
701 fn drop(&mut self) {
702 if let Some(tx) = self.0.take() {
703 let _ = tx.send(());
704 }
705 }
706 }
707 let _slot = SlotRelease(Some(slot_tx));
708
709 let host_ctx = crate::ssh_context::SshContext {
710 alias: &alias,
711 config_path: &config_path,
712 askpass: askpass.as_deref(),
713 bw_session: bw_session.as_deref(),
714 has_tunnel,
715 };
716 let guard = execute_host(run_id, &host_ctx, &command, &tx);
717
718 if let Some(g) = guard {
720 guards.lock().unwrap_or_else(|e| e.into_inner()).push(g);
721 }
722
723 let c = completed.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
724 let _ = tx.send(crate::event::AppEvent::SnippetProgress {
725 run_id,
726 completed: c,
727 total,
728 });
729 });
731 worker_handles.push(handle);
732 }
733
734 for handle in worker_handles {
736 let _ = handle.join();
737 }
738 } else {
739 for (i, (alias, askpass)) in askpass_map.into_iter().enumerate() {
741 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
742 break;
743 }
744
745 let has_tunnel = tunnel_aliases.contains(&alias);
746 let host_ctx = crate::ssh_context::SshContext {
747 alias: &alias,
748 config_path: &config_path,
749 askpass: askpass.as_deref(),
750 bw_session: bw_session.as_deref(),
751 has_tunnel,
752 };
753 let guard = execute_host(run_id, &host_ctx, &command, &tx);
754
755 if let Some(g) = guard {
756 guards.lock().unwrap_or_else(|e| e.into_inner()).push(g);
757 }
758
759 let _ = tx.send(crate::event::AppEvent::SnippetProgress {
760 run_id,
761 completed: i + 1,
762 total,
763 });
764 }
765 }
766
767 let _ = tx.send(crate::event::AppEvent::SnippetAllDone { run_id });
768 })
770 .expect("failed to spawn snippet coordinator");
771}
772
773pub fn run_snippet(
778 alias: &str,
779 config_path: &Path,
780 command: &str,
781 askpass: Option<&str>,
782 bw_session: Option<&str>,
783 capture: bool,
784 has_active_tunnel: bool,
785) -> anyhow::Result<SnippetResult> {
786 let mut cmd = base_ssh_command(
787 alias,
788 config_path,
789 command,
790 askpass,
791 bw_session,
792 has_active_tunnel,
793 );
794 cmd.stdin(Stdio::inherit());
795
796 if capture {
797 cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
798 } else {
799 cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
800 }
801
802 if capture {
803 let output = cmd
804 .output()
805 .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
806
807 Ok(SnippetResult {
808 status: output.status,
809 stdout: String::from_utf8_lossy(&output.stdout).to_string(),
810 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
811 })
812 } else {
813 let status = cmd
814 .status()
815 .map_err(|e| anyhow::anyhow!("Failed to run ssh for '{}': {}", alias, e))?;
816
817 Ok(SnippetResult {
818 status,
819 stdout: String::new(),
820 stderr: String::new(),
821 })
822 }
823}
824
825#[cfg(test)]
826#[path = "snippet_tests.rs"]
827mod tests;