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)?;
236 if let Some(value) = raw_user_variables.get(DOTR_PROFILE_KEY) {
244 variables.insert(DOTR_PROFILE_KEY.to_string(), value.clone());
245 }
246 let (profile, created) = Self::get_profile_from_config(
247 conf,
248 profile_name,
249 create_profile_if_missing,
250 &variables,
251 )?;
252 variables.insert(
253 DOTR_PROFILE_KEY.to_string(),
254 toml::Value::String(profile.name.clone()),
255 );
256 let bitwarden_note = resolve_bitwarden_note(
261 env::var(BITWARDEN_NOTE_OVERRIDE_KEY).ok(),
262 &raw_user_variables,
263 profile.bitwarden_note.as_deref(),
264 conf.bitwarden_note.as_deref(),
265 );
266 variables.insert(
267 BITWARDEN_NOTE_OVERRIDE_KEY.to_string(),
268 toml::Value::String(bitwarden_note),
269 );
270 Ok((
271 Self {
272 working_dir: working_dir.to_path_buf(),
273 variables,
274 user_variables: Table::new(),
275 profile,
276 symlink: conf.symlink,
277 },
278 created,
279 ))
280 }
281
282 pub fn get_profile_from_config(
283 conf: &Config,
284 pname: &Option<String>,
285 create_if_missing: bool,
286 variables: &Table,
287 ) -> anyhow::Result<(Profile, bool)> {
288 let profile_name = match pname {
289 Some(name) => name.clone(),
290 None => {
291 if let Ok(env_p_name) = variables
292 .get(DOTR_PROFILE_KEY)
293 .and_then(|v| v.as_str())
294 .ok_or_else(|| anyhow::anyhow!("DOTR_PROFILE variable must be a string"))
295 {
296 env_p_name.to_string()
297 } else {
298 "default".to_string()
299 }
300 }
301 };
302
303 let profile = conf.profiles.get(&profile_name);
304 if let Some(prof) = profile {
305 return Ok((prof.clone(), false));
306 } else if !create_if_missing && profile_name != "default" {
307 anyhow::bail!("Profile {} not found", profile_name);
308 }
309 Ok((Profile::new(&profile_name), true))
310 }
311
312 pub fn get_variables(&self) -> &Table {
313 &self.variables
314 }
315
316 pub fn get_user_variables(&self) -> &Table {
317 &self.user_variables
318 }
319
320 pub fn get_context_variables(&self) -> Table {
321 let mut context_vars = self.variables.clone();
322 context_vars.extend(self.profile.variables.clone());
323 context_vars.extend(self.user_variables.clone());
324 context_vars
325 }
326
327 pub fn extend_variables(&mut self, new_vars: Table) {
328 self.variables.extend(new_vars);
329 }
330
331 pub fn print_variables(&self) {
332 let variables = &self.get_context_variables();
333 println!("User Variables:");
334 if variables.is_empty() {
335 println!(" (none)");
336 } else {
337 for (key, value) in variables.iter() {
338 print_variable(key, value, 1);
339 }
340 }
341 }
342}
343
344pub fn print_variable(key: &str, value: &toml::Value, level: usize) {
345 let indent = " ".repeat(level);
346 match value {
347 toml::Value::String(s) => {
348 println!("{}{} = {}", indent, key, s);
349 }
350 toml::Value::Integer(i) => {
351 println!("{}{} = {}", indent, key, i);
352 }
353 toml::Value::Float(f) => {
354 println!("{}{} = {}", indent, key, f);
355 }
356 toml::Value::Boolean(b) => {
357 println!("{}{} = {}", indent, key, b);
358 }
359 toml::Value::Table(t) => {
360 println!("{}{} =", indent, key);
361 for (k, v) in t.iter() {
362 print_variable(k, v, level + 1);
363 }
364 }
365 toml::Value::Array(arr) => {
366 println!("{}{} = [", indent, key);
367 for v in arr.iter() {
368 let item_indent = " ".repeat(level + 1);
369 match v {
370 toml::Value::String(s) => {
371 println!("{}- {}", item_indent, s);
372 }
373 toml::Value::Integer(i) => {
374 println!("{}- {}", item_indent, i);
375 }
376 toml::Value::Float(f) => {
377 println!("{}- {}", item_indent, f);
378 }
379 toml::Value::Boolean(b) => {
380 println!("{}- {}", item_indent, b);
381 }
382 toml::Value::Table(_) | toml::Value::Array(_) => {
383 println!("{}-", item_indent);
384 print_variable("", v, level + 2);
385 }
386 _ => {
387 println!("{}- {:?}", item_indent, v);
388 }
389 }
390 }
391 println!("{}]", indent);
392 }
393 _ => {
394 println!("{}{} = {:?}", indent, key, value);
395 }
396 }
397}
398
399fn resolve_bitwarden_note(
404 env_override: Option<String>,
405 user_variables: &Table,
406 profile_note: Option<&str>,
407 config_note: Option<&str>,
408) -> String {
409 use crate::prompt_store::DEFAULT_BITWARDEN_NOTE;
410
411 env_override
412 .or_else(|| {
413 user_variables
414 .get(BITWARDEN_NOTE_OVERRIDE_KEY)
415 .and_then(|v| v.as_str())
416 .map(|s| s.to_string())
417 })
418 .or_else(|| profile_note.map(|s| s.to_string()))
419 .or_else(|| config_note.map(|s| s.to_string()))
420 .unwrap_or_else(|| DEFAULT_BITWARDEN_NOTE.to_string())
421}
422
423fn get_prompted_variables<R: io::BufRead, W: io::Write>(
424 prompt: &str,
425 mut reader: R,
426 mut writer: W,
427) -> anyhow::Result<String> {
428 writer.write_all(format!("{}\n>>> ", prompt).as_bytes())?;
430 writer.flush()?;
431 let mut input = String::new();
432 reader.read_line(&mut input)?;
433 Ok(input)
434}