dnslib/cli/
completions.rs

1#[cfg(any(
2    feature = "technitium",
3    feature = "pangolin",
4    feature = "cloudflare",
5    feature = "unifi",
6    feature = "pihole"
7))]
8pub fn generate_completions(shell: clap_complete::Shell) {
9    use clap::CommandFactory;
10    use clap_complete::generate;
11    use std::io::{self, Write};
12
13    use crate::cli::Cli;
14
15    let mut cmd = Cli::command();
16    let bin_name = std::env::current_exe()
17        .ok()
18        .and_then(|p| p.file_name().map(|n| n.to_string_lossy().into_owned()))
19        .unwrap_or_else(|| cmd.get_name().to_string());
20    let fn_name = bin_name.replace('-', "_");
21
22    let mut out = io::stdout();
23
24    // For zsh, patch the generated output so --server specs point at our
25    // dynamic helper instead of the default (_default) completer.
26    if shell == clap_complete::Shell::Zsh {
27        let mut buf: Vec<u8> = Vec::new();
28        generate(shell, &mut cmd, &bin_name, &mut buf);
29        let raw = String::from_utf8_lossy(&buf);
30        let patched = patch_zsh_server_completion(&raw, &fn_name);
31        out.write_all(patched.as_bytes()).ok();
32        let helper = format!(
33            "\n# Dynamic --server completion from config\n\
34             _{fn_name}_server_ids() {{\n\
35             \tlocal -a ids=(\"${{(@f)$({bin_name} _servers 2>/dev/null)}}\")\n\
36             \t_describe 'server ID' ids\n\
37             }}\n"
38        );
39        out.write_all(helper.as_bytes()).ok();
40        return;
41    }
42
43    generate(shell, &mut cmd, &bin_name, &mut out);
44
45    let dynamic = match shell {
46        clap_complete::Shell::Fish => format!(
47            "\n# Dynamic --server completion from config\n\
48             complete -e -c {bin_name} -l server\n\
49             complete -c {bin_name} -l server -r -a '({bin_name} _servers 2>/dev/null)'\n"
50        ),
51        clap_complete::Shell::Bash => format!(
52            "\n# Dynamic --server completion from config\n\
53             __{fn_name}_complete() {{\n\
54             \tlocal cur prev\n\
55             \tcur=\"${{COMP_WORDS[COMP_CWORD]}}\"\n\
56             \tprev=\"${{COMP_WORDS[COMP_CWORD-1]}}\"\n\
57             \tif [[ \"$cur\" == --server=* ]]; then\n\
58             \t\tlocal value=\"${{cur#--server=}}\"\n\
59             \t\tmapfile -t COMPREPLY < <(compgen -P \"--server=\" -W \"$({bin_name} _servers 2>/dev/null)\" -- \"$value\")\n\
60             \t\treturn\n\
61             \tfi\n\
62             \tif [[ \"$prev\" == \"--server\" ]]; then\n\
63             \t\tmapfile -t COMPREPLY < <(compgen -W \"$({bin_name} _servers 2>/dev/null)\" -- \"$cur\")\n\
64             \t\treturn\n\
65             \tfi\n\
66             \t_{fn_name} \"$@\"\n\
67             }}\n\
68             complete -F __{fn_name}_complete {bin_name}\n"
69        ),
70        _ => String::new(),
71    };
72
73    if !dynamic.is_empty() {
74        out.write_all(dynamic.as_bytes()).ok();
75    }
76}
77
78#[cfg(any(
79    feature = "technitium",
80    feature = "pangolin",
81    feature = "cloudflare",
82    feature = "unifi",
83    feature = "pihole"
84))]
85fn patch_zsh_server_completion(raw: &str, fn_name: &str) -> String {
86    let helper = format!(":_{fn_name}_server_ids'");
87    let patched: String = raw
88        .lines()
89        .map(|line| {
90            if line.contains("'--server=[") || line.contains("'*--server=[") {
91                line.replace(":_default'", &helper)
92            } else {
93                line.to_string()
94            }
95        })
96        .collect::<Vec<_>>()
97        .join("\n");
98
99    if raw.ends_with('\n') {
100        patched + "\n"
101    } else {
102        patched
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::patch_zsh_server_completion;
109
110    #[test]
111    fn zsh_patch_updates_repeatable_and_single_server_options_only() {
112        let raw = "\
113'*--server=[DNS server ID from the config file]:SERVERS:_default' \\
114'--server=[A configured server entry to query]:SERVER:_default' \\
115'--server-name=[TLS SNI server name]:SERVER_NAME:_default' \\
116";
117
118        let patched = patch_zsh_server_completion(raw, "dns");
119
120        assert!(
121            patched.contains(
122                "'*--server=[DNS server ID from the config file]:SERVERS:_dns_server_ids'"
123            )
124        );
125        assert!(
126            patched
127                .contains("'--server=[A configured server entry to query]:SERVER:_dns_server_ids'")
128        );
129        assert!(patched.contains("'--server-name=[TLS SNI server name]:SERVER_NAME:_default'"));
130        assert!(patched.ends_with('\n'));
131    }
132}
dnslib/cli/completions.rs

dnslib/cli/
completions.rs
