pacaptr/
exec.rs

1//! APIs for spawning subprocesses and handling their results.
2
3use 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/// Different ways in which a [`Cmd`] shall be dealt with.
34#[derive(Copy, Clone, Default, Debug)]
35pub enum Mode {
36    /// Solely prints out the command that should be executed and stops.
37    PrintCmd,
38
39    /// Silently collects all the `stdout`/`stderr` combined. Prints nothing.
40    Mute,
41
42    /// Prints out the command which should be executed, runs it and collects
43    /// its `stdout`/`stderr` combined.
44    ///
45    /// This is potentially dangerous as it destroys the colored `stdout`. Use
46    /// it only if really necessary.
47    CheckAll {
48        /// Whether the log output should be suppressed.
49        quiet: bool,
50    },
51
52    /// Prints out the command which should be executed, runs it and collects
53    /// its `stderr`.
54    ///
55    /// This will work with a colored `stdout`.
56    CheckErr {
57        /// Whether the log output should be suppressed.
58        quiet: bool,
59    },
60
61    /// A CUSTOM prompt implemented by a `pacaptr` module itself.
62    ///
63    /// Prints out the command which should be executed, runs it and collects
64    /// its `stderr`. Also, this will ask for confirmation before
65    /// proceeding.
66    #[default]
67    Prompt,
68}
69
70/// The status code type returned by a [`Cmd`],
71pub type StatusCode = i32;
72
73/// Returns a [`Result`] for a [`Cmd`] according to if its exit status code
74/// indicates an error.
75///
76/// # Errors
77/// This function might return one of the following errors:
78///
79/// - [`Error::CmdStatusCodeError`], when `status` is `Some(n)` where `n != 0`.
80/// - [`Error::CmdInterruptedError`], when `status` is `None`.
81#[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
90/// The type for captured `stdout`, and if set to [`Mode::CheckAll`], mixed with
91/// captured `stderr`.
92pub type Output = Vec<u8>;
93
94/// A command to be executed, provided in `command-flags-keywords` form.
95#[must_use]
96#[derive(Debug, Clone, Default)]
97pub struct Cmd {
98    /// Flag indicating if a **normal admin** needs to run this command with
99    /// `sudo`.
100    pub sudo: bool,
101
102    /// The "command" part of the command string, e.g. `brew install`.
103    pub cmd: Vec<String>,
104
105    /// The "flags" part of the command string, e.g. `--dry-run`.
106    pub flags: Vec<String>,
107
108    /// The "keywords" part of the command string, e.g. `curl fish`.
109    pub kws: Vec<String>,
110}
111
112impl Cmd {
113    /// Makes a new [`Cmd`] instance with the given [`cmd`](Cmd::cmd) part.
114    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    /// Makes a new [`Cmd`] instance with the given [`cmd`](Cmd::cmd) part,
122    /// setting [`sudo`](field@Cmd::sudo) to `true`.
123    pub(crate) fn with_sudo(cmd: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
124        Self::new(cmd).sudo(true)
125    }
126
127    /// Overrides the value of [`flags`](field@Cmd::flags).
128    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    /// Overrides the value of [`kws`](field@Cmd::kws).
134    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    /// Overrides the value of [`sudo`](field@Cmd::sudo).
140    pub(crate) const fn sudo(mut self, sudo: bool) -> Self {
141        self.sudo = sudo;
142        self
143    }
144
145    /// Determines if this command actually needs to run with `sudo -S`.
146    ///
147    /// If a **normal admin** needs to run it with `sudo`, and we are not
148    /// `root`, then this is the case.
149    #[must_use]
150    fn should_sudo(&self) -> bool {
151        self.sudo && !is_root()
152    }
153
154    /// Converts a [`Cmd`] object into an [`Exec`].
155    #[must_use]
156    fn build(self) -> Exec {
157        // ! Special fix for `zypper`: `zypper install -y curl` is accepted,
158        // ! but not `zypper install curl -y`.
159        // ! So we place the flags first, and then keywords.
160        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
180/// Takes contents from an input stream and copy to an output stream (optional)
181/// and a [`Vec<u8>`], then returns the [`Vec<u8>`].
182///
183/// Helper to implement [`Cmd::exec_checkerr`] and [`Cmd::exec_checkall`].
184///
185/// # Arguments
186///
187/// * `src` - The input stream to read from.
188/// * `out` - The optional output stream to write to.
189async 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    /// Executes a [`Cmd`] and returns its output.
225    ///
226    /// The exact behavior depends on the [`Mode`] passed in (see the definition
227    /// of [`Mode`] for more info).
228    #[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    /// Inner implementation of [`Cmd::exec_checkerr`] (if `merge` is `false`)
253    /// and [`Cmd::exec_checkall`] (otherwise).
254    #[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    /// Executes a [`Cmd`] and returns its `stdout` and `stderr`.
304    ///
305    /// If `mute` is `false`, then normal `stdout/stderr` output will be printed
306    /// to `stdout` too.
307    #[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    /// Executes a [`Cmd`] and collects its `stderr`.
313    ///
314    /// If `mute` is `false`, then its `stderr` output will be printed to
315    /// `stderr` too.
316    #[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    /// Executes a [`Cmd`] and collects its `stderr`.
322    ///
323    /// If `mute` is `false`, then its `stderr` output will be printed to
324    /// `stderr` too.
325    ///
326    /// This function behaves just like [`exec_checkerr`](Cmd::exec_checkerr),
327    /// but in addition, the user will be prompted if (s)he wishes to
328    /// continue with the command execution.
329    #[doc = docs_errors_exec!()]
330    async fn exec_prompt(self, mute: bool) -> Result<Output> {
331        /// If the user has skipped all the prompts with `yes`.
332        static ALL: AtomicBool = AtomicBool::new(false);
333
334        // The answer obtained from the prompt.
335        // The only Atomic* we're dealing with is `ALL`, so `Ordering::Relaxed` is fine.
336        // See: <https://marabos.nl/atomics/memory-ordering.html#relaxed>
337        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                // The default answer is `Yes`.
348                0 => true,
349                // You can also say `All` to answer `Yes` to all the other questions that follow.
350                1 => {
351                    ALL.store(true, Ordering::Relaxed);
352                    true
353                }
354                // Or you can say `No`.
355                2 => false,
356                // ! I didn't put a `None` option because you can just Ctrl-C it if you want.
357                _ => 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
376/// Gives a prompt and returns the index of the user choice.
377fn 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/// Finds all lines in the given `text` that matches all the `patterns`.
396///
397/// We suppose that all patterns are legal regular expressions.
398/// An error message will be returned if this is not the case.
399#[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/// Prints the result of [`grep`] line by line.
412#[doc = docs_errors_grep!()]
413pub fn grep_print(text: &str, patterns: &[&str]) -> Result<()> {
414    grep_print_with_header(text, patterns, 0)
415}
416
417/// Prints the result of [`grep`] line by line, with `header_lines` of header
418/// prepended.
419/// If `header_lines >= text.lines().count()`, then `text` is printed without
420/// changes.
421#[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/// Checks if an executable exists by name (consult `$PATH`) or by path.
434///
435/// To check by one parameter only, pass `""` to the other one.
436#[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/// Checks if the current user is root or admin.
442#[cfg(windows)]
443#[must_use]
444pub fn is_root() -> bool {
445    is_elevated::is_elevated()
446}
447
448/// Checks if the current user is root or admin.
449#[cfg(unix)]
450#[must_use]
451pub fn is_root() -> bool {
452    nix::unistd::Uid::current().is_root()
453}
454
455/// Turns an [`AsyncRead`] into a [`Stream`].
456///
457/// _Shamelessly copied from [`StackOverflow`](https://stackoverflow.com/a/59327560)._
458fn into_bytes(reader: impl AsyncRead) -> impl Stream<Item = io::Result<Bytes>> {
459    FramedRead::new(reader, BytesCodec::new()).map_ok(BytesMut::freeze)
460}