Skip to main content

lean_ctx/hook_handlers/
search_rewrite.rs

1//! Search and directory-listing command rewriting.
2//!
3//! Extracted from `hook_handlers/mod.rs` (#660 LOC gate) to keep the main
4//! module under the 1500-line budget. All `rewrite_*` functions plus shared
5//! helpers (`shell_tokenize`, `shell_quote`) live here.
6
7use super::file_rewrite::is_outside_project_path;
8
9/// Rewrites `grep`/`egrep`/`fgrep`/`rg` (and PowerShell `Select-String`/`sls`,
10/// #561) to `lean-ctx grep <pattern> [path]` when the invocation is simple enough
11/// to map losslessly; complex flag combos fall through to the `lean-ctx -c` wrap.
12pub(super) fn rewrite_search_command(cmd: &str, binary: &str) -> Option<String> {
13    let parts = shell_tokenize(cmd);
14    #[allow(clippy::match_same_arms)]
15    match parts.first().map(String::as_str) {
16        // fgrep uses fixed-string matching; lean-ctx grep is regex-only → always -c wrap
17        Some("fgrep") => None,
18        Some("grep" | "egrep") => rewrite_grep(&parts, binary),
19        Some("rg") => rewrite_rg(&parts, binary),
20        Some("Select-String" | "sls") => rewrite_select_string(&parts, binary),
21        _ => None,
22    }
23}
24
25/// Flags that are purely cosmetic and safe to strip: lean-ctx grep always shows
26/// line numbers, filenames, and searches recursively. Flags that change search
27/// semantics (-i, -w, -l, -c, -F, --include/--exclude) are NOT here — they
28/// cause fall-through to `lean-ctx -c` for correct native grep behavior.
29const GREP_SAFE_FLAGS: &[&str] = &[
30    "-n",
31    "--line-number",
32    "-r",
33    "-R",
34    "--recursive",
35    "-H",
36    "--with-filename",
37    "-s",
38    "--no-messages",
39    "--color=auto",
40    "--color=always",
41    "--color=never",
42    "--color",
43];
44
45/// Flags that take a value argument (next token is consumed as value).
46/// Only context flags are here — they cause fall-through to `-c` wrap.
47/// All other value-carrying flags (--include, -m, etc.) are unknown → fall-through.
48const GREP_VALUE_FLAGS: &[&str] = &[
49    "-A",
50    "--after-context",
51    "-B",
52    "--before-context",
53    "-C",
54    "--context",
55];
56
57/// Rewrites `grep [-nirlcwHRs] [--include=...] <pattern> [path...]` to
58/// `lean-ctx grep <pattern> [path]`. Complex invocations (pipes as stdin,
59/// unsupported flags, multiple paths) fall through to the `lean-ctx -c` wrap
60/// via the `is_rewritable` fallback.
61fn rewrite_grep(parts: &[String], binary: &str) -> Option<String> {
62    let mut pattern: Option<String> = None;
63    let mut path: Option<String> = None;
64    let mut has_context_flags = false;
65    let mut i = 1;
66
67    while i < parts.len() {
68        let arg = &parts[i];
69
70        if arg == "--" {
71            i += 1;
72            continue;
73        }
74
75        // Combined short flags like -rn: validate each char is safe to strip
76        if arg.starts_with('-') && !arg.starts_with("--") && arg.len() > 2 {
77            let chars = &arg[1..];
78            if chars.chars().all(|c| "nrRHs".contains(c)) {
79                i += 1;
80                continue;
81            }
82            return None;
83        }
84
85        // --flag=value style
86        if arg.starts_with("--") && arg.contains('=') {
87            let flag_name = arg.split('=').next().unwrap_or("");
88            if GREP_SAFE_FLAGS.contains(&flag_name) || GREP_VALUE_FLAGS.contains(&flag_name) {
89                i += 1;
90                continue;
91            }
92            return None;
93        }
94
95        // Known flags (with or without value)
96        if arg.starts_with('-') {
97            if GREP_VALUE_FLAGS.contains(&arg.as_str()) {
98                has_context_flags |= matches!(
99                    arg.as_str(),
100                    "-A" | "-B" | "-C" | "--after-context" | "--before-context" | "--context"
101                );
102                i += 2;
103                continue;
104            }
105            if GREP_SAFE_FLAGS.contains(&arg.as_str()) {
106                i += 1;
107                continue;
108            }
109            return None;
110        }
111
112        if pattern.is_none() {
113            pattern = Some(arg.clone());
114        } else if path.is_none() {
115            path = Some(arg.clone());
116        } else {
117            return None;
118        }
119        i += 1;
120    }
121
122    let pattern = pattern?;
123
124    if has_context_flags {
125        return None;
126    }
127
128    match &path {
129        Some(p) if is_outside_project_path(p) => None,
130        Some(p) => Some(format!(
131            "{binary} grep {} {}",
132            shell_quote(&pattern),
133            shell_quote(p)
134        )),
135        None => Some(format!("{binary} grep {}", shell_quote(&pattern))),
136    }
137}
138
139/// Rewrites `rg [flags] <pattern> [path]` to `lean-ctx grep <pattern> [path]`.
140/// Supports common flags that don't alter the fundamental search semantics.
141fn rewrite_rg(parts: &[String], binary: &str) -> Option<String> {
142    if parts.len() < 2 {
143        return None;
144    }
145
146    const RG_SAFE_SHORT: &str = "nsSHu";
147    const RG_SAFE_LONG: &[&str] = &[
148        "--line-number",
149        "--no-ignore",
150        "--hidden",
151        "--no-heading",
152        "--with-filename",
153        "--follow",
154        "--unrestricted",
155        "--color=auto",
156        "--color=always",
157        "--color=never",
158        "--color=ansi",
159        "--no-line-number",
160    ];
161    const RG_VALUE_FLAGS: &[&str] = &[
162        "-A",
163        "--after-context",
164        "-B",
165        "--before-context",
166        "-C",
167        "--context",
168    ];
169
170    let mut pattern: Option<String> = None;
171    let mut path: Option<String> = None;
172    let mut has_context_flags = false;
173    let mut i = 1;
174
175    while i < parts.len() {
176        let arg = &parts[i];
177
178        if arg == "--" {
179            i += 1;
180            continue;
181        }
182
183        if arg.starts_with("--") && arg.contains('=') {
184            let flag_name = arg.split('=').next().unwrap_or("");
185            if RG_SAFE_LONG.contains(&arg.as_str())
186                || RG_SAFE_LONG.contains(&flag_name)
187                || RG_VALUE_FLAGS.contains(&flag_name)
188            {
189                i += 1;
190                continue;
191            }
192            return None;
193        }
194
195        if arg.starts_with("--") {
196            if RG_SAFE_LONG.contains(&arg.as_str()) {
197                i += 1;
198                continue;
199            }
200            if RG_VALUE_FLAGS.contains(&arg.as_str()) {
201                has_context_flags |= matches!(
202                    arg.as_str(),
203                    "--after-context" | "--before-context" | "--context"
204                );
205                i += 2;
206                continue;
207            }
208            return None;
209        }
210
211        if arg.starts_with('-') && arg.len() >= 2 {
212            let flag_str = &arg[..2];
213            if RG_VALUE_FLAGS.contains(&flag_str) {
214                has_context_flags |= matches!(flag_str, "-A" | "-B" | "-C");
215                if arg.len() > 2 {
216                    i += 1;
217                } else {
218                    i += 2;
219                }
220                continue;
221            }
222            let chars = &arg[1..];
223            if chars.chars().all(|c| RG_SAFE_SHORT.contains(c)) {
224                i += 1;
225                continue;
226            }
227            return None;
228        }
229
230        if pattern.is_none() {
231            pattern = Some(arg.clone());
232        } else if path.is_none() {
233            path = Some(arg.clone());
234        } else {
235            return None;
236        }
237        i += 1;
238    }
239
240    let pattern = pattern?;
241
242    if has_context_flags {
243        return None;
244    }
245
246    match &path {
247        Some(p) if is_outside_project_path(p) => None,
248        Some(p) => Some(format!(
249            "{binary} grep {} {}",
250            shell_quote(&pattern),
251            shell_quote(p)
252        )),
253        None => Some(format!("{binary} grep {}", shell_quote(&pattern))),
254    }
255}
256
257/// Maps `Select-String`/`sls` to `lean-ctx grep`, honoring `-Pattern` and
258/// `-Path`/`-LiteralPath` plus the positional `<pattern> [path]` form.
259fn rewrite_select_string(parts: &[String], binary: &str) -> Option<String> {
260    let mut pattern: Option<String> = None;
261    let mut path: Option<String> = None;
262    let mut i = 1;
263    while i < parts.len() {
264        if let Some(flag) = parts[i].strip_prefix('-') {
265            let value = parts.get(i + 1);
266            match flag.to_ascii_lowercase().as_str() {
267                "pattern" => pattern = Some(value?.clone()),
268                "path" | "literalpath" => path = Some(value?.clone()),
269                _ => return None,
270            }
271            i += 2;
272        } else if pattern.is_none() {
273            pattern = Some(parts[i].clone());
274            i += 1;
275        } else if path.is_none() {
276            path = Some(parts[i].clone());
277            i += 1;
278        } else {
279            return None;
280        }
281    }
282    let pattern = shell_quote(&pattern?);
283    match path {
284        Some(p) if is_outside_project_path(&p) => None,
285        Some(p) => Some(format!("{binary} grep {pattern} {}", shell_quote(&p))),
286        None => Some(format!("{binary} grep {pattern}")),
287    }
288}
289
290/// Rewrites simple `ls [path]` (and PowerShell `Get-ChildItem`/`gci`, #561) to
291/// `lean-ctx ls [path]`.
292pub(super) fn rewrite_dir_list_command(cmd: &str, binary: &str) -> Option<String> {
293    let parts = shell_tokenize(cmd);
294    match parts.first().map(String::as_str) {
295        Some("ls") => match parts.len() {
296            1 => Some(format!("{binary} ls")),
297            2 if !parts[1].starts_with('-') => {
298                Some(format!("{binary} ls {}", shell_quote(&parts[1])))
299            }
300            _ => None,
301        },
302        Some("Get-ChildItem" | "gci") => rewrite_get_childitem(&parts, binary),
303        _ => None,
304    }
305}
306
307fn rewrite_get_childitem(parts: &[String], binary: &str) -> Option<String> {
308    let mut path: Option<String> = None;
309    let mut i = 1;
310    while i < parts.len() {
311        if let Some(flag) = parts[i].strip_prefix('-') {
312            let value = parts.get(i + 1);
313            match flag.to_ascii_lowercase().as_str() {
314                "path" | "literalpath" => path = Some(value?.clone()),
315                _ => return None,
316            }
317            i += 2;
318        } else if path.is_none() {
319            path = Some(parts[i].clone());
320            i += 1;
321        } else {
322            return None;
323        }
324    }
325    match path {
326        Some(p) => Some(format!("{binary} ls {}", shell_quote(&p))),
327        None => Some(format!("{binary} ls")),
328    }
329}
330
331/// Tokenize a shell command respecting single/double quotes and backslash escapes.
332pub fn shell_tokenize(input: &str) -> Vec<String> {
333    let mut tokens = Vec::new();
334    let mut current = String::new();
335    let mut chars = input.chars().peekable();
336    let mut in_single = false;
337    let mut in_double = false;
338
339    while let Some(c) = chars.next() {
340        match c {
341            '\'' if !in_double => in_single = !in_single,
342            '"' if !in_single => in_double = !in_double,
343            '\\' if !in_single => {
344                if let Some(next) = chars.next() {
345                    current.push(next);
346                }
347            }
348            c if c.is_whitespace() && !in_single && !in_double => {
349                if !current.is_empty() {
350                    tokens.push(std::mem::take(&mut current));
351                }
352            }
353            _ => current.push(c),
354        }
355    }
356    if !current.is_empty() {
357        tokens.push(current);
358    }
359    tokens
360}
361
362/// Quote a path/arg for shell if it contains spaces or special chars.
363pub fn shell_quote(s: &str) -> String {
364    if s.contains(|c: char| {
365        c.is_whitespace()
366            || matches!(
367                c,
368                '\'' | '"'
369                    | '\\'
370                    | '|'
371                    | '&'
372                    | ';'
373                    | '$'
374                    | '`'
375                    | '('
376                    | ')'
377                    | '*'
378                    | '?'
379                    | '>'
380                    | '<'
381                    | '#'
382                    | '!'
383                    | '['
384                    | ']'
385                    | '{'
386                    | '}'
387                    | '~'
388            )
389    }) {
390        format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\""))
391    } else {
392        s.to_string()
393    }
394}