falcon_cli/
router.rs

1use crate::help::CliHelpScreen;
2use crate::CliCommand;
3use crate::*;
4use std::collections::HashMap;
5use std::env;
6use strsim::levenshtein;
7
8pub struct CliRouter {
9    pub commands: HashMap<String, Box<dyn CliCommand>>,
10    pub shortcuts: HashMap<String, String>,
11    pub value_flags: HashMap<String, Vec<String>>,
12    pub categories: HashMap<String, (String, String)>,
13}
14
15pub struct CliRequest {
16    pub cmd_alias: String,
17    pub is_help: bool,
18    pub args: Vec<String>,
19    pub flags: Vec<String>,
20    pub value_flags: HashMap<String, String>,
21    pub shortcuts: Vec<String>,
22}
23
24impl Default for CliRouter {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl CliRouter {
31    pub fn new() -> Self {
32        Self {
33            commands: HashMap::new(),
34            shortcuts: HashMap::new(),
35            value_flags: HashMap::new(),
36            categories: HashMap::new(),
37        }
38    }
39
40    /// Link a struct / impl that inherits the CliRouter trait to a command name.  Takes three arguments,
41    /// the full name of the command, a vector of available shortcuts, and a vector of
42    /// long-form flags (prefixed with dashes (--)) for which a value is expected.
43    pub fn add<T: CliCommand + Default + 'static>(
44        &mut self,
45        alias: &str,
46        shortcuts: Vec<&str>,
47        value_flags: Vec<&str>,
48    ) {
49        // Add to list of commands
50        let cmd = Box::<T>::default();
51        self.commands.insert(alias.to_string(), cmd);
52
53        // Add shortcuts
54        for shortcut in shortcuts {
55            self.shortcuts
56                .insert(shortcut.to_string(), alias.to_string());
57        }
58
59        // Add value flags
60        let flags: Vec<String> = value_flags.iter().map(|s| s.to_string()).collect();
61        self.value_flags.insert(alias.to_string(), flags);
62    }
63
64    /// Taking arguments passed via command line into account,  checks all routes that were added and
65    /// tries to determine the correct impl to execute.  This function is automatically
66    /// executed by the cli_run() function and should never be manually executed.
67    pub fn lookup(&self) -> (&Box<dyn CliCommand>, CliRequest) {
68        // Get args
69        let mut cmdargs: Vec<String> = env::args().collect();
70        cmdargs.remove(0);
71        if cmdargs.is_empty() {
72            cli_error!("You did not specify a command to run.  Please specify a command or use 'help' or '-h' to view a list of all available commands.");
73            std::process::exit(0);
74        }
75
76        // Blank variables
77        let mut extra_args: Vec<String> = Vec::new();
78        let mut cmd_alias = String::new();
79
80        // Check if help
81        let mut is_help: bool = false;
82        if cmdargs.len() > 0 && (cmdargs[0] == "help" || cmdargs[0] == "-h") {
83            is_help = true;
84            cmdargs.remove(0);
85        }
86
87        // Check for help index
88        if is_help && cmdargs.is_empty() {
89            CliHelpScreen::render_index(self);
90        }
91
92        // Check routing table for command
93        loop {
94            // Check for zero cmdargs
95            if cmdargs.is_empty() && is_help {
96                break;
97            } else if cmdargs.is_empty() {
98                cli_error!("No command exists with that name.  Use either 'help' or the -h flag to view a list of all available commands.");
99                std::process::exit(0);
100            }
101            let alias = cmdargs.iter().map(|v| v.to_string()).filter(|a| !a.starts_with("-")).collect::<Vec<String>>().join(" ").to_string();
102
103            if self.commands.contains_key(&alias) {
104                cmd_alias = alias;
105                break;
106            } else if self.shortcuts.contains_key(&alias) {
107                cmd_alias = self.shortcuts.get(&alias).unwrap().to_string();
108                break;
109            } else if is_help && self.categories.contains_key(&alias) {
110                CliHelpScreen::render_category(self, &alias);
111            } else if let Some(found_cmd) = self.lookup_similar(&alias) {
112                let confirm_msg = format!("No command with that name exists, but a similar command with the name '{}' does exist.  Is this the command you wish to run?", found_cmd);
113                if cli_confirm(&confirm_msg) {
114                    cmd_alias = found_cmd.to_string();
115                    break;
116                }
117            }
118            extra_args.insert(0, cmdargs.pop().unwrap());
119        }
120
121        // Set variables
122        let mut args: Vec<String> = Vec::new();
123        let mut flags: Vec<String> = Vec::new();
124        let mut value_flags: HashMap<String, String> = HashMap::new();
125        let flag_values = self.value_flags.get(&cmd_alias).unwrap_or(&Vec::new()).clone();
126
127        // Get flags
128        while !extra_args.is_empty() {
129            let chk_arg = extra_args[0].to_string();
130            extra_args.remove(0);
131
132            if chk_arg.starts_with("--") {
133                let arg = chk_arg.trim_start_matches('-').to_string();
134                if flag_values.contains(&arg) {
135                    value_flags.insert(arg, extra_args[0].to_string());
136                    extra_args.remove(0);
137                } else {
138                    flags.push(arg);
139                }
140            } else if chk_arg.starts_with('-') {
141                let arg = chk_arg.trim_start_matches('-').to_string();
142                for c in arg.chars() {
143                    flags.push(c.to_string());
144                }
145            } else {
146                args.push(chk_arg);
147            }
148        }
149
150        // Get all shortcuts
151        let shortcuts: Vec<String> = self
152            .shortcuts
153            .iter()
154            .filter_map(|(key, value)| {
155                if *value == cmd_alias {
156                    Some(key.to_string())
157                } else {
158                    None
159                }
160            })
161            .collect();
162
163        let cmd = self.commands.get(&cmd_alias).unwrap();
164        let req = CliRequest {
165            cmd_alias,
166            is_help,
167            args,
168            flags,
169            value_flags,
170            shortcuts,
171        };
172
173        (cmd, req)
174    }
175
176    /// Never needs to be manually executed, and is used when a full match for the command name can not
177    /// be found.  Uses the levenshtein to see if any commands closely resemble the
178    /// command name given in case of typo.
179    fn lookup_similar(&self, chk_cmd: &String) -> Option<&String> {
180        let mut distance = 4;
181        let mut res = None;
182
183        for cmd in self.commands.keys() {
184            let num = levenshtein(chk_cmd, cmd);
185            if num < distance {
186                distance = num;
187                res = Some(cmd)
188            }
189        }
190
191        res
192    }
193
194    /// Add a new top level category that contains CLI commands.  Used for organization and to display proper help screens.
195    pub fn add_category(&mut self, alias: &str, name: &str, description: &str) {
196        self.categories.insert(
197            alias.to_string(),
198            (name.to_string(), description.to_string()),
199        );
200    }
201}