1use std::{
4 process::Stdio,
5 sync::atomic::{AtomicBool, Ordering},
6};
7
8use bytes::{Bytes, BytesMut};
9use dialoguer::FuzzySelect;
10use futures::prelude::*;
11use indoc::indoc;
12use itertools::{Itertools, chain};
13use regex::{RegexSet, RegexSetBuilder};
14use tap::prelude::*;
15use tokio::{
16 io::{self, AsyncRead, AsyncWrite},
17 process::Command as Exec,
18 task::JoinHandle,
19};
20#[allow(clippy::wildcard_imports)]
21use tokio_util::{
22 codec::{BytesCodec, FramedRead},
23 compat::*,
24 either::Either,
25};
26use which::which;
27
28use crate::{
29 error::{Error, Result},
30 print::{println_quoted, prompt, question_theme},
31};
32
33#[derive(Copy, Clone, Default, Debug)]
35pub enum Mode {
36 PrintCmd,
38
39 Mute,
41
42 CheckAll {
48 quiet: bool,
50 },
51
52 CheckErr {
57 quiet: bool,
59 },
60
61 #[default]
67 Prompt,
68}
69
70pub type StatusCode = i32;
72
73#[allow(clippy::missing_const_for_fn)]
82fn exit_result(code: Option<StatusCode>, output: Output) -> Result<Output> {
83 match code {
84 Some(0) => Ok(output),
85 Some(code) => Err(Error::CmdStatusCodeError { code, output }),
86 None => Err(Error::CmdInterruptedError),
87 }
88}
89
90pub type Output = Vec<u8>;
93
94#[must_use]
96#[derive(Debug, Clone, Default)]
97pub struct Cmd {
98 pub sudo: bool,
101
102 pub cmd: Vec<String>,
104
105 pub flags: Vec<String>,
107
108 pub kws: Vec<String>,
110}
111
112impl Cmd {
113 pub(crate) fn new(cmd: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
115 Self {
116 cmd: cmd.into_iter().map(|s| s.as_ref().into()).collect(),
117 ..Self::default()
118 }
119 }
120
121 pub(crate) fn with_sudo(cmd: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
124 Self::new(cmd).sudo(true)
125 }
126
127 pub(crate) fn flags(mut self, flags: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
129 self.flags = flags.into_iter().map(|s| s.as_ref().into()).collect();
130 self
131 }
132
133 pub(crate) fn kws(mut self, kws: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
135 self.kws = kws.into_iter().map(|s| s.as_ref().into()).collect();
136 self
137 }
138
139 pub(crate) const fn sudo(mut self, sudo: bool) -> Self {
141 self.sudo = sudo;
142 self
143 }
144
145 #[must_use]
150 fn should_sudo(&self) -> bool {
151 self.sudo && !is_root()
152 }
153
154 #[must_use]
156 fn build(self) -> Exec {
157 if self.should_sudo() {
161 Exec::new("sudo").tap_mut(|builder| {
162 builder
163 .arg("-S")
164 .args(&self.cmd)
165 .args(&self.flags)
166 .args(&self.kws);
167 })
168 } else {
169 let (cmd, subcmd) = self
170 .cmd
171 .split_first()
172 .expect("failed to build Cmd, command is empty");
173 Exec::new(cmd).tap_mut(|builder| {
174 builder.args(subcmd).args(&self.flags).args(&self.kws);
175 })
176 }
177 }
178}
179
180async fn exec_tee(
190 src: impl Stream<Item = io::Result<Bytes>> + Send,
191 out: Option<impl AsyncWrite + Send>,
192) -> Result<Vec<u8>> {
193 let mut buf = Vec::<u8>::new();
194 let buf_sink = (&mut buf).into_sink();
195
196 let sink = if let Some(out) = out {
197 let out_sink = out.compat_write().into_sink();
198 buf_sink.fanout(out_sink).left_sink()
199 } else {
200 buf_sink.right_sink()
201 };
202
203 src.forward(sink).await?;
204 Ok(buf)
205}
206
207macro_rules! docs_errors_exec {
208 () => {
209 indoc! {"
210 # Errors
211 This function might return one of the following errors:
212
213 - [`Error::CmdJoinError`]
214 - [`Error::CmdNoHandleError`]
215 - [`Error::CmdSpawnError`]
216 - [`Error::CmdWaitError`]
217 - [`Error::CmdStatusCodeError`]
218 - [`Error::CmdInterruptedError`]
219 "}
220 };
221}
222
223impl Cmd {
224 #[doc = docs_errors_exec!()]
229 pub(crate) async fn exec(self, mode: Mode) -> Result<Output> {
230 match mode {
231 Mode::PrintCmd => {
232 println_quoted(&*prompt::CANCELED, &self);
233 Ok(Output::default())
234 }
235 Mode::Mute => self.exec_checkall(true).await,
236 Mode::CheckAll { quiet } => {
237 if !quiet {
238 println_quoted(&*prompt::RUNNING, &self);
239 }
240 self.exec_checkall(false).await
241 }
242 Mode::CheckErr { quiet } => {
243 if !quiet {
244 println_quoted(&*prompt::RUNNING, &self);
245 }
246 self.exec_checkerr(false).await
247 }
248 Mode::Prompt => self.exec_prompt(false).await,
249 }
250 }
251
252 #[doc = docs_errors_exec!()]
255 async fn exec_check_output(self, mute: bool, merge: bool) -> Result<Output> {
256 use Error::{CmdJoinError, CmdNoHandleError, CmdSpawnError, CmdWaitError};
257 use tokio_stream::StreamExt;
258
259 fn make_reader(
260 src: Option<impl AsyncRead>,
261 name: &str,
262 ) -> Result<impl Stream<Item = io::Result<Bytes>>> {
263 src.map(into_bytes).ok_or_else(|| CmdNoHandleError {
264 handle: name.into(),
265 })
266 }
267
268 let mut child = self
269 .build()
270 .stderr(Stdio::piped())
271 .tap_deref_mut(|cmd| {
272 if merge {
273 cmd.stdout(Stdio::piped());
274 }
275 })
276 .spawn()
277 .map_err(CmdSpawnError)?;
278
279 let stderr_reader = make_reader(child.stderr.take(), "stderr")?;
280 let mut reader = if merge {
281 let stdout_reader = make_reader(child.stdout.take(), "stdout")?;
282 StreamExt::merge(stdout_reader, stderr_reader).left_stream()
283 } else {
284 stderr_reader.right_stream()
285 };
286
287 let mut out = if merge {
288 Either::Left(io::stdout())
289 } else {
290 Either::Right(io::stderr())
291 };
292
293 let code: JoinHandle<Result<Option<i32>>> = tokio::spawn(async move {
294 let status = child.wait().await.map_err(CmdWaitError)?;
295 Ok(status.code())
296 });
297
298 let output = exec_tee(&mut reader, (!mute).then_some(&mut out)).await?;
299 let code = code.await.map_err(CmdJoinError)??;
300 exit_result(code, output)
301 }
302
303 #[doc = docs_errors_exec!()]
308 async fn exec_checkall(self, mute: bool) -> Result<Output> {
309 self.exec_check_output(mute, true).await
310 }
311
312 #[doc = docs_errors_exec!()]
317 async fn exec_checkerr(self, mute: bool) -> Result<Output> {
318 self.exec_check_output(mute, false).await
319 }
320
321 #[doc = docs_errors_exec!()]
330 async fn exec_prompt(self, mute: bool) -> Result<Output> {
331 static ALL: AtomicBool = AtomicBool::new(false);
333
334 let proceed = ALL.load(Ordering::Relaxed) || {
338 println_quoted(&*prompt::PENDING, &self);
339 let answer = tokio::task::block_in_place(move || {
340 prompt(
341 "Proceed",
342 "with the previous command?",
343 &["Yes", "All", "No"],
344 )
345 })?;
346 match answer {
347 0 => true,
349 1 => {
351 ALL.store(true, Ordering::Relaxed);
352 true
353 }
354 2 => false,
356 _ => unreachable!(),
358 }
359 };
360 if !proceed {
361 return Ok(Output::default());
362 }
363 println_quoted(&*prompt::RUNNING, &self);
364 self.exec_checkerr(mute).await
365 }
366}
367
368impl std::fmt::Display for Cmd {
369 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
370 let sudo: &str = if self.should_sudo() { "sudo -S " } else { "" };
371 let cmd = chain!(&self.cmd, &self.flags, &self.kws).join(" ");
372 write!(f, "{sudo}{cmd}")
373 }
374}
375
376fn prompt(prompt: &str, question: &str, expected: &[&str]) -> Result<usize> {
378 Ok(FuzzySelect::with_theme(&question_theme(prompt))
379 .with_prompt(question)
380 .items(expected)
381 .default(0)
382 .interact()?)
383}
384
385macro_rules! docs_errors_grep {
386 () => {
387 indoc! {"
388 # Errors
389 Returns an [`Error::OtherError`] when any of the
390 regex patterns is ill-formed.
391 "}
392 };
393}
394
395#[doc = docs_errors_grep!()]
400pub fn grep<'t>(text: &'t str, patterns: &[&str]) -> Result<Vec<&'t str>> {
401 let patterns: RegexSet = RegexSetBuilder::new(patterns)
402 .case_insensitive(true)
403 .build()
404 .map_err(|e| Error::OtherError(format!("ill-formed patterns found: {e:?}")))?;
405 Ok(text
406 .lines()
407 .filter(|line| patterns.matches(line).into_iter().count() == patterns.len())
408 .collect())
409}
410
411#[doc = docs_errors_grep!()]
413pub fn grep_print(text: &str, patterns: &[&str]) -> Result<()> {
414 grep_print_with_header(text, patterns, 0)
415}
416
417#[doc = docs_errors_grep!()]
422pub fn grep_print_with_header(text: &str, patterns: &[&str], header_lines: usize) -> Result<()> {
423 let lns = text.lines().collect_vec();
424 let (header, rest) = lns.split_at(header_lines);
425 header
426 .iter()
427 .copied()
428 .chain(grep(&rest.join("\n"), patterns)?)
429 .for_each(|ln| println!("{ln}"));
430 Ok(())
431}
432
433#[must_use]
437pub fn is_exe(name: &str, path: &str) -> bool {
438 (!path.is_empty() && which(path).is_ok()) || (!name.is_empty() && which(name).is_ok())
439}
440
441#[cfg(windows)]
443#[must_use]
444pub fn is_root() -> bool {
445 is_elevated::is_elevated()
446}
447
448#[cfg(unix)]
450#[must_use]
451pub fn is_root() -> bool {
452 nix::unistd::Uid::current().is_root()
453}
454
455fn into_bytes(reader: impl AsyncRead) -> impl Stream<Item = io::Result<Bytes>> {
459 FramedRead::new(reader, BytesCodec::new()).map_ok(BytesMut::freeze)
460}