korasi_cli/
util.rs

1//! IO Utilities wrapper to allow automock for requests and user input prompts.
2
3use std::{
4    fmt::{self, Display},
5    io::Write,
6    iter::Map,
7    path::{Path, PathBuf},
8};
9
10use aws_sdk_ec2::types::{Image, Instance, InstanceStateName, InstanceType};
11use ignore::Walk;
12use inquire::{InquireError, MultiSelect, Select};
13
14use crate::ec2::{EC2Error, EC2Impl as EC2};
15
16#[derive(Default)]
17pub struct UtilImpl;
18
19impl UtilImpl {
20    /// Utility to perform a GET request and return the body as UTF-8, or an appropriate EC2Error.
21    pub async fn do_get(url: &str) -> Result<String, EC2Error> {
22        reqwest::get(url)
23            .await
24            .map_err(|e| EC2Error::new(format!("Could not request ip from {url}: {e:?}")))?
25            .error_for_status()
26            .map_err(|e| EC2Error::new(format!("Failure status from {url}: {e:?}")))?
27            .text_with_charset("utf-8")
28            .await
29            .map_err(|e| EC2Error::new(format!("Failed to read response from {url}: {e:?}")))
30    }
31
32    pub fn write_secure(path: &PathBuf, material: String, mode: u32) -> Result<(), EC2Error> {
33        let mut file = open_file_with_perm(path, mode)?;
34        file.write(material.as_bytes())
35            .map_err(|e| EC2Error::new(format!("Failed to write to {path:?} ({e:?})")))?;
36        Ok(())
37    }
38}
39
40#[cfg(unix)]
41fn open_file_with_perm(path: &PathBuf, mode: u32) -> Result<std::fs::File, EC2Error> {
42    use std::os::unix::fs::OpenOptionsExt;
43    std::fs::OpenOptions::new()
44        .mode(mode)
45        .write(true)
46        .create(true)
47        .open(path)
48        .map_err(|e| EC2Error::new(format!("Failed to create {path:?} ({e:?})")))
49}
50
51#[cfg(not(unix))]
52fn open_file(path: &PathBuf) -> Result<File, EC2Error> {
53    fs::File::create(path.clone())
54        .map_err(|e| EC2Error::new(format!("Failed to create {path:?} ({e:?})")))?
55}
56
57/// Image doesn't impl Display, which is necessary for inquire to use it in a Select.
58/// This wraps Image and provides a Display impl.
59#[derive(PartialEq, Debug)]
60pub struct ScenarioImage(pub Image);
61impl From<Image> for ScenarioImage {
62    fn from(value: Image) -> Self {
63        ScenarioImage(value)
64    }
65}
66
67impl Display for ScenarioImage {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        write!(
70            f,
71            "{}: {}",
72            self.0.name().unwrap_or("(unknown)"),
73            self.0.description().unwrap_or("unknown")
74        )
75    }
76}
77
78#[derive(Debug, Default, Clone)]
79pub struct SelectOption {
80    pub name: String,
81    pub instance_id: String,
82    pub public_dns_name: Option<String>,
83    state: Option<InstanceStateName>,
84    instance_type: Option<InstanceType>,
85}
86
87impl fmt::Display for SelectOption {
88    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
89        let status = self.state.as_ref().unwrap().clone();
90        write!(
91            f,
92            "name = {}, type = {}, instance_id = {}, status = {}",
93            self.name,
94            self.instance_type.as_ref().unwrap(),
95            self.instance_id,
96            status
97        )
98    }
99}
100
101impl From<Instance> for SelectOption {
102    fn from(value: Instance) -> Self {
103        let mut opt = SelectOption {
104            state: value.state().unwrap().name().cloned(),
105            instance_id: value.instance_id().unwrap().to_string(),
106            public_dns_name: value.public_dns_name().map(str::to_string),
107            ..SelectOption::default()
108        };
109
110        opt.instance_type = value.instance_type().cloned();
111        for t in value.tags() {
112            if t.key() == Some("Name") {
113                opt.name = t.value().unwrap().to_owned();
114            }
115        }
116
117        opt
118    }
119}
120
121/// Express list of instance ids as a comma separated string.
122pub fn ids_to_str(ids: Vec<SelectOption>) -> String {
123    ids.iter()
124        .map(|i| i.instance_id.to_owned())
125        .collect::<Vec<_>>()
126        .join(",")
127}
128
129pub async fn multi_select_instances(
130    ec2: &EC2,
131    prompt: &str,
132    statuses: Vec<InstanceStateName>,
133) -> Result<Vec<SelectOption>, InquireError> {
134    // Get all instances tagged by this tool.
135    let instances = ec2.describe_instance(statuses).await.unwrap();
136    let options: Vec<SelectOption> = instances.into_iter().map(|i| i.into()).collect();
137
138    if options.len() == 1 {
139        return Ok(vec![options[0].to_owned()]);
140    }
141    MultiSelect::new(prompt, options)
142        .with_vim_mode(true)
143        .prompt()
144}
145
146pub async fn select_instance(
147    ec2: &EC2,
148    prompt: &str,
149    statuses: Vec<InstanceStateName>,
150) -> Result<SelectOption, InquireError> {
151    let instances = ec2.describe_instance(statuses).await.unwrap();
152    let options: Vec<SelectOption> = instances.into_iter().map(|i| i.into()).collect();
153
154    if options.len() == 1 {
155        return Ok(options[0].to_owned());
156    }
157    Select::new(prompt, options).with_vim_mode(true).prompt()
158}
159
160pub fn calc_prefix(pth: PathBuf) -> std::io::Result<PathBuf> {
161    Ok(pth.parent().unwrap_or(Path::new("")).to_path_buf())
162}
163
164pub fn biject_paths<'a>(
165    src_path: &str,
166    prefix: &'a str,
167    dst_folder: &'a str,
168) -> Map<
169    Walk,
170    impl FnMut(
171            Result<ignore::DirEntry, ignore::Error>,
172        ) -> Result<(PathBuf, PathBuf, bool), ignore::Error>
173        + 'a,
174> {
175    Walk::new(src_path).map(move |result| match result {
176        Ok(entry) => {
177            let is_dir = match entry.metadata() {
178                Ok(ent) => ent.is_dir(),
179                _ => false,
180            };
181            let local_pth = entry.path().to_path_buf();
182            let mut rel_pth = entry
183                .path()
184                .to_str()
185                .unwrap()
186                .strip_prefix(prefix)
187                .unwrap()
188                .chars();
189            rel_pth.next();
190            let transformed = PathBuf::from(dst_folder).join(rel_pth.as_str());
191
192            tracing::info!("uploaded path = {:?}", transformed);
193
194            Ok((local_pth, transformed, is_dir))
195        }
196        Err(err) => Err(err),
197    })
198}
199
200#[cfg(test)]
201mod tests {
202    use std::{
203        fs::remove_file,
204        path::{Path, PathBuf},
205    };
206
207    use crate::util::biject_paths;
208
209    use super::{calc_prefix, open_file_with_perm};
210
211    #[test]
212    fn open_readonly_file() {
213        let pk_file = "pk.pem";
214
215        assert!(
216            !Path::new(pk_file).exists(),
217            "Test pk file should not exist before test."
218        );
219        let _ = open_file_with_perm(&pk_file.into(), 0o400);
220        let meta = std::fs::metadata(pk_file).unwrap();
221        assert!(
222            meta.permissions().readonly(),
223            "ssh PK file should be readonly."
224        );
225        let _ = remove_file(pk_file);
226    }
227
228    #[test]
229    fn calc_src_prefix() {
230        let _ = std::fs::remove_dir("../outside-cwd");
231
232        let cwd = std::env::current_dir().unwrap();
233        std::fs::create_dir("../outside-cwd").unwrap();
234
235        let cases = [
236            ("/", PathBuf::from("")),
237            ("README.md", cwd.clone()),
238            ("src/main.rs", cwd.join("src")),
239            ("../outside-cwd", cwd.parent().unwrap().to_path_buf()),
240        ];
241
242        for (input, expected) in cases {
243            println!("input = {input}");
244            let canon_pth = std::fs::canonicalize(input).unwrap();
245            let got = calc_prefix(canon_pth);
246            assert!(
247                got.is_ok(),
248                "Failed to canonicalize path = {}, Err = {}",
249                input,
250                got.unwrap_err()
251            );
252            pretty_assertions::assert_eq!(got.unwrap(), expected);
253        }
254
255        std::fs::remove_dir("../outside-cwd").unwrap();
256    }
257
258    #[test]
259    fn calc_remote_paths() {
260        let cwd = std::env::current_dir().unwrap();
261
262        let cases = [
263            (
264                // Paths are unchanged
265                cwd.as_path().to_str().unwrap(),
266                "",
267                "/home/foobar",
268            ),
269            (
270                // Paths prefixes are replaced
271                cwd.as_path().to_str().unwrap(),
272                cwd.parent().unwrap().to_str().unwrap(),
273                "/home/foobar",
274            ),
275        ];
276
277        for (x, y, z) in cases {
278            for result in biject_paths(x, y, z) {
279                match result {
280                    Ok(entry) => {
281                        println!("entry = {:?}", entry);
282                    }
283                    Err(err) => {
284                        println!("err = {}", err);
285                    }
286                }
287            }
288            println!();
289        }
290    }
291}
korasi_cli/util.rs

korasi_cli/
util.rs
