aws_unlock/
aws_lock.rs

1use std::{
2    collections::HashMap,
3    io::{stdin, stdout, Write},
4};
5
6use anyhow::{bail, Result};
7use itertools::Itertools;
8
9use crate::aws_profile::{AwsFile, ProfileName};
10
11#[derive(Debug)]
12pub struct AwsLockGuard<'a> {
13    target_profiles: &'a [ProfileName],
14}
15
16impl<'a> AwsLockGuard<'a> {
17    pub fn unlock(
18        target_profiles: &'a [ProfileName],
19        error_if_not_exist: bool,
20        warn_on_production: bool,
21    ) -> Result<Self> {
22        modify_lock_status(
23            target_profiles,
24            error_if_not_exist,
25            warn_on_production,
26            false,
27        )?;
28        Ok(Self { target_profiles })
29    }
30
31    pub fn lock(self) {
32        // Drop re-locks profiles. No need to do anything.
33    }
34}
35
36impl Drop for AwsLockGuard<'_> {
37    fn drop(&mut self) {
38        let _ = modify_lock_status(self.target_profiles, false, false, true);
39    }
40}
41
42pub fn check_current_lock_status(
43    target_profiles: &[ProfileName],
44) -> Result<(Vec<ProfileName>, Vec<ProfileName>)> {
45    let mut locked_profiles = vec![];
46    let mut unlocked_profiles = vec![];
47
48    let mut aws_file = AwsFile::open()?;
49    let profiles = aws_file.parse()?;
50
51    // check all target profiles exist
52    let mut unknown_profiles = vec![];
53    for profile in target_profiles {
54        if !profiles.iter().any(|p| p.name == *profile) {
55            unknown_profiles.push(profile.clone());
56        }
57    }
58    if !unknown_profiles.is_empty() {
59        bail!(
60            "some target profiles not found: {}",
61            unknown_profiles
62                .iter()
63                .map(|s| format!("'{s}'"))
64                .format(", ")
65        )
66    }
67
68    for profile in profiles {
69        if target_profiles.contains(&profile.name) {
70            if profile.is_locked {
71                locked_profiles.push(profile.name)
72            } else {
73                unlocked_profiles.push(profile.name)
74            }
75        }
76    }
77
78    Ok((locked_profiles, unlocked_profiles))
79}
80
81fn modify_lock_status(
82    target_profiles: &[ProfileName],
83    error_if_not_exist: bool,
84    warn_on_production: bool,
85    lock: bool,
86) -> Result<()> {
87    let mut aws_file = AwsFile::open()?;
88
89    let mut profiles = aws_file.parse()?;
90    let profile_indices: HashMap<_, _> = profiles
91        .iter()
92        .enumerate()
93        .map(|(index, profile)| (profile.name.clone(), index))
94        .collect();
95
96    if error_if_not_exist {
97        // Check profiles exist if non-existence is explicit error
98        let unknown_profiles: Vec<_> = target_profiles
99            .iter()
100            .filter(|name| !profile_indices.contains_key(name))
101            .collect();
102
103        if !unknown_profiles.is_empty() {
104            bail!(
105                "unknown profiles: {}",
106                unknown_profiles
107                    .into_iter()
108                    .map(|s| format!("'{s}'"))
109                    .format(", ")
110            );
111        }
112    }
113
114    if warn_on_production {
115        // Warn if target profile contains production profile
116        let production_profiles: Vec<_> = target_profiles
117            .iter()
118            .filter(|name| profile_indices.contains_key(name))
119            .filter(|name| profiles[profile_indices[name]].is_production)
120            .collect();
121
122        if !production_profiles.is_empty() {
123            print!(
124                "You are unlocking production profiles: {}. Are you sure? (y/N) ",
125                production_profiles
126                    .into_iter()
127                    .map(|s| format!("'{s}'"))
128                    .format(", ")
129            );
130            stdout().flush()?;
131            let mut buf = String::new();
132            stdin().read_line(&mut buf)?;
133            if !["y", "Y"].contains(&buf.trim()) {
134                bail!("Unlocking production profiles cancelled by user");
135            }
136        }
137    }
138
139    // Lock target profiles
140    target_profiles
141        .iter()
142        .filter(|name| profile_indices.contains_key(name))
143        .for_each(|name| profiles[profile_indices[name]].is_locked = lock);
144
145    // Write to file
146    aws_file.write(&profiles)?;
147    aws_file.flush()?;
148
149    Ok(())
150}