1#[cfg(feature = "internal-commands")]
5use crate::commands::{FrameworkCommand, HelloCommand, HelpCommand, PingCommand};
6use crate::output::hook;
7
8use crate::command::Command;
9use crate::error::ModCliError;
10use std::collections::{HashMap, HashSet};
11
12type PreHookFn = dyn Fn(&str, &[String]) + Send + Sync;
14type PostHookFn = dyn Fn(&str, &[String], Result<(), &str>) + Send + Sync;
15type ErrorFmtFn = dyn Fn(&crate::error::ModCliError) -> String + Send + Sync;
16type VisibilityPolicyFn = dyn Fn(&dyn Command, &HashSet<String>) -> bool + Send + Sync;
17type AuthorizePolicyFn =
18 dyn Fn(&dyn Command, &HashSet<String>, &[String]) -> Result<(), String> + Send + Sync;
19
20pub struct CommandRegistry {
38 prefix: String,
39 commands: HashMap<String, Box<dyn Command>>,
40 aliases: HashMap<String, String>,
41 caps: HashSet<String>,
42 visibility_policy: Option<Box<VisibilityPolicyFn>>,
43 authorize_policy: Option<Box<AuthorizePolicyFn>>,
44 pre_hook: Option<Box<PreHookFn>>, post_hook: Option<Box<PostHookFn>>, error_formatter: Option<Box<ErrorFmtFn>>,
47 #[cfg(feature = "dispatch-cache")]
48 cache: std::sync::Mutex<Option<(String, String)>>,
49}
50
51impl Default for CommandRegistry {
52 fn default() -> Self {
53 Self::new()
54 }
55}
56
57impl CommandRegistry {
58 pub fn new() -> Self {
60 let mut reg = Self {
61 prefix: String::new(),
62 commands: HashMap::new(),
63 aliases: HashMap::new(),
64 caps: HashSet::new(),
65 visibility_policy: None,
66 authorize_policy: None,
67 pre_hook: None,
68 post_hook: None,
69 error_formatter: None,
70 #[cfg(feature = "dispatch-cache")]
71 cache: std::sync::Mutex::new(None),
72 };
73
74 #[cfg(feature = "custom-commands")]
75 reg.load_custom_commands();
76
77 #[cfg(feature = "internal-commands")]
78 reg.load_internal_commands();
79
80 reg
81 }
82
83 pub fn set_prefix(&mut self, prefix: &str) {
86 self.prefix = prefix.to_string();
87 }
88
89 pub fn get_prefix(&self) -> &str {
92 &self.prefix
93 }
94
95 #[inline(always)]
98 pub fn get(&self, name: &str) -> Option<&dyn Command> {
99 self.commands.get(name).map(|b| b.as_ref())
100 }
101
102 #[inline(always)]
105 pub fn register(&mut self, cmd: Box<dyn Command>) {
106 let name = cmd.name().to_string();
108 self.commands.insert(name.clone(), cmd);
109
110 for &alias in self.commands[&name].aliases() {
112 if !self.commands.contains_key(alias) {
114 self.aliases.insert(alias.to_string(), name.clone());
116 }
117 }
118 }
119
120 pub fn all(&self) -> impl Iterator<Item = &Box<dyn Command>> {
123 self.commands.values()
124 }
125
126 pub fn grant_cap<S: Into<String>>(&mut self, cap: S) {
128 self.caps.insert(cap.into());
129 }
130 pub fn revoke_cap(&mut self, cap: &str) {
131 self.caps.remove(cap);
132 }
133 pub fn has_cap(&self, cap: &str) -> bool {
134 self.caps.contains(cap)
135 }
136 pub fn set_caps<I, S>(&mut self, caps: I)
137 where
138 I: IntoIterator<Item = S>,
139 S: Into<String>,
140 {
141 self.caps.clear();
142 for c in caps {
143 self.caps.insert(c.into());
144 }
145 }
146
147 pub fn set_visibility_policy<F>(&mut self, f: F)
148 where
149 F: Fn(&dyn Command, &HashSet<String>) -> bool + Send + Sync + 'static,
150 {
151 self.visibility_policy = Some(Box::new(f));
152 }
153
154 pub fn set_authorize_policy<F>(&mut self, f: F)
155 where
156 F: Fn(&dyn Command, &HashSet<String>, &[String]) -> Result<(), String>
157 + Send
158 + Sync
159 + 'static,
160 {
161 self.authorize_policy = Some(Box::new(f));
162 }
163
164 pub fn set_pre_hook<F>(&mut self, f: F)
165 where
166 F: Fn(&str, &[String]) + Send + Sync + 'static,
167 {
168 self.pre_hook = Some(Box::new(f));
169 }
170
171 pub fn set_post_hook<F>(&mut self, f: F)
172 where
173 F: Fn(&str, &[String], Result<(), &str>) + Send + Sync + 'static,
174 {
175 self.post_hook = Some(Box::new(f));
176 }
177
178 pub fn set_error_formatter<F>(&mut self, f: F)
179 where
180 F: Fn(&crate::error::ModCliError) -> String + Send + Sync + 'static,
181 {
182 self.error_formatter = Some(Box::new(f));
183 }
184
185 #[inline(always)]
186 pub fn is_visible(&self, cmd: &dyn Command) -> bool {
187 if let Some(ref pol) = self.visibility_policy {
188 return pol(cmd, &self.caps);
189 }
190 if cmd.hidden() {
191 return false;
192 }
193 cmd.required_caps().iter().all(|c| self.caps.contains(*c))
194 }
195
196 #[inline(always)]
197 pub fn is_authorized(&self, cmd: &dyn Command, args: &[String]) -> Result<(), String> {
198 if let Some(ref pol) = self.authorize_policy {
199 return pol(cmd, &self.caps, args);
200 }
201 if cmd.required_caps().iter().all(|c| self.caps.contains(*c)) {
202 Ok(())
203 } else {
204 Err("Not authorized".into())
205 }
206 }
207
208 #[inline(always)]
227 pub fn execute(&self, cmd: &str, args: &[String]) {
228 if let Err(err) = self.try_execute(cmd, args) {
229 if let Some(ref fmt) = self.error_formatter {
230 hook::error(&fmt(&err));
231 } else {
232 match err {
233 ModCliError::InvalidUsage(msg) => hook::error(&format!("Invalid usage: {msg}")),
234 ModCliError::UnknownCommand(name) => hook::unknown(&format!(
235 "[{name}]. Type `help` or `--help` for a list of available commands."
236 )),
237 other => hook::error(&format!("{other}")),
238 }
239 }
240 }
241 }
242
243 #[inline(always)]
265 pub fn try_execute(&self, cmd: &str, args: &[String]) -> Result<(), ModCliError> {
266 if let Some(ref pre) = self.pre_hook {
267 pre(cmd, args);
268 }
269 let token: &str = if !self.prefix.is_empty() && cmd.len() > self.prefix.len() + 1 {
271 let (maybe_prefix, rest_with_colon) = cmd.split_at(self.prefix.len());
272 if maybe_prefix == self.prefix && rest_with_colon.as_bytes().first() == Some(&b':') {
273 &rest_with_colon[1..]
274 } else {
275 cmd
276 }
277 } else {
278 cmd
279 };
280
281 #[cfg(feature = "dispatch-cache")]
282 if let Ok(guard) = self.cache.lock() {
283 if let Some((ref t, ref p)) = *guard {
284 if t == token {
285 if let Some(command) = self.commands.get(p.as_str()) {
286 if let Err(err) = command.validate(args) {
287 return Err(ModCliError::InvalidUsage(err));
288 }
289 command.execute_with(args, self);
290 return Ok(());
291 }
292 }
293 }
294 }
295
296 if let Some(command) = self.commands.get(token) {
298 if let Err(err) = self.is_authorized(command.as_ref(), args) {
299 return Err(ModCliError::InvalidUsage(err));
300 }
301 if let Err(err) = command.validate(args) {
302 return Err(ModCliError::InvalidUsage(err));
303 }
304 command.execute_with(args, self);
305 #[cfg(feature = "dispatch-cache")]
306 if let Ok(mut guard) = self.cache.lock() {
307 *guard = Some((token.to_string(), token.to_string()));
308 }
309 if let Some(ref post) = self.post_hook {
310 post(cmd, args, Ok(()));
311 }
312 return Ok(());
313 }
314
315 if let Some(primary) = self.aliases.get(token) {
317 if let Some(command) = self.commands.get(primary.as_str()) {
318 if let Err(err) = self.is_authorized(command.as_ref(), args) {
319 return Err(ModCliError::InvalidUsage(err));
320 }
321 if let Err(err) = command.validate(args) {
322 return Err(ModCliError::InvalidUsage(err));
323 }
324 command.execute_with(args, self);
325 #[cfg(feature = "dispatch-cache")]
326 if let Ok(mut guard) = self.cache.lock() {
327 *guard = Some((token.to_string(), primary.clone()));
328 }
329 if let Some(ref post) = self.post_hook {
330 post(cmd, args, Ok(()));
331 }
332 return Ok(());
333 }
334 }
335
336 if !args.is_empty() {
338 let combined = format!("{token}:{}", args[0]);
339 if let Some(command) = self.commands.get(combined.as_str()) {
340 let rest = &args[1..];
341 if let Err(err) = self.is_authorized(command.as_ref(), rest) {
342 return Err(ModCliError::InvalidUsage(err));
343 }
344 if let Err(err) = command.validate(rest) {
345 return Err(ModCliError::InvalidUsage(err));
346 }
347 command.execute_with(rest, self);
348 if let Some(ref post) = self.post_hook {
349 post(cmd, args, Ok(()));
350 }
351 return Ok(());
352 }
353 }
354 let err = ModCliError::UnknownCommand(cmd.to_string());
355 if let Some(ref post) = self.post_hook {
356 post(cmd, args, Err("unknown"));
357 }
358 Err(err)
359 }
360
361 #[cfg(feature = "internal-commands")]
362 pub fn load_internal_commands(&mut self) {
363 self.register(Box::new(PingCommand));
364 self.register(Box::new(HelloCommand));
365 self.register(Box::new(FrameworkCommand));
366 self.register(Box::new(HelpCommand::new()));
367 }
368
369 pub fn len(&self) -> usize {
372 self.commands.len()
373 }
374
375 pub fn is_empty(&self) -> bool {
376 self.commands.is_empty()
377 }
378
379 #[cfg(feature = "custom-commands")]
380 pub fn load_custom_commands(&mut self) {
381 }
383}