wash_cli/
util.rs

1use std::{
2    fs::File,
3    io::Read,
4    path::{Path, PathBuf},
5};
6
7use anyhow::{Context, Result};
8use term_table::{Table, TableStyle};
9use wash_lib::{
10    config::{cfg_dir, DEFAULT_NATS_TIMEOUT_MS},
11    plugin::{subcommand::SubcommandRunner, PLUGIN_DIR},
12};
13
14const MAX_TERMINAL_WIDTH: usize = 120;
15
16pub fn format_optional(value: Option<String>) -> String {
17    value.unwrap_or_else(|| "N/A".into())
18}
19
20/// Returns value from an argument that may be a file path or the value itself
21pub fn extract_arg_value(arg: &str) -> Result<String> {
22    match File::open(arg) {
23        Ok(mut f) => {
24            let mut value = String::new();
25            f.read_to_string(&mut value)
26                .with_context(|| format!("Failed to read file {}", &arg))?;
27            Ok(value)
28        }
29        Err(_) => Ok(arg.to_string()),
30    }
31}
32
33pub fn default_timeout_ms() -> u64 {
34    DEFAULT_NATS_TIMEOUT_MS
35}
36
37/// Transform a json string (e.g. "{"hello": "world"}") into msgpack bytes
38pub fn json_str_to_msgpack_bytes(payload: &str) -> Result<Vec<u8>> {
39    let json: serde_json::Value =
40        serde_json::from_str(payload).context("failed to encode string as JSON")?;
41    rmp_serde::to_vec_named(&json).context("failed to encode JSON as msgpack")
42}
43
44/// Ensure that the given plugin directory exists or return the default plugin dir path
45pub async fn ensure_plugin_dir(dir: Option<impl AsRef<Path>>) -> Result<PathBuf> {
46    let plugin_dir = match dir.map(|dir| dir.as_ref().to_path_buf()) {
47        Some(dir) => dir,
48        None => cfg_dir()?.join(PLUGIN_DIR),
49    };
50
51    if !tokio::fs::try_exists(&plugin_dir).await.unwrap_or(false) {
52        tokio::fs::create_dir(&plugin_dir).await?;
53    }
54    Ok(plugin_dir)
55}
56
57/// Helper for loading plugins from the given directory
58pub async fn load_plugins(plugin_dir: impl AsRef<Path>) -> anyhow::Result<SubcommandRunner> {
59    let mut readdir = tokio::fs::read_dir(&plugin_dir)
60        .await
61        .context("Unable to read plugin directory")?;
62
63    let mut plugins = SubcommandRunner::new().context("Could not initialize plugin runner")?;
64
65    // We load each plugin separately so we only warn if a plugin fails to load
66    while let Some(entry) = readdir.next_entry().await.transpose() {
67        let entry = match entry {
68            Ok(entry) => entry,
69            Err(e) => {
70                eprintln!("WARN: Could not read plugin directory entry. Skipping: {e:?}");
71                continue;
72            }
73        };
74
75        if entry
76            .file_type()
77            .await
78            .map(|ft| ft.is_file())
79            .unwrap_or(false)
80        {
81            if let Err(e) = plugins.add_plugin(entry.path()).await {
82                eprintln!("WARN: Couldn't load plugin, skipping: {:?}", e);
83            }
84        }
85    }
86
87    Ok(plugins)
88}
89
90use once_cell::sync::OnceCell;
91static BIN_STR: OnceCell<char> = OnceCell::new();
92
93fn msgpack_to_json(mval: rmpv::Value) -> serde_json::Value {
94    use rmpv::Value as RV;
95    use serde_json::Value as JV;
96    match mval {
97        RV::String(s) => JV::String(s.to_string()),
98        RV::Boolean(b) => JV::Bool(b),
99        RV::Array(v) => JV::Array(v.into_iter().map(msgpack_to_json).collect::<Vec<_>>()),
100        RV::F64(f) => JV::from(f),
101        RV::F32(f) => JV::from(f),
102        RV::Integer(i) => match (i.is_u64(), i.is_i64()) {
103            (true, _) => JV::from(i.as_u64().unwrap()),
104            (_, true) => JV::from(i.as_i64().unwrap()),
105            _ => JV::from(0u64),
106        },
107        RV::Map(vkv) => JV::Object(
108            vkv.into_iter()
109                .map(|(k, v)| {
110                    (
111                        k.as_str().unwrap_or_default().to_string(),
112                        msgpack_to_json(v),
113                    )
114                })
115                .collect::<serde_json::Map<_, _>>(),
116        ),
117        RV::Binary(v) => match BIN_STR.get().unwrap() {
118            's' => JV::String(String::from_utf8_lossy(&v).into_owned()),
119            '2' => serde_json::json!({
120                "str": String::from_utf8_lossy(&v),
121                "bin": v,
122            }),
123            /*'b'|*/ _ => JV::Array(v.into_iter().map(JV::from).collect::<Vec<_>>()),
124        },
125        RV::Ext(i, v) => serde_json::json!({
126            "type": i,
127            "data": v
128        }),
129        RV::Nil => JV::Bool(false),
130    }
131}
132
133/// transform msgpack bytes into json
134pub fn msgpack_to_json_val(msg: Vec<u8>, bin_str: char) -> serde_json::Value {
135    use bytes::Buf;
136
137    BIN_STR.set(bin_str).unwrap();
138
139    let bytes = bytes::Bytes::from(msg);
140    if let Ok(v) = rmpv::decode::value::read_value(&mut bytes.reader()) {
141        msgpack_to_json(v)
142    } else {
143        serde_json::json!({ "error": "Could not decode data" })
144    }
145}
146
147pub fn configure_table_style(table: &mut Table<'_>, num_rows: usize) {
148    table.style = empty_table_style();
149    table.separate_rows = false;
150
151    // Sets the max column width to ensure the table evenly fills the terminal width
152    table.max_column_width = termsize::get()
153        // Just slightly reducing the terminal width to account for padding
154        .map(|size| size.cols.saturating_sub(4) as usize)
155        .unwrap_or(MAX_TERMINAL_WIDTH)
156        / num_rows;
157}
158
159fn empty_table_style() -> TableStyle {
160    TableStyle {
161        top_left_corner: ' ',
162        top_right_corner: ' ',
163        bottom_left_corner: ' ',
164        bottom_right_corner: ' ',
165        outer_left_vertical: ' ',
166        outer_right_vertical: ' ',
167        outer_bottom_horizontal: ' ',
168        outer_top_horizontal: ' ',
169        intersection: ' ',
170        vertical: ' ',
171        horizontal: ' ',
172    }
173}
174
175mod test {
176    #[test]
177    fn test_safe_base64_parse_option() {
178        let base64_option = "config_b64=eyJhZGRyZXNzIjogIjAuMC4wLjA6ODA4MCJ9Cg==".to_string();
179        let mut expected = std::collections::HashMap::new();
180        expected.insert(
181            "config_b64".to_string(),
182            "eyJhZGRyZXNzIjogIjAuMC4wLjA6ODA4MCJ9Cg==".to_string(),
183        );
184        let output = wash_lib::cli::input_vec_to_hashmap(vec![base64_option]).unwrap();
185        assert_eq!(expected, output);
186    }
187}