rustutils_env/
lib.rs

1use clap::Parser;
2use rustutils_runnable::Runnable;
3use std::env;
4use std::error::Error;
5use std::ffi::{OsStr, OsString};
6use std::io::Write;
7use std::os::unix::ffi::OsStrExt;
8use std::process::{Command, ExitStatus};
9
10/// Set environment variables and run command.
11///
12/// When no command is supplied, print the current environment variables.
13#[derive(Parser, Clone, Debug)]
14#[clap(author, version, about, long_about = None)]
15pub struct Env {
16    /// Start with an empty environment.
17    #[clap(short, long)]
18    ignore_environment: bool,
19    /// When printing the current environment, separate the variables with a NUL character.
20    #[clap(short = '0', long)]
21    null: bool,
22    /// When printing the current environment, output it as JSON.
23    #[clap(short, long, conflicts_with = "null")]
24    json: bool,
25    /// Unset variables.
26    #[clap(short, long)]
27    unset: Vec<OsString>,
28    /// Environment variables to set, in the shape of `NAME=value`.
29    variables: Vec<OsString>,
30    /// Command to run, using the specified environment.
31    #[clap(last = true)]
32    command: Vec<OsString>,
33}
34
35/// Parse an environment variable definition.
36///
37/// Environment variables are specified in the form `VARIABLE=value`. Typically, the variables are
38/// uppercased, but this is not mandatory.
39pub fn parse_variable(var: &OsStr) -> Option<(&OsStr, &OsStr)> {
40    let var = var.as_bytes();
41    var.iter()
42        .enumerate()
43        .skip(1)
44        .filter(|(_index, byte)| **byte == b'=')
45        .map(|(index, _byte)| index)
46        .last()
47        .map(|pos| {
48            (
49                OsStr::from_bytes(&var[..pos]),
50                OsStr::from_bytes(&var[pos + 1..]),
51            )
52        })
53}
54
55#[test]
56fn can_parse_variable() {
57    assert_eq!(parse_variable(OsStr::from_bytes(b"")), None);
58    assert_eq!(parse_variable(OsStr::from_bytes(b"=")), None);
59    assert_eq!(parse_variable(OsStr::from_bytes(b"VAR")), None);
60    assert_eq!(parse_variable(OsStr::from_bytes(b"=VAR")), None);
61    assert_eq!(
62        parse_variable(OsStr::from_bytes(b"VAR=value")),
63        Some((OsStr::from_bytes(b"VAR"), OsStr::from_bytes(b"value")))
64    );
65    assert_eq!(
66        parse_variable(OsStr::from_bytes(b"VAR=")),
67        Some((OsStr::from_bytes(b"VAR"), OsStr::from_bytes(b"")))
68    );
69}
70
71impl Env {
72    /// Setup environment
73    pub fn setup(&self) {
74        // unset every environment variable
75        if self.ignore_environment {
76            for (name, _) in env::vars_os() {
77                env::remove_var(name);
78            }
79        }
80
81        // unset request environment variables
82        for variable in &self.unset {
83            env::remove_var(variable);
84        }
85    }
86
87    pub fn parse_arguments(&self) -> Result<usize, Box<dyn Error>> {
88        let count = self
89            .variables
90            .iter()
91            .map(|variable| parse_variable(variable))
92            .take_while(|parsed| parsed.is_some())
93            .filter_map(|parsed| parsed)
94            .inspect(|(name, value)| env::set_var(name, value))
95            .count();
96        Ok(count)
97    }
98
99    pub fn run_command(&self, offset: usize) -> Result<Option<ExitStatus>, Box<dyn Error>> {
100        let mut args = self
101            .variables
102            .iter()
103            .skip(offset)
104            .chain(self.command.iter());
105        if let Some(arg) = args.next() {
106            let mut command = Command::new(arg);
107            let result = command.args(args).spawn()?.wait()?;
108            Ok(Some(result))
109        } else {
110            Ok(None)
111        }
112    }
113
114    pub fn run(&self) -> Result<Option<ExitStatus>, Box<dyn Error>> {
115        self.setup();
116        let offset = self.parse_arguments()?;
117        let result = self.run_command(offset)?;
118        if result.is_none() {
119            self.show()?;
120        }
121        Ok(result)
122    }
123
124    pub fn show(&self) -> Result<(), Box<dyn Error>> {
125        if self.json {
126            self.show_json()
127        } else {
128            let separator = match self.null {
129                true => OsStr::new("\0"),
130                false => OsStr::new("\n"),
131            };
132            self.show_normal(separator)
133        }
134    }
135
136    pub fn show_json(&self) -> Result<(), Box<dyn Error>> {
137        let mut environment: std::collections::BTreeMap<String, String> = Default::default();
138        for (name, value) in env::vars_os() {
139            if let (Some(name), Some(value)) = (name.to_str(), value.to_str()) {
140                environment.insert(name.to_string(), value.to_string());
141            }
142        }
143        let mut stdout = std::io::stdout();
144        stdout.write_all(serde_json::to_string(&environment)?.as_bytes())?;
145        Ok(())
146    }
147    pub fn show_normal(&self, separator: &OsStr) -> Result<(), Box<dyn Error>> {
148        let mut stdout = std::io::stdout();
149        for (name, value) in env::vars_os() {
150            stdout.write_all(name.as_bytes())?;
151            stdout.write_all(b"=")?;
152            stdout.write_all(value.as_bytes())?;
153            stdout.write_all(separator.as_bytes())?;
154        }
155        Ok(())
156    }
157}
158
159impl Runnable for Env {
160    fn run(&self) -> Result<(), Box<dyn Error>> {
161        self.run()?;
162        Ok(())
163    }
164}
rustutils_env/lib.rs

rustutils_env/
lib.rs
