1extern crate snailquote;
27
28use std::collections::BTreeMap;
29use std::fs::File;
30use std::io::{self, Read, Write};
31use std::path::{Path, PathBuf};
32use std::str;
33
34use snailquote::{unescape, escape};
35
36
37pub struct EnvFile {
39 pub path: PathBuf,
41 pub store: BTreeMap<String, String>,
43}
44
45fn parse_line(entry: &[u8]) -> Option<(String, String)> {
46 str::from_utf8(entry).ok().and_then(|l| {
47 let line = l.trim();
48 if line.starts_with('#') {
50 return None;
51 }
52 let vline = line.as_bytes();
53 vline.iter().position(|&x| x == b'=').and_then(|pos| {
54 str::from_utf8(&vline[..pos]).ok().and_then(|x| {
55 str::from_utf8(&vline[pos+1..]).ok().and_then(|right| {
56 unescape(right).ok().map(|y| (x.to_owned(), y))
58 })
59 })
60 })
61 })
62}
63
64impl EnvFile {
65 pub fn new<P: Into<PathBuf>>(path: P) -> io::Result<Self> {
67 let path = path.into();
68 let data = read(&path)?;
69 let mut store = BTreeMap::new();
70
71 let values = data.split(|&x| x == b'\n').flat_map(parse_line);
72
73 for (key, value) in values {
74 store.insert(key, value);
75 }
76
77 Ok(EnvFile { path, store })
78 }
79
80 pub fn update(&mut self, key: &str, value: &str) -> &mut Self {
82 self.store.insert(key.into(), value.into());
83 self
84 }
85
86 pub fn get(&self, key: &str) -> Option<&str> {
88 self.store.get(key).as_ref().map(|x| x.as_str())
89 }
90
91 pub fn write(&mut self) -> io::Result<()> {
96 let mut buffer = Vec::with_capacity(1024);
97 for (key, value) in &self.store {
98 buffer.extend_from_slice(key.as_bytes());
99 buffer.push(b'=');
100 let v = escape(value.as_str()).into_owned();
102 buffer.extend_from_slice(v.as_bytes());
103 buffer.push(b'\n');
104 }
105
106 write(&self.path, &buffer)
107 }
108}
109
110fn open<P: AsRef<Path>>(path: P) -> io::Result<File> {
111 File::open(&path).map_err(|why| io::Error::new(
112 io::ErrorKind::Other,
113 format!("unable to open file at {:?}: {}", path.as_ref(), why)
114 ))
115}
116
117fn create<P: AsRef<Path>>(path: P) -> io::Result<File> {
118 File::create(&path).map_err(|why| io::Error::new(
119 io::ErrorKind::Other,
120 format!("unable to create file at {:?}: {}", path.as_ref(), why)
121 ))
122}
123
124fn read<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
125 open(path).and_then(|mut file| {
126 let mut buffer = Vec::with_capacity(file.metadata().ok().map_or(0, |x| x.len()) as usize);
127 file.read_to_end(&mut buffer).map(|_| buffer)
128 })
129}
130
131fn write<P: AsRef<Path>, C: AsRef<[u8]>>(path: P, contents: C) -> io::Result<()> {
132 create(path).and_then(|mut file| file.write_all(contents.as_ref()))
133}
134
135#[cfg(test)]
136mod tests {
137 extern crate tempdir;
138 use super::*;
139 use self::tempdir::TempDir;
140 use std::collections::BTreeMap;
141 use std::io::Write;
142
143 const SAMPLE: &str = r#"DOUBLE_QUOTED_STRING="This is a 'double-quoted' string"
144EFI_UUID=DFFD-D047
145HOSTNAME=pop-testing
146KBD_LAYOUT=us
147KBD_MODEL=
148KBD_VARIANT=
149 LANG=en_US.UTF-8
150OEM_MODE=0
151# Intentional blank line
152
153# Should ignore = operator in comment
154RECOVERY_UUID=PARTUUID=asdfasd7asdf7sad-asdfa
155ROOT_UUID=2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c
156SINGLE_QUOTED_STRING='This is a single-quoted string'
157"#;
158
159 const SAMPLE_CLEANED: &str = r#"DOUBLE_QUOTED_STRING="This is a 'double-quoted' string"
160EFI_UUID=DFFD-D047
161HOSTNAME=pop-testing
162KBD_LAYOUT=us
163KBD_MODEL=
164KBD_VARIANT=
165LANG=en_US.UTF-8
166OEM_MODE=0
167RECOVERY_UUID=PARTUUID=asdfasd7asdf7sad-asdfa
168ROOT_UUID=2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c
169SINGLE_QUOTED_STRING='This is a single-quoted string'
170"#;
171
172 #[test]
173 fn env_file_read() {
174 let tempdir = TempDir::new("distinst_test").unwrap();
175 let path = &tempdir.path().join("recovery.conf");
176
177 {
178 let mut file = create(path).unwrap();
179 file.write_all(SAMPLE.as_bytes()).unwrap();
180 }
181
182 let env = EnvFile::new(path).unwrap();
183 assert_eq!(&env.store, &{
184 let mut map = BTreeMap::new();
185 map.insert("HOSTNAME".into(), "pop-testing".into());
186 map.insert("LANG".into(), "en_US.UTF-8".into());
187 map.insert("KBD_LAYOUT".into(), "us".into());
188 map.insert("KBD_MODEL".into(), "".into());
189 map.insert("KBD_VARIANT".into(), "".into());
190 map.insert("EFI_UUID".into(), "DFFD-D047".into());
191 map.insert("RECOVERY_UUID".into(), "PARTUUID=asdfasd7asdf7sad-asdfa".into());
192 map.insert("ROOT_UUID".into(), "2ef950c2-5ce6-4ae0-9fb9-a8c7468fa82c".into());
193 map.insert("OEM_MODE".into(), "0".into());
194 map.insert("DOUBLE_QUOTED_STRING".into(), "This is a 'double-quoted' string".into());
195 map.insert("SINGLE_QUOTED_STRING".into(), "This is a single-quoted string".into());
196 map
197 });
198 }
199
200 #[test]
201 fn env_file_write() {
202 let tempdir = TempDir::new("distinst_test").unwrap();
203 let path = &tempdir.path().join("recovery.conf");
204
205 {
206 let mut file = create(path).unwrap();
207 file.write_all(SAMPLE.as_bytes()).unwrap();
208 }
209
210 let mut env = EnvFile::new(path).unwrap();
211 env.write().unwrap();
212 let copy: &[u8] = &read(path).unwrap();
213
214 assert_eq!(copy, SAMPLE_CLEANED.as_bytes(), "Expected '{}' == '{}'", String::from_utf8_lossy(copy), SAMPLE_CLEANED);
215 }
216}