1use std::{
2 collections::HashMap,
3 env, fs,
4 io::{self},
5 path::{Path, PathBuf},
6};
7
8use serde::Serialize;
9use toml::Table;
10
11use crate::{
12 config::Config,
13 profile::Profile,
14 prompt_store,
15 utils::{LogLevel, cprintln},
16};
17
18#[cfg(test)]
19mod tests;
20
21const BITWARDEN_NOTE_OVERRIDE_KEY: &str = "DOTR_BITWARDEN_NOTE";
25
26const DOTR_PROFILE_KEY: &str = "DOTR_PROFILE";
27
28#[derive(Debug, Clone, Serialize)]
29pub struct Context {
30 pub working_dir: PathBuf,
31 variables: Table,
32 user_variables: Table,
33 pub profile: Profile,
34 pub symlink: bool,
35}
36
37impl Context {
38 pub fn get_variable(&self, key: &str) -> Option<&toml::Value> {
39 self.variables.get(key)
40 }
41
42 pub fn get_user_variable(&self, key: &str) -> Option<&toml::Value> {
43 self.user_variables.get(key)
44 }
45
46 pub fn get_profile_variable(&self, key: &str) -> Option<&toml::Value> {
47 self.profile.variables.get(key)
48 }
49
50 pub fn get_context_variable(&self, key: &str) -> Option<&toml::Value> {
51 self.get_user_variable(key).or_else(|| {
52 self.get_profile_variable(key)
53 .or_else(|| self.get_variable(key))
54 })
55 }
56
57 pub fn set_profile(&mut self, profile: Profile) {
58 self.profile = profile;
59 }
60
61 pub fn get_prompted_variables(
62 &mut self,
63 conf: &Config,
64 packages: &Option<Vec<String>>,
65 ) -> anyhow::Result<Table> {
66 self.get_prompted_variables_with_io(
67 conf,
68 packages,
69 &mut std::io::stdin().lock(),
70 &mut std::io::stdout(),
71 )
72 }
73
74 fn get_backend(
78 profile: &Profile,
79 conf: &Config,
80 variables: &Table,
81 ) -> Box<dyn prompt_store::PromptStoreBackend> {
82 let backend_type = profile
83 .prompt_backend
84 .or(conf.prompt_backend)
85 .unwrap_or_default();
86 match backend_type {
87 prompt_store::PromptBackendType::File => Box::new(prompt_store::FileStore::new()),
88 prompt_store::PromptBackendType::Keychain => {
89 Box::new(prompt_store::KeychainStore::new())
90 }
91 prompt_store::PromptBackendType::Bitwarden => {
92 let note = variables
93 .get(BITWARDEN_NOTE_OVERRIDE_KEY)
94 .and_then(|v| v.as_str())
95 .unwrap_or(prompt_store::DEFAULT_BITWARDEN_NOTE)
96 .to_string();
97 Box::new(prompt_store::BitwardenStore::new(note))
98 }
99 }
100 }
101
102 pub fn get_prompts(
103 profile: &Profile,
104 conf: &Config,
105 packages: &Option<Vec<String>>,
106 ) -> HashMap<String, String> {
107 let mut prompts = conf.prompts.clone();
109 for (key, prompt) in profile.prompts.iter() {
110 prompts.insert(key.clone(), prompt.clone());
111 }
112 if let Ok(filtered_packages) = conf.filter_packages(profile, packages, false) {
113 for package in filtered_packages.values() {
114 for (key, prompt) in package.prompts.iter() {
115 prompts.insert(key.clone(), prompt.clone());
116 }
117 }
118 }
119 prompts
120 }
121
122 fn get_prompted_variables_with_io<R: io::BufRead, W: io::Write>(
123 &mut self,
124 conf: &Config,
125 packages: &Option<Vec<String>>,
126 reader: &mut R,
127 writer: &mut W,
128 ) -> Result<Table, anyhow::Error> {
129 let prompts = Self::get_prompts(&self.profile, conf, packages);
131 let missing_keys: Vec<String> = prompts
132 .keys()
133 .filter(|key| !self.user_variables.contains_key(*key))
134 .cloned()
135 .collect();
136
137 let mut store = self.user_variables.clone();
138 if missing_keys.is_empty() {
139 return Ok(store);
140 }
141
142 let mut backend = Self::get_backend(&self.profile, conf, &self.variables);
148 let existing = backend.get(&self.working_dir, &missing_keys)?;
149 let mut dirty = false;
150
151 for key in &missing_keys {
152 if let Some(value) = existing.get(key) {
153 store.insert(key.clone(), value.clone());
154 continue;
155 }
156 let message = &prompts[key];
157 match get_prompted_variables(message, &mut *reader, &mut *writer) {
158 Ok(input) => {
159 store.insert(key.clone(), toml::Value::String(input));
160 dirty = true;
161 }
162 Err(e) => {
163 cprintln(
164 &format!("Error getting prompted variable '{}': {}", key, e),
165 &LogLevel::Warning,
166 );
167 }
168 }
169 }
170
171 if dirty {
172 backend.save(&self.working_dir, &store)?;
173 }
174 self.user_variables = store.clone();
175 Ok(store)
176 }
177
178 pub fn save_to_uservariables(
179 &mut self,
180 key: &str,
181 val: toml::Value,
182 ) -> Result<(), anyhow::Error> {
183 let mut on_disk = Self::parse_uservariables(&self.working_dir)?;
184 on_disk.insert(key.to_string(), val.clone());
185 let toml_string = toml::to_string(&on_disk)?;
186 let path = self.working_dir.join(".uservariables.toml");
187 fs::write(&path, toml_string)?;
188 self.user_variables.insert(key.to_string(), val);
189 Ok(())
190 }
191
192 pub fn parse_uservariables(cwd: &Path) -> Result<Table, anyhow::Error> {
193 let path = cwd.join(".uservariables.toml");
194 if path.exists() {
195 let content = fs::read_to_string(&path)?;
196 let table: Table = toml::de::from_str(&content).map_err(|e| {
197 anyhow::anyhow!(
198 "Failed to parse .uservariables.toml at '{}': {}",
199 path.display(),
200 e
201 )
202 })?;
203 Ok(table)
204 } else {
205 Ok(Table::new())
206 }
207 }
208
209 pub fn new(
210 working_dir: &Path,
211 conf: &Config,
212 profile_name: &Option<String>,
213 create_profile_if_missing: bool,
214 ) -> Result<(Self, bool), anyhow::Error> {
215 let mut variables = conf.variables.clone();
216 for (key, value) in std::env::vars() {
217 variables.insert(key, toml::Value::String(value));
218 }
219 let raw_user_variables = Self::parse_uservariables(working_dir)?;
233 if let Some(value) = raw_user_variables.get(DOTR_PROFILE_KEY) {
241 variables.insert(DOTR_PROFILE_KEY.to_string(), value.clone());
242 }
243 let (profile, created) = Self::get_profile_from_config(
244 conf,
245 profile_name,
246 create_profile_if_missing,
247 &variables,
248 )?;
249 variables.insert(
250 DOTR_PROFILE_KEY.to_string(),
251 toml::Value::String(profile.name.clone()),
252 );
253 let bitwarden_note = resolve_bitwarden_note(
258 env::var(BITWARDEN_NOTE_OVERRIDE_KEY).ok(),
259 &raw_user_variables,
260 profile.bitwarden_note.as_deref(),
261 conf.bitwarden_note.as_deref(),
262 );
263 variables.insert(
264 BITWARDEN_NOTE_OVERRIDE_KEY.to_string(),
265 toml::Value::String(bitwarden_note),
266 );
267 Ok((
268 Self {
269 working_dir: working_dir.to_path_buf(),
270 variables,
271 user_variables: raw_user_variables,
272 profile,
273 symlink: conf.symlink,
274 },
275 created,
276 ))
277 }
278
279 pub fn get_profile_from_config(
280 conf: &Config,
281 pname: &Option<String>,
282 create_if_missing: bool,
283 variables: &Table,
284 ) -> anyhow::Result<(Profile, bool)> {
285 let profile_name = match pname {
286 Some(name) => name.clone(),
287 None => {
288 if let Ok(env_p_name) = variables
289 .get(DOTR_PROFILE_KEY)
290 .and_then(|v| v.as_str())
291 .ok_or_else(|| anyhow::anyhow!("DOTR_PROFILE variable must be a string"))
292 {
293 env_p_name.to_string()
294 } else {
295 "default".to_string()
296 }
297 }
298 };
299
300 let profile = conf.profiles.get(&profile_name);
301 if let Some(prof) = profile {
302 return Ok((prof.clone(), false));
303 } else if !create_if_missing && profile_name != "default" {
304 anyhow::bail!("Profile {} not found", profile_name);
305 }
306 Ok((Profile::new(&profile_name), true))
307 }
308
309 pub fn get_variables(&self) -> &Table {
310 &self.variables
311 }
312
313 pub fn get_user_variables(&self) -> &Table {
314 &self.user_variables
315 }
316
317 pub fn get_context_variables(&self) -> Table {
318 let mut context_vars = self.variables.clone();
319 context_vars.extend(self.profile.variables.clone());
320 context_vars.extend(self.user_variables.clone());
321 context_vars
322 }
323
324 pub fn extend_variables(&mut self, new_vars: Table) {
325 self.variables.extend(new_vars);
326 }
327
328 pub fn print_variables(&self) {
329 let variables = &self.get_context_variables();
330 println!("User Variables:");
331 if variables.is_empty() {
332 println!(" (none)");
333 } else {
334 for (key, value) in variables.iter() {
335 print_variable(key, value, 1);
336 }
337 }
338 }
339}
340
341pub fn print_variable(key: &str, value: &toml::Value, level: usize) {
342 let indent = " ".repeat(level);
343 match value {
344 toml::Value::String(s) => {
345 println!("{}{} = {}", indent, key, s);
346 }
347 toml::Value::Integer(i) => {
348 println!("{}{} = {}", indent, key, i);
349 }
350 toml::Value::Float(f) => {
351 println!("{}{} = {}", indent, key, f);
352 }
353 toml::Value::Boolean(b) => {
354 println!("{}{} = {}", indent, key, b);
355 }
356 toml::Value::Table(t) => {
357 println!("{}{} =", indent, key);
358 for (k, v) in t.iter() {
359 print_variable(k, v, level + 1);
360 }
361 }
362 toml::Value::Array(arr) => {
363 println!("{}{} = [", indent, key);
364 for v in arr.iter() {
365 let item_indent = " ".repeat(level + 1);
366 match v {
367 toml::Value::String(s) => {
368 println!("{}- {}", item_indent, s);
369 }
370 toml::Value::Integer(i) => {
371 println!("{}- {}", item_indent, i);
372 }
373 toml::Value::Float(f) => {
374 println!("{}- {}", item_indent, f);
375 }
376 toml::Value::Boolean(b) => {
377 println!("{}- {}", item_indent, b);
378 }
379 toml::Value::Table(_) | toml::Value::Array(_) => {
380 println!("{}-", item_indent);
381 print_variable("", v, level + 2);
382 }
383 _ => {
384 println!("{}- {:?}", item_indent, v);
385 }
386 }
387 }
388 println!("{}]", indent);
389 }
390 _ => {
391 println!("{}{} = {:?}", indent, key, value);
392 }
393 }
394}
395
396fn resolve_bitwarden_note(
401 env_override: Option<String>,
402 user_variables: &Table,
403 profile_note: Option<&str>,
404 config_note: Option<&str>,
405) -> String {
406 use crate::prompt_store::DEFAULT_BITWARDEN_NOTE;
407
408 env_override
409 .or_else(|| {
410 user_variables
411 .get(BITWARDEN_NOTE_OVERRIDE_KEY)
412 .and_then(|v| v.as_str())
413 .map(|s| s.to_string())
414 })
415 .or_else(|| profile_note.map(|s| s.to_string()))
416 .or_else(|| config_note.map(|s| s.to_string()))
417 .unwrap_or_else(|| DEFAULT_BITWARDEN_NOTE.to_string())
418}
419
420fn get_prompted_variables<R: io::BufRead, W: io::Write>(
421 prompt: &str,
422 mut reader: R,
423 mut writer: W,
424) -> anyhow::Result<String> {
425 writer.write_all(format!("{}\n>>> ", prompt).as_bytes())?;
427 writer.flush()?;
428 let mut input = String::new();
429 reader.read_line(&mut input)?;
430 Ok(input)
431}