From 2b55068fbfd8d9e1fd52541b4f10d9c765cceed8 Mon Sep 17 00:00:00 2001 From: xuezhixin Date: Thu, 21 May 2026 17:25:00 +0800 Subject: [PATCH] feat: add secret detection with three-option dialog and vault redaction Add regex-based secret scanner (API keys, JWTs, passwords) and an in-memory SecretVault that redacts detected secrets to semantic placeholders ($SECRET_*) before sending to the LLM, then restores them for command execution and re-redacts command output. - SecretScanner with RegexSet for efficient multi-pattern matching - SecretVault with redact/restore/redact_output round-trip - Three-option dialog: redact to env vars, send plaintext, or abort - Both TUI (inquire) and SSH (raw terminal I/O) dialog paths - BashTool integration: restore before execution, redact output before returning to AI, restore in preflight for accurate security checks --- Cargo.lock | 3 + config/security_policy.yaml | 10 + crates/aish-i18n/locales/en-US.yaml | 10 + crates/aish-i18n/locales/zh-CN.yaml | 10 + crates/aish-pty/Cargo.toml | 1 + crates/aish-pty/src/lib.rs | 8 + crates/aish-pty/src/persistent.rs | 279 +++++++++++- crates/aish-pty/src/session_interceptor.rs | 16 +- crates/aish-security/Cargo.toml | 2 + crates/aish-security/src/decision.rs | 24 + crates/aish-security/src/lib.rs | 1 + crates/aish-security/src/manager.rs | 62 +++ crates/aish-security/src/policy.rs | 30 +- crates/aish-security/src/sandbox/degraded.rs | 1 + crates/aish-security/src/secret/mod.rs | 7 + crates/aish-security/src/secret/patterns.rs | 172 ++++++++ crates/aish-security/src/secret/scanner.rs | 199 +++++++++ crates/aish-security/src/secret/vault.rs | 442 +++++++++++++++++++ crates/aish-shell/src/app.rs | 122 ++++- crates/aish-shell/src/tui.rs | 45 +- crates/aish-tools/src/bash.rs | 71 ++- 21 files changed, 1486 insertions(+), 29 deletions(-) create mode 100644 crates/aish-security/src/secret/mod.rs create mode 100644 crates/aish-security/src/secret/patterns.rs create mode 100644 crates/aish-security/src/secret/scanner.rs create mode 100644 crates/aish-security/src/secret/vault.rs diff --git a/Cargo.lock b/Cargo.lock index 70c43ba..f0f56c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -167,6 +167,7 @@ dependencies = [ "aish-core", "aish-hosts", "aish-i18n", + "aish-security", "base64 0.22.1", "chrono", "futures", @@ -198,11 +199,13 @@ name = "aish-security" version = "0.3.0-beta.3" dependencies = [ "libc", + "regex", "serde", "serde_json", "serde_yaml", "tempfile", "thiserror 2.0.18", + "tracing", ] [[package]] diff --git a/config/security_policy.yaml b/config/security_policy.yaml index 78065c4..692081d 100644 --- a/config/security_policy.yaml +++ b/config/security_policy.yaml @@ -80,3 +80,13 @@ rules: risk: LOW reason: "临时区代码和项目文件，允许 AI 修改" +# Secret detection patterns (optional). +# Uncomment and customize to add your own patterns. +# Built-in patterns cover: OpenAI, Anthropic, AWS, Google, GitHub, Stripe, +# Slack, Fireworks API keys; JWT tokens; URL-embedded passwords. +# +# secret_patterns: +# - name: "My Internal Token" +# pattern: "int_[a-zA-Z0-9]{32}" +# secret_type: api_key + diff --git a/crates/aish-i18n/locales/en-US.yaml b/crates/aish-i18n/locales/en-US.yaml index c939ccc..2899b0a 100644 --- a/crates/aish-i18n/locales/en-US.yaml +++ b/crates/aish-i18n/locales/en-US.yaml @@ -382,6 +382,16 @@ shell: alternative: generic: "Please manually review this command and verify impact before execution." etc_manual: "If you must modify files under /etc, prefer a manual change process with backups/change management." + secret: + title: "Security Warning" + detected: "Detected sensitive data in input:\n {reasons}\n\nChoose how to proceed:" + redact: "Replace with env vars (safe)" + allow: "Send plaintext to AI" + abort: "Cancel" + aborted: "Aborted: sensitive data detected." + redacted: "Replaced {count} secret(s) with env var placeholders (stored in memory only, not sent to AI)" + restored: "(Restored {count} env var(s) from in-memory vault, not sent to AI)" + option_esc: "[Enter to confirm, ↑↓ to select, Esc to cancel]" options_header: "Options" option: approve: "Approve" diff --git a/crates/aish-i18n/locales/zh-CN.yaml b/crates/aish-i18n/locales/zh-CN.yaml index 88f7273..164eb0a 100644 --- a/crates/aish-i18n/locales/zh-CN.yaml +++ b/crates/aish-i18n/locales/zh-CN.yaml @@ -382,6 +382,16 @@ shell: alternative: generic: "请人工复核该命令并在确认影响后再执行。" etc_manual: "如需修改 /etc 下文件，建议采用人工变更并配合变更管理或备份策略。" + secret: + title: "安全警告" + detected: "检测到输入中包含敏感数据：\n {reasons}\n\n请选择处理方式：" + redact: "替换为环境变量发送 (安全)" + allow: "明文发送给 AI" + abort: "取消" + aborted: "已取消：检测到敏感数据。" + redacted: "已将 {count} 个敏感数据替换为环境变量占位符（仅保存在内存中，未发送给 AI）" + restored: "（已还原 {count} 个环境变量，仅保存在内存中，未发送给 AI）" + option_esc: "[Enter 确认, ↑↓ 选择, Esc 取消]" options_header: "选项" option: approve: "允许" diff --git a/crates/aish-pty/Cargo.toml b/crates/aish-pty/Cargo.toml index 7c7a35e..bee136a 100644 --- a/crates/aish-pty/Cargo.toml +++ b/crates/aish-pty/Cargo.toml @@ -7,6 +7,7 @@ edition.workspace = true aish-core.workspace = true aish-hosts.workspace = true aish-i18n.workspace = true +aish-security.workspace = true nix.workspace = true tokio.workspace = true serde.workspace = true diff --git a/crates/aish-pty/src/lib.rs b/crates/aish-pty/src/lib.rs index 6d0ea92..6074a55 100644 --- a/crates/aish-pty/src/lib.rs +++ b/crates/aish-pty/src/lib.rs @@ -24,6 +24,14 @@ pub mod session_interceptor; pub mod state_capture; pub mod types; +/// Result returned by the SSH secret-check closure. +pub struct SshSecretCheckResult { + /// Formatted warning message (title + detected secrets). + pub warning: String, + /// Detected secret matches for vault redaction. + pub detected_secrets: Vec, +} + pub use command_state::CommandState; pub use control::{decode_control_chunk, encode_control_event, BackendControlEvent}; pub use executor::PtyExecutor; diff --git a/crates/aish-pty/src/persistent.rs b/crates/aish-pty/src/persistent.rs index 4db486d..d33d3ea 100644 --- a/crates/aish-pty/src/persistent.rs +++ b/crates/aish-pty/src/persistent.rs @@ -417,6 +417,10 @@ impl PersistentPty { command: &str, ai_callback: Option>, shared_host: Option>>>, + secret_check: Option< + std::sync::Arc Option + Send + Sync>, + >, + secret_vault: Option>>, ) -> aish_core::Result<(i32, String, String)> { let is_session = is_session_command(command); let mut interceptor = if is_session { @@ -653,7 +657,7 @@ impl PersistentPty { if probe_for_ai { probe_for_ai = false; if let Some(question) = pending_ai_question.take() { - let resp = interceptor.call_ai(question); + let resp = interceptor.call_ai(question, secret_vault.as_ref()); interceptor.finish_ai(); if let Some(response) = resp { pending_response = Some(response); @@ -906,6 +910,29 @@ impl PersistentPty { _ => false, }; if approved { + // Restore secret placeholders in the AI-generated command + let mut cmd_restored = cmd.clone(); + if let Some(ref vault) = secret_vault { + let vault_guard = vault.lock().unwrap(); + let (restored, count) = vault_guard.restore(cmd); + if count > 0 { + let mut rargs = std::collections::HashMap::new(); + rargs.insert("count".to_string(), count.to_string()); + let msg = aish_i18n::t_with_args( + "shell.security.secret.restored", + &rargs, + ); + let info = format!("\x1b[2m{}\x1b[0m\r\n", msg); + unsafe { + libc::write( + libc::STDOUT_FILENO, + info.as_ptr() as *const libc::c_void, + info.len(), + ); + } + cmd_restored = restored; + } + } // Show "Running..." feedback let running_msg = format!( "\x1b[90m{}\x1b[0m\r\n", @@ -918,7 +945,7 @@ impl PersistentPty { running_msg.len(), ); } - let safe_cmd = close_unclosed_heredoc(cmd); + let safe_cmd = close_unclosed_heredoc(&cmd_restored); skip_echo_cmd = Some(safe_cmd.clone()); let mut inject = safe_cmd.as_bytes().to_vec(); inject.push(b'\r'); @@ -1172,7 +1199,7 @@ impl PersistentPty { 1, ); }, - crate::StdinAction::TriggerAi(question) => { + crate::StdinAction::TriggerAi(mut question) => { // Reset hard-cancel flag for a fresh AI session ai_cancelled = false; // Clear stdin shadow — the AI prefix @@ -1244,6 +1271,80 @@ impl PersistentPty { ); } + // Security gate: secret detection in AI input + if let Some(ref checker) = secret_check { + if let Some(result) = checker(&question) { + let choice = show_secret_dialog( + &result.warning, + libc::STDIN_FILENO, + ); + match choice { + SshSecretChoice::Abort => { + let abort_msg = format!( + "\x1b[33m{}\x1b[0m\r\n", + aish_i18n::t( + "shell.security.secret.aborted" + ) + ); + unsafe { + libc::write( + libc::STDOUT_FILENO, + abort_msg.as_ptr() + as *const libc::c_void, + abort_msg.len(), + ); + libc::write( + self.master_fd, + b"\r".as_ptr() as *const libc::c_void, + 1, + ); + } + interceptor.finish_ai(); + continue; + } + SshSecretChoice::Redact => { + if let Some(ref vault) = secret_vault { + let mut vault_guard = vault.lock().unwrap(); + let redacted = vault_guard.redact( + &result.detected_secrets, + &question, + ); + let count = result.detected_secrets.len(); + let mut rargs = + std::collections::HashMap::new(); + rargs.insert( + "count".to_string(), + count.to_string(), + ); + let msg = aish_i18n::t_with_args( + "shell.security.secret.redacted", + &rargs, + ); + let info = + format!("\x1b[33m{}\x1b[0m\r\n", msg); + unsafe { + libc::write( + libc::STDOUT_FILENO, + info.as_ptr() + as *const libc::c_void, + info.len(), + ); + } + // Replace question with redacted version + question = redacted; + } else { + // Cannot honor "Redact" safely without vault support. + interceptor.finish_ai(); + continue; + } + } + SshSecretChoice::Allow => { + // Proceed with original question unchanged + } + } + } + } + // Check for dossier commands before invoking AI let question_trimmed = question.trim().to_string(); let dossier_result = handle_dossier_command( @@ -1324,7 +1425,7 @@ impl PersistentPty { continue; } else { // No probe needed, call AI now. - interceptor.call_ai(question) + interceptor.call_ai(question, secret_vault.as_ref()) } }; interceptor.finish_ai(); @@ -1494,7 +1595,8 @@ impl PersistentPty { if probe_for_ai { probe_for_ai = false; if let Some(question) = pending_ai_question.take() { - let resp = interceptor.call_ai(question); + let resp = interceptor + .call_ai(question, secret_vault.as_ref()); interceptor.finish_ai(); if let Some(response) = resp { pending_response = Some(response); @@ -1831,7 +1933,29 @@ impl PersistentPty { running_msg.len(), ); } - let safe_cmd = close_unclosed_heredoc(cmd); + // Restore secret placeholders before execution. + let mut cmd_restored = cmd.clone(); + if let Some(ref vault) = secret_vault { + let (restored, count) = vault.lock().unwrap().restore(cmd); + if count > 0 { + let mut rargs = std::collections::HashMap::new(); + rargs.insert("count".to_string(), count.to_string()); + let msg = aish_i18n::t_with_args( + "shell.security.secret.restored", + &rargs, + ); + let info = format!("\x1b[2m{}\x1b[0m\r\n", msg); + unsafe { + libc::write( + libc::STDOUT_FILENO, + info.as_ptr() as *const libc::c_void, + info.len(), + ); + } + cmd_restored = restored; + } + } + let safe_cmd = close_unclosed_heredoc(&cmd_restored); skip_echo_cmd = Some(safe_cmd.clone()); let mut inject = safe_cmd.as_bytes().to_vec(); inject.push(b'\r'); @@ -2276,6 +2400,143 @@ impl PersistentPty { } } +// ---- secret detection confirmation for SSH sessions ---- + +/// Render the warning + options + help into a byte buffer. +/// Returns the buffer and the exact number of `\r\n` sequences it contains, +/// which equals the number of cursor-up moves needed to return to the first row. +fn render_secret_confirmation(warning: &str, options: &[&str], cursor: usize) -> (Vec, usize) { + let mut out = Vec::new(); + + // Warning message — replace \n with \r\n for terminal correctness + let warning_display = warning.replace('\n', "\r\n"); + out.extend_from_slice(warning_display.as_bytes()); + out.extend_from_slice(b"\r\n"); + + // Options with cursor highlight (inquire-style) + for (i, opt) in options.iter().enumerate() { + if i == cursor { + out.extend_from_slice(b"\x1b[36m> \x1b[1m"); + } else { + out.extend_from_slice(b" "); + } + out.extend_from_slice(opt.as_bytes()); + out.extend_from_slice(b"\x1b[0m\r\n"); + } + + // Help line (no trailing \r\n — cursor stays on this row) + let esc_hint = aish_i18n::t("shell.security.secret.option_esc"); + out.extend_from_slice(b"\x1b[2m"); + out.extend_from_slice(esc_hint.as_bytes()); + out.extend_from_slice(b"\x1b[0m"); + + // Count \r\n occurrences = exact up-movement needed from last row to first row + let up_moves = out.windows(2).filter(|w| w == b"\r\n").count(); + + (out, up_moves) +} + +/// Choice for SSH secret dialog. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SshSecretChoice { + Redact, + Allow, + Abort, +} + +/// Display a three-option secret dialog with up/down selection (SSH path). +fn show_secret_dialog(warning: &str, stdin_fd: libc::c_int) -> SshSecretChoice { + let option_redact = aish_i18n::t("shell.security.secret.redact"); + let option_allow = aish_i18n::t("shell.security.secret.allow"); + let option_abort = aish_i18n::t("shell.security.secret.abort"); + let options = [ + option_redact.as_str(), + option_allow.as_str(), + option_abort.as_str(), + ]; + let choice_values = [ + SshSecretChoice::Redact, + SshSecretChoice::Allow, + SshSecretChoice::Abort, + ]; + // Default cursor on "Redact" (index 0, safest) + let mut cursor: usize = 0; + + let (buf, up_moves) = render_secret_confirmation(warning, &options, cursor); + unsafe { + libc::write( + libc::STDOUT_FILENO, + buf.as_ptr() as *const libc::c_void, + buf.len(), + ); + } + + loop { + match read_byte(stdin_fd) { + Some(byte) => match byte { + 0x03 => { + unsafe { + libc::write( + libc::STDOUT_FILENO, + b"^C\r\n".as_ptr() as *const _, + b"^C\r\n".len(), + ); + } + return SshSecretChoice::Abort; + } + 0x1B => { + if stdin_poll(stdin_fd, 50_000) { + let next = read_byte(stdin_fd); + if next == Some(b'[') { + if let Some(final_byte) = consume_csi(stdin_fd) { + match final_byte { + b'A' => { + if cursor > 0 { + cursor -= 1; + } + } + b'B' => { + if cursor < options.len() - 1 { + cursor += 1; + } + } + _ => {} + } + let (buf, up) = + render_secret_confirmation(warning, &options, cursor); + let mut redraw = Vec::new(); + redraw + .extend_from_slice(format!("\x1b[{}A\r\x1b[J", up).as_bytes()); + redraw.extend_from_slice(&buf); + unsafe { + libc::write( + libc::STDOUT_FILENO, + redraw.as_ptr() as *const libc::c_void, + redraw.len(), + ); + } + } + continue; + } + } + unsafe { + libc::write(libc::STDOUT_FILENO, b"\r\n".as_ptr() as *const _, 2); + } + return SshSecretChoice::Abort; + } + b'\r' | b'\n' => { + unsafe { + libc::write(libc::STDOUT_FILENO, b"\r\n".as_ptr() as *const _, 2); + } + return choice_values[cursor]; + } + _ => {} + }, + None => return SshSecretChoice::Abort, + } + } +} + // ---- ask_user helpers for SSH sessions ---- /// Drain trailing bytes (e.g. `\n` or `\r`) from stdin after a single-byte @@ -2544,7 +2805,11 @@ fn read_line_from_stdin_raw( // Ctrl+C → always cancel 0x03 => { unsafe { - libc::write(libc::STDOUT_FILENO, b"^C\r\n".as_ptr() as *const _, 5); + libc::write( + libc::STDOUT_FILENO, + b"^C\r\n".as_ptr() as *const _, + b"^C\r\n".len(), + ); } return crate::AskUserAnswer::Cancelled; } diff --git a/crates/aish-pty/src/session_interceptor.rs b/crates/aish-pty/src/session_interceptor.rs index d2a218c..973c028 100644 --- a/crates/aish-pty/src/session_interceptor.rs +++ b/crates/aish-pty/src/session_interceptor.rs @@ -323,9 +323,17 @@ impl SessionInterceptor { /// Run the AI callback. The callback returns an AiResponse containing /// an optional command and display text (or None on error). - pub fn call_ai(&self, question: String) -> Option { + pub fn call_ai( + &self, + question: String, + secret_vault: Option<&std::sync::Arc>>, + ) -> Option { self.ai_callback.as_ref().and_then(|cb| { - let recent = self.recent_output(4000); + let mut recent = self.recent_output(4000); + if let Some(vault) = secret_vault { + let (redacted, _) = vault.lock().unwrap().redact_output(&recent); + recent = redacted; + } cb(AiQuery { question, recent_output: recent, @@ -564,7 +572,7 @@ mod tests { let mut ic = SessionInterceptor::new(Some(noop_callback())); ic.feed_stdin(b';'); ic.feed_stdin(b'\r'); - let resp = ic.call_ai("test".to_string()); + let resp = ic.call_ai("test".to_string(), None); assert!(resp.is_some()); let r = resp.unwrap(); assert_eq!(r.command, Some("echo test".to_string())); @@ -573,7 +581,7 @@ mod tests { #[test] fn test_call_ai_returns_none() { let ic = SessionInterceptor::new(Some(noop_callback_no_cmd())); - let cmd = ic.call_ai("test".to_string()); + let cmd = ic.call_ai("test".to_string(), None); assert!(cmd.is_none()); } diff --git a/crates/aish-security/Cargo.toml b/crates/aish-security/Cargo.toml index 49cb4c3..683fb47 100644 --- a/crates/aish-security/Cargo.toml +++ b/crates/aish-security/Cargo.toml @@ -9,6 +9,8 @@ serde.workspace = true serde_json.workspace = true serde_yaml.workspace = true thiserror.workspace = true +regex.workspace = true +tracing.workspace = true [dev-dependencies] tempfile.workspace = true diff --git a/crates/aish-security/src/decision.rs b/crates/aish-security/src/decision.rs index 1928593..9a45921 100644 --- a/crates/aish-security/src/decision.rs +++ b/crates/aish-security/src/decision.rs @@ -96,6 +96,8 @@ pub struct SecurityAnalysis { pub sandbox_off_action: Option, #[serde(default)] pub sandbox: SandboxStatus, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub detected_secrets: Option>, } impl Default for SecurityAnalysis { @@ -113,6 +115,7 @@ impl Default for SecurityAnalysis { matched_paths: Vec::new(), sandbox_off_action: None, sandbox: SandboxStatus::default(), + detected_secrets: None, } } } @@ -196,4 +199,25 @@ mod tests { assert_eq!(SandboxStatus::default().error, None); assert_eq!(SandboxStatus::default().exit_code, None); } + + #[test] + fn security_analysis_default_has_no_detected_secrets() { + let analysis = SecurityAnalysis::default(); + assert!(analysis.detected_secrets.is_none()); + } + + #[test] + fn security_analysis_serializes_detected_secrets_roundtrip() { + let mut analysis = SecurityAnalysis::default(); + analysis.detected_secrets = Some(vec![crate::secret::SecretMatch { + pattern_name: "Test Key".to_string(), + start: 5, + end: 20, + secret_type: crate::secret::SecretType::ApiKey, + }]); + let json = serde_json::to_string(&analysis).unwrap(); + let deserialized: SecurityAnalysis = serde_json::from_str(&json).unwrap(); + assert!(deserialized.detected_secrets.is_some()); + assert_eq!(deserialized.detected_secrets.unwrap().len(), 1); + } } diff --git a/crates/aish-security/src/lib.rs b/crates/aish-security/src/lib.rs index a04ae73..076b56b 100644 --- a/crates/aish-security/src/lib.rs +++ b/crates/aish-security/src/lib.rs @@ -4,6 +4,7 @@ pub mod manager; pub mod policy; pub mod risk; mod sandbox; +pub mod secret; pub mod sudo; pub mod types; diff --git a/crates/aish-security/src/manager.rs b/crates/aish-security/src/manager.rs index 3952225..51b2fc1 100644 --- a/crates/aish-security/src/manager.rs +++ b/crates/aish-security/src/manager.rs @@ -11,6 +11,7 @@ use crate::sandbox::degraded::{analyze_sandbox_degraded, SandboxDegradedDetails} use crate::sandbox::error::{SandboxError, SandboxReason}; use crate::sandbox::ipc::client::SandboxRunner; use crate::sandbox::types::SandboxRunRequest; +use crate::secret::SecretScanner; use crate::SandboxClient; const DEFAULT_SANDBOX_SOCKET_PATH: &str = "/run/aish/sandbox.sock"; @@ -87,6 +88,7 @@ pub struct SecurityManager { policy: SecurityPolicy, fallback_engine: FallbackRuleEngine, sandbox_runner: Option>, + secret_scanner: SecretScanner, } impl SecurityManager { @@ -95,10 +97,12 @@ impl SecurityManager { let sandbox_runner = policy.enable_sandbox.then(|| { Arc::new(SandboxClient::new(DEFAULT_SANDBOX_SOCKET_PATH)) as Arc }); + let secret_scanner = SecretScanner::new(&policy.secret_patterns); Self { policy, fallback_engine, sandbox_runner, + secret_scanner, } } @@ -106,6 +110,10 @@ impl SecurityManager { &self.policy } + pub fn secret_scanner(&self) -> &SecretScanner { + &self.secret_scanner + } + pub fn with_sandbox_client(mut self, client: SandboxClient) -> Self { self.sandbox_runner = Some(Arc::new(client)); self @@ -159,6 +167,23 @@ impl SecurityManager { decision_from_risk(level, analysis) } + /// Check AI input text for sensitive data. + /// Returns `Confirm` if secrets are detected, `Allow` otherwise. + pub fn check_ai_input(&self, input: &str) -> SecurityDecision { + let secrets = self.secret_scanner.scan(input); + if secrets.is_empty() { + return SecurityDecision::allow(RiskLevel::Low, SecurityAnalysis::default()); + } + + let analysis = SecurityAnalysis { + risk_level: RiskLevel::Medium, + reasons: secrets.iter().map(|s| s.format_reason()).collect(), + detected_secrets: Some(secrets), + ..SecurityAnalysis::default() + }; + SecurityDecision::confirm(RiskLevel::Medium, analysis) + } + fn analyze_with_sandbox(&self, command: &str, request: &SecurityRequest) -> SecurityAnalysis { let cwd = request .cwd() @@ -226,6 +251,7 @@ impl std::fmt::Debug for SecurityManager { .field("policy", &self.policy) .field("fallback_engine", &self.fallback_engine) .field("sandbox_runner", &self.sandbox_runner.is_some()) + .field("secret_scanner", &"") .finish() } } @@ -577,4 +603,40 @@ mod tests { Some("/repo") ); } + + #[test] + fn check_ai_input_allows_clean_text() { + let manager = SecurityManager::new(SecurityPolicy::default()); + let decision = manager.check_ai_input("list files in current directory"); + assert!(decision.allow); + assert!(!decision.require_confirmation); + } + + #[test] + fn check_ai_input_confirms_on_secret() { + let manager = SecurityManager::new(SecurityPolicy::default()); + let decision = manager.check_ai_input( + "use key sk-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ123456", + ); + assert!(decision.allow); + assert!(decision.require_confirmation); + let secrets = decision.analysis.detected_secrets.unwrap(); + assert!(!secrets.is_empty()); + assert_eq!(secrets[0].secret_type, crate::secret::SecretType::ApiKey); + } + + #[test] + fn check_ai_input_with_custom_pattern() { + let mut policy = SecurityPolicy::default(); + policy.secret_patterns = vec![crate::secret::CustomPattern { + name: "MyToken".to_string(), + pattern: r"mytoken_[a-z]{8}".to_string(), + secret_type: crate::secret::SecretType::Token, + }]; + let manager = SecurityManager::new(policy); + let decision = manager.check_ai_input("here is mytoken_abcdefgh ok"); + assert!(decision.require_confirmation); + let secrets = decision.analysis.detected_secrets.unwrap(); + assert_eq!(secrets[0].pattern_name, "MyToken"); + } } diff --git a/crates/aish-security/src/policy.rs b/crates/aish-security/src/policy.rs index 04c584d..552a6db 100644 --- a/crates/aish-security/src/policy.rs +++ b/crates/aish-security/src/policy.rs @@ -66,6 +66,8 @@ pub struct SecurityPolicy { pub invalid_fallback_rules: Vec, #[serde(default)] pub validation_issues: Vec, + #[serde(default)] + pub secret_patterns: Vec, } impl Default for SecurityPolicy { @@ -80,6 +82,7 @@ impl Default for SecurityPolicy { rules: default_rules(), invalid_fallback_rules: Vec::new(), validation_issues: Vec::new(), + secret_patterns: Vec::new(), } } } @@ -428,7 +431,7 @@ pub fn load_policy(config_path: Option<&Path>) -> SecurityPolicy { .filter(|item| mapping_get(item, "path").is_some()) .collect(); - let (invalid_fallback_rules, validation_issues) = parse_invalid_fallback_rules(&v2_items); + let (invalid_fallback_rules, mut validation_issues) = parse_invalid_fallback_rules(&v2_items); let valid_items: Vec<&Mapping> = v2_items .iter() @@ -448,6 +451,30 @@ pub fn load_policy(config_path: Option<&Path>) -> SecurityPolicy { let mut rules = default_rules(); rules.extend(parse_v2_rules(&valid_items)); + let secret_patterns: Vec = root + .and_then(|m| mapping_get(m, "secret_patterns")) + .and_then(Value::as_sequence) + .map(|seq| { + seq.iter() + .filter_map(|v| { + let yaml_str = serde_yaml::to_string(v).unwrap_or_default(); + match serde_yaml::from_value(v.clone()) { + Ok(pattern) => Some(pattern), + Err(e) => { + validation_issues.push(ValidationIssue { + rule_id: None, + field: "secret_patterns".to_string(), + value: Some(yaml_str.trim().to_string()), + message: Some(format!("Invalid custom pattern: {e}")), + }); + None + } + } + }) + .collect() + }) + .unwrap_or_default(); + SecurityPolicy { enable_sandbox, sandbox_off_action, @@ -458,6 +485,7 @@ pub fn load_policy(config_path: Option<&Path>) -> SecurityPolicy { rules, invalid_fallback_rules, validation_issues, + secret_patterns, } } diff --git a/crates/aish-security/src/sandbox/degraded.rs b/crates/aish-security/src/sandbox/degraded.rs index 93a8aa0..16adc3f 100644 --- a/crates/aish-security/src/sandbox/degraded.rs +++ b/crates/aish-security/src/sandbox/degraded.rs @@ -154,6 +154,7 @@ fn analysis_from_fallback( matched_paths: assessment.matched_paths.clone(), sandbox_off_action: Some(policy.sandbox_off_action), sandbox: sandbox_status(reason, details), + detected_secrets: None, } } diff --git a/crates/aish-security/src/secret/mod.rs b/crates/aish-security/src/secret/mod.rs new file mode 100644 index 0000000..9363e9e --- /dev/null +++ b/crates/aish-security/src/secret/mod.rs @@ -0,0 +1,7 @@ +mod patterns; +mod scanner; +mod vault; + +pub use patterns::{CustomPattern, SecretMatch, SecretPattern, SecretType}; +pub use scanner::SecretScanner; +pub use vault::SecretVault; diff --git a/crates/aish-security/src/secret/patterns.rs b/crates/aish-security/src/secret/patterns.rs new file mode 100644 index 0000000..584c448 --- /dev/null +++ b/crates/aish-security/src/secret/patterns.rs @@ -0,0 +1,172 @@ +use serde::{Deserialize, Serialize}; + +/// Category of a detected secret. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SecretType { + ApiKey, + Token, + Password, + Credential, +} + +impl Default for SecretType { + fn default() -> Self { + Self::ApiKey + } +} + +pub fn default_secret_type() -> SecretType { + SecretType::default() +} + +/// A single secret match found in input text. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SecretMatch { + /// Human-readable name, e.g. "Anthropic API Key". + pub pattern_name: String, + /// Byte offset where the match starts. + pub start: usize, + /// Byte offset where the match ends (exclusive). + pub end: usize, + /// Category of the secret. + pub secret_type: SecretType, +} + +impl SecretMatch { + pub fn format_reason(&self) -> String { + format!( + "{} detected at position {}..{}", + self.pattern_name, self.start, self.end + ) + } +} + +/// Internal representation of a compiled pattern. +pub struct SecretPattern { + pub name: &'static str, + pub regex: &'static str, + pub secret_type: SecretType, +} + +/// User-defined pattern from config. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CustomPattern { + pub name: String, + pub pattern: String, + #[serde(default = "default_secret_type")] + pub secret_type: SecretType, +} + +pub fn builtin_patterns() -> Vec { + vec![ + SecretPattern { + name: "OpenAI API Key", + regex: r"\bsk-[a-zA-Z0-9]{48}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Anthropic API Key", + regex: r"\bsk-ant-api\d{0,2}-[a-zA-Z0-9\-]{80,120}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Generic SK API Key", + regex: r"\bsk-[a-zA-Z0-9\-]{10,100}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Google API Key", + regex: r"\bAIza[0-9A-Za-z-_]{35}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "AWS Access Key", + regex: r"\b(AKIA|A3T|AGPA|AIDA|AROA|AIPA|ANPA|ANVA|ASIA)[A-Z0-9]{12,}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "GitHub Classic PAT", + regex: r"\bghp_[A-Za-z0-9_]{36}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "GitHub Fine-Grained PAT", + regex: r"\bgithub_pat_[A-Za-z0-9_]{82}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "GitHub OAuth Token", + regex: r"\bgho_[A-Za-z0-9_]{36}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Stripe Key", + regex: r"\b[rs]k_(test|live)_[0-9a-zA-Z]{24}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Slack App Token", + regex: r"\bxapp-[0-9]+-[A-Za-z0-9_]+-[0-9]+-[a-f0-9]+\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "Fireworks API Key", + regex: r"\bfw_[a-zA-Z0-9]{24}\b", + secret_type: SecretType::ApiKey, + }, + SecretPattern { + name: "JWT", + regex: r"\bey[a-zA-Z0-9_\-=]{10,}\.[a-zA-Z0-9_\-=]{10,}\.[a-zA-Z0-9_\-=]{10,}\b", + secret_type: SecretType::Token, + }, + SecretPattern { + name: "URL Embedded Password", + regex: r"://[^/\s:]+:[^/\s@]+@", + secret_type: SecretType::Password, + }, + ] +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builtin_patterns_are_valid_regex() { + for p in builtin_patterns() { + assert!( + regex::Regex::new(p.regex).is_ok(), + "builtin pattern '{}' has invalid regex: {}", + p.name, + p.regex + ); + } + } + + #[test] + fn builtin_patterns_are_not_empty() { + assert!(!builtin_patterns().is_empty()); + } + + #[test] + fn secret_type_default_is_api_key() { + assert_eq!(default_secret_type(), SecretType::ApiKey); + } + + #[test] + fn stripe_pattern_matches_rk_and_sk_prefix() { + let re = regex::Regex::new( + builtin_patterns() + .iter() + .find(|p| p.name == "Stripe Key") + .unwrap() + .regex, + ) + .unwrap(); + let key_24 = "abcdefghijklmnopqrstuvwx"; // exactly 24 chars + assert!(re.is_match(&format!("sk_test_{key_24}"))); + assert!(re.is_match(&format!("rk_live_{key_24}"))); + assert!(!re.is_match(&format!("xk_test_{key_24}"))); + } +} diff --git a/crates/aish-security/src/secret/scanner.rs b/crates/aish-security/src/secret/scanner.rs new file mode 100644 index 0000000..d7b1cc1 --- /dev/null +++ b/crates/aish-security/src/secret/scanner.rs @@ -0,0 +1,199 @@ +use regex::{Regex, RegexSet}; + +use crate::secret::patterns::{builtin_patterns, CustomPattern, SecretMatch, SecretType}; + +#[derive(Clone)] +pub struct SecretScanner { + set: RegexSet, + compiled: Vec, + pattern_names: Vec, + pattern_types: Vec, +} + +impl SecretScanner { + /// Build scanner from user-defined custom patterns (built-in patterns always included). + /// Invalid custom patterns are logged and skipped. + pub fn new(custom: &[CustomPattern]) -> Self { + let builtins = builtin_patterns(); + let mut all_regexes: Vec = Vec::with_capacity(builtins.len() + custom.len()); + let mut names: Vec = Vec::with_capacity(all_regexes.len()); + let mut types: Vec = Vec::with_capacity(all_regexes.len()); + + for p in &builtins { + all_regexes.push(p.regex.to_string()); + names.push(p.name.to_string()); + types.push(p.secret_type); + } + + for c in custom { + match regex::Regex::new(&c.pattern) { + Ok(_) => { + all_regexes.push(c.pattern.clone()); + names.push(c.name.clone()); + types.push(c.secret_type); + } + Err(e) => { + tracing::warn!("skipping invalid custom secret pattern '{}': {}", c.name, e); + } + } + } + + let set = match RegexSet::new(&all_regexes) { + Ok(s) => s, + Err(e) => { + tracing::error!("failed to compile secret RegexSet: {e}"); + return Self::empty(); + } + }; + + let compiled: Vec = all_regexes + .iter() + .map(|p| { + Regex::new(p) + .unwrap_or_else(|e| panic!("regex that passed RegexSet failed Regex::new: {e}")) + }) + .collect(); + + Self { + set, + compiled, + pattern_names: names, + pattern_types: types, + } + } + + /// Scanner that never matches anything, used as a safe fallback. + fn empty() -> Self { + let set = RegexSet::new(&["^.$_never_match_"]).unwrap(); + Self { + set, + compiled: Vec::new(), + pattern_names: Vec::new(), + pattern_types: Vec::new(), + } + } + + /// Scan input text for secrets. Returns all matches. + pub fn scan(&self, input: &str) -> Vec { + let matched_indices: Vec<_> = self.set.matches(input).into_iter().collect(); + if matched_indices.is_empty() { + return Vec::new(); + } + + let mut results = Vec::new(); + for idx in matched_indices { + let Some(re) = self.compiled.get(idx) else { + continue; + }; + for mat in re.find_iter(input) { + results.push(SecretMatch { + pattern_name: self.pattern_names[idx].clone(), + start: mat.start(), + end: mat.end(), + secret_type: self.pattern_types[idx], + }); + } + } + + results.sort_by(|a, b| { + a.start.cmp(&b.start).then_with(|| a.end.cmp(&b.end)) // narrower span first + }); + // Deduplicate overlapping matches: keep the first (narrowest/most-specific) match + // for each overlapping region. Built-in patterns are ordered from most specific + // to least specific, so earlier pattern indices produce better results. + let mut deduped: Vec = Vec::new(); + for m in results { + if let Some(prev) = deduped.last() { + if m.start < prev.end { + continue; + } + } + deduped.push(m); + } + deduped + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::secret::patterns::{CustomPattern, SecretType}; + + fn default_scanner() -> SecretScanner { + SecretScanner::new(&[]) + } + + #[test] + fn scan_returns_empty_for_clean_input() { + let scanner = default_scanner(); + assert!(scanner.scan("hello world, no secrets here").is_empty()); + } + + #[test] + fn scan_detects_openai_key() { + let scanner = default_scanner(); + // sk- + exactly 48 alphanumeric characters (OpenAI pattern requires exactly 48). + let input = "please use sk-abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuv for auth"; + let matches = scanner.scan(input); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].pattern_name, "OpenAI API Key"); + assert_eq!(matches[0].secret_type, SecretType::ApiKey); + } + + #[test] + fn scan_detects_url_embedded_password() { + let scanner = default_scanner(); + let input = "connect to https://admin:s3cret@db.example.com:5432"; + let matches = scanner.scan(input); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].pattern_name, "URL Embedded Password"); + assert_eq!(matches[0].secret_type, SecretType::Password); + } + + #[test] + fn scan_detects_multiple_secrets() { + let scanner = default_scanner(); + // sk- + exactly 48 alphanumeric chars for OpenAI key. + let sk_key = format!("sk-{}", "a".repeat(48)); + let jwt = format!("ey{}.{}.{}", "A".repeat(12), "B".repeat(12), "C".repeat(12)); + let input = format!("key={sk_key} token={jwt}"); + let matches = scanner.scan(&input); + assert!(matches.len() >= 2); + } + + #[test] + fn scan_with_custom_pattern() { + let custom = vec![CustomPattern { + name: "Test Token".to_string(), + pattern: r"test_[a-zA-Z0-9]{16}".to_string(), + secret_type: SecretType::Token, + }]; + let scanner = SecretScanner::new(&custom); + let matches = scanner.scan("my token is test_abcdefghijklmnop"); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].pattern_name, "Test Token"); + assert_eq!(matches[0].secret_type, SecretType::Token); + } + + #[test] + fn scan_skips_invalid_custom_pattern() { + let custom = vec![CustomPattern { + name: "Bad".to_string(), + pattern: "[invalid(".to_string(), + secret_type: SecretType::ApiKey, + }]; + let scanner = SecretScanner::new(&custom); + assert!(scanner.scan("hello").is_empty()); + } + + #[test] + fn scan_deduplicates_overlapping_matches_to_most_specific() { + let scanner = default_scanner(); + // OpenAI key (48 chars) also matches Generic SK (10-100 chars). + // Dedup should keep only the more specific OpenAI match. + let input = "key=sk-abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuv"; + let matches = scanner.scan(input); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].pattern_name, "OpenAI API Key"); + } +} diff --git a/crates/aish-security/src/secret/vault.rs b/crates/aish-security/src/secret/vault.rs new file mode 100644 index 0000000..29906cb --- /dev/null +++ b/crates/aish-security/src/secret/vault.rs @@ -0,0 +1,442 @@ +use std::collections::HashMap; + +use crate::secret::patterns::SecretMatch; + +/// Internal entry tracking a single secret's placeholder and original value. +#[derive(Debug)] +struct SecretEntry { + placeholder: String, + original: String, +} + +/// In-memory vault that maps detected secrets to stable placeholders. +/// +/// Typical usage: +/// 1. Call [`SecretVault::redact`] with scanner results to replace secret +/// values with `$SECRET_…` placeholders. +/// 2. Pass the redacted text to an external service. +/// 3. Call [`SecretVault::restore`] on the response to put original values back. +#[derive(Debug, Default)] +pub struct SecretVault { + /// Maps original secret value → entry metadata. + entries: HashMap, + /// Tracks how many distinct values have been seen per pattern name, + /// so that collisions receive `_2`, `_3`, … suffixes. + value_counter: HashMap, +} + +/// Convert a human-readable pattern name into a placeholder token. +/// +/// Rules: +/// - Convert to uppercase +/// - Replace non-alphanumeric characters with `_` +/// - Collapse consecutive `_` into one +/// - Strip leading/trailing `_` +/// - Prepend `$SECRET_` +/// +/// Examples: +/// - `"OpenAI API Key"` → `"$SECRET_OPENAI_API_KEY"` +/// - `"JWT"` → `"$SECRET_JWT"` +/// - `"URL Embedded Password"` → `"$SECRET_URL_EMBEDDED_PASSWORD"` +pub fn pattern_name_to_placeholder(name: &str) -> String { + let upper = name.to_uppercase(); + let mut normalized = String::with_capacity(upper.len()); + let mut prev_underscore = false; + + for ch in upper.chars() { + if ch.is_ascii_alphanumeric() { + normalized.push(ch); + prev_underscore = false; + } else { + if !prev_underscore && !normalized.is_empty() { + normalized.push('_'); + prev_underscore = true; + } + } + } + + // Strip trailing underscore + if normalized.ends_with('_') { + normalized.pop(); + } + + format!("$SECRET_{normalized}") +} + +impl SecretVault { + /// Create a new, empty vault. + pub fn new() -> Self { + Self::default() + } + + /// Number of secrets currently stored in the vault. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Whether the vault contains any secrets. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Return (or create) a unique placeholder for the given pattern name and value. + /// + /// - If the exact value was already stored, return its existing placeholder. + /// - If the pattern name was seen before with a *different* value, append + /// `_2`, `_3`, … to distinguish. + /// - Otherwise create a new placeholder from the pattern name. + fn find_unique_placeholder(&mut self, pattern_name: &str, value: &str) -> String { + // Same value → reuse existing placeholder + if let Some(entry) = self.entries.get(value) { + return entry.placeholder.clone(); + } + + let base = pattern_name_to_placeholder(pattern_name); + let count = self + .value_counter + .entry(pattern_name.to_string()) + .or_insert(0); + *count += 1; + + let placeholder = if *count == 1 { + base + } else { + format!("{base}_{}", *count) + }; + + placeholder + } + + /// Replace detected secrets in `input` with placeholder tokens. + /// + /// Matches are processed right-to-left so that byte offsets remain valid + /// as earlier portions of the string are untouched until their turn. + /// + /// Returns the redacted string with all secrets replaced by `$SECRET_…` tokens. + pub fn redact(&mut self, secrets: &[SecretMatch], input: &str) -> String { + if secrets.is_empty() { + return input.to_string(); + } + + // Collect (placeholder, start, end) tuples first, using original offsets. + let mut replacements: Vec<(String, usize, usize)> = Vec::with_capacity(secrets.len()); + + for sm in secrets { + let value = &input[sm.start..sm.end]; + let placeholder = self.find_unique_placeholder(&sm.pattern_name, value); + + // Store the entry if this is a new value + if !self.entries.contains_key(value) { + self.entries.insert( + value.to_string(), + SecretEntry { + placeholder: placeholder.clone(), + original: value.to_string(), + }, + ); + } + + replacements.push((placeholder, sm.start, sm.end)); + } + + // Sort by start position descending (right-to-left) so that replacing + // later spans first does not shift earlier byte offsets. + replacements.sort_by(|a, b| b.1.cmp(&a.1)); + + let mut result = input.to_string(); + for (placeholder, start, end) in &replacements { + result.replace_range(*start..*end, placeholder); + } + + result + } + + /// Restore placeholders in `text` back to their original secret values. + /// + /// Returns `(restored_text, count)` where `count` is the number of + /// placeholders that were successfully replaced. + pub fn restore(&self, text: &str) -> (String, usize) { + let mut result = text.to_string(); + let mut count = 0; + + // Sort by placeholder length descending so longer (more specific) + // placeholders are replaced first, avoiding prefix collisions + // (e.g. $SECRET_KEY_2 before $SECRET_KEY). + let mut entries: Vec<&SecretEntry> = self.entries.values().collect(); + entries.sort_by(|a, b| b.placeholder.len().cmp(&a.placeholder.len())); + + for entry in entries { + let occurrences = result.matches(&entry.placeholder).count(); + if occurrences > 0 { + result = result.replace(&entry.placeholder, &entry.original); + count += occurrences; + } + } + + (result, count) + } + + /// Redact secret values from command output before sending to AI. + /// + /// Scans `text` for all stored original values and replaces them with + /// their corresponding placeholders. Returns `(redacted_text, count)`. + /// + /// Sort by original value length descending to avoid short secrets + /// matching substrings of longer ones. + pub fn redact_output(&self, text: &str) -> (String, usize) { + let mut result = text.to_string(); + let mut count = 0; + + let mut entries: Vec<&SecretEntry> = self.entries.values().collect(); + entries.sort_by(|a, b| b.original.len().cmp(&a.original.len())); + + for entry in entries { + let occurrences = result.matches(&entry.original).count(); + if occurrences > 0 { + result = result.replace(&entry.original, &entry.placeholder); + count += occurrences; + } + } + + (result, count) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::secret::patterns::SecretType; + + /// Helper to build a SecretMatch quickly. + fn make_match(name: &str, start: usize, end: usize, st: SecretType) -> SecretMatch { + SecretMatch { + pattern_name: name.to_string(), + start, + end, + secret_type: st, + } + } + + // ── pattern_name_to_placeholder ────────────────────────────────────── + + #[test] + fn placeholder_converts_openai_api_key() { + assert_eq!( + pattern_name_to_placeholder("OpenAI API Key"), + "$SECRET_OPENAI_API_KEY" + ); + } + + #[test] + fn placeholder_converts_single_word() { + assert_eq!(pattern_name_to_placeholder("JWT"), "$SECRET_JWT"); + } + + #[test] + fn placeholder_collapses_multiple_spaces() { + assert_eq!( + pattern_name_to_placeholder(" foo bar "), + "$SECRET_FOO_BAR" + ); + } + + #[test] + fn placeholder_strips_leading_trailing_non_alnum() { + assert_eq!( + pattern_name_to_placeholder("--hello-world--"), + "$SECRET_HELLO_WORLD" + ); + } + + #[test] + fn placeholder_handles_special_chars() { + assert_eq!( + pattern_name_to_placeholder("key/v2 (beta)"), + "$SECRET_KEY_V2_BETA" + ); + } + + // ── redact ─────────────────────────────────────────────────────────── + + #[test] + fn redact_single_secret() { + let mut vault = SecretVault::new(); + let input = "my key is sk-abc123xyz789 use it wisely"; + // "sk-abc123xyz789" starts at 10, ends at 25 + let secrets = vec![make_match("OpenAI API Key", 10, 25, SecretType::ApiKey)]; + + let redacted = vault.redact(&secrets, input); + assert_eq!(redacted, "my key is $SECRET_OPENAI_API_KEY use it wisely"); + } + + #[test] + fn redact_multiple_different_secrets() { + let mut vault = SecretVault::new(); + let input = "key=sk-aaaa token=sk-bbbb"; + // "sk-aaaa" is at 4..11, "sk-bbbb" is at 18..25 + let secrets = vec![ + make_match("OpenAI API Key", 4, 11, SecretType::ApiKey), + make_match("Generic SK API Key", 18, 25, SecretType::ApiKey), + ]; + + let redacted = vault.redact(&secrets, input); + assert_eq!( + redacted, + "key=$SECRET_OPENAI_API_KEY token=$SECRET_GENERIC_SK_API_KEY" + ); + } + + #[test] + fn redact_same_value_reuses_placeholder() { + let mut vault = SecretVault::new(); + let input = "first=sk-aaaa second=sk-aaaa"; + // Both "sk-aaaa" occurrences: 6..13 and 21..28 + let secrets = vec![ + make_match("OpenAI API Key", 6, 13, SecretType::ApiKey), + make_match("OpenAI API Key", 21, 28, SecretType::ApiKey), + ]; + + let redacted = vault.redact(&secrets, input); + assert_eq!( + redacted, + "first=$SECRET_OPENAI_API_KEY second=$SECRET_OPENAI_API_KEY" + ); + } + + #[test] + fn redact_same_type_different_value_gets_suffix() { + let mut vault = SecretVault::new(); + let input = "first=sk-aaaa second=sk-bbbb"; + // "sk-aaaa" at 6..13, "sk-bbbb" at 21..28 + let secrets = vec![ + make_match("OpenAI API Key", 6, 13, SecretType::ApiKey), + make_match("OpenAI API Key", 21, 28, SecretType::ApiKey), + ]; + + let redacted = vault.redact(&secrets, input); + assert_eq!( + redacted, + "first=$SECRET_OPENAI_API_KEY second=$SECRET_OPENAI_API_KEY_2" + ); + } + + #[test] + fn redact_no_secrets_returns_input_unchanged() { + let mut vault = SecretVault::new(); + let input = "nothing to see here"; + let redacted = vault.redact(&[], input); + assert_eq!(redacted, input); + } + + // ── restore ────────────────────────────────────────────────────────── + + #[test] + fn restore_unknown_placeholder_passes_through() { + let vault = SecretVault::new(); + let text = "hello $SECRET_UNKNOWN world"; + let (restored, count) = vault.restore(text); + assert_eq!(restored, text); + assert_eq!(count, 0); + } + + #[test] + fn roundtrip_redact_restore_preserves_original() { + let mut vault = SecretVault::new(); + let input = "key=sk-abc123xyz789 done"; + let secrets = vec![make_match("OpenAI API Key", 4, 19, SecretType::ApiKey)]; + + let redacted = vault.redact(&secrets, input); + assert_ne!(redacted, input); + + let (restored, count) = vault.restore(&redacted); + assert_eq!(restored, input); + assert_eq!(count, 1); + } + + #[test] + fn roundtrip_multiple_secrets() { + let mut vault = SecretVault::new(); + let input = "a=sk-aaaa b=sk-bbbb c=sk-aaaa"; + // "sk-aaaa" at 2..9, "sk-bbbb" at 12..19, "sk-aaaa" at 22..29 + let secrets = vec![ + make_match("OpenAI API Key", 2, 9, SecretType::ApiKey), + make_match("OpenAI API Key", 12, 19, SecretType::ApiKey), + make_match("OpenAI API Key", 22, 29, SecretType::ApiKey), + ]; + + let redacted = vault.redact(&secrets, input); + // Two distinct values: first occurrence gets base, second gets _2, + // third reuses the same placeholder as the first (same value). + assert!(redacted.contains("$SECRET_OPENAI_API_KEY")); + assert!(redacted.contains("$SECRET_OPENAI_API_KEY_2")); + + let (restored, count) = vault.restore(&redacted); + assert_eq!(restored, input); + assert_eq!(count, 3); + } + + // ── redact_output ──────────────────────────────────────────────────── + + #[test] + fn redact_output_removes_secret_from_command_output() { + let mut vault = SecretVault::new(); + let input = "key=sk-abc123xyz789 done"; + let secrets = vec![make_match("OpenAI API Key", 4, 19, SecretType::ApiKey)]; + + let _redacted = vault.redact(&secrets, input); + + // Simulate command output containing the real secret. + let output = "sk-abc123xyz789"; + let (result, count) = vault.redact_output(output); + assert_eq!(result, "$SECRET_OPENAI_API_KEY"); + assert_eq!(count, 1); + } + + #[test] + fn redact_output_preserves_non_secret_text() { + let vault = SecretVault::new(); + let (result, count) = vault.redact_output("hello world"); + assert_eq!(result, "hello world"); + assert_eq!(count, 0); + } + + #[test] + fn redact_output_handles_multiple_secrets() { + let mut vault = SecretVault::new(); + let input = "a=sk-aaaa b=sk-bbbb"; + let secrets = vec![ + make_match("OpenAI API Key", 2, 9, SecretType::ApiKey), + make_match("OpenAI API Key", 12, 19, SecretType::ApiKey), + ]; + let _redacted = vault.redact(&secrets, input); + + let output = "results: sk-aaaa and sk-bbbb"; + let (result, count) = vault.redact_output(output); + assert!(result.contains("$SECRET_OPENAI_API_KEY")); + assert!(!result.contains("sk-aaaa")); + assert!(!result.contains("sk-bbbb")); + assert_eq!(count, 2); + } + + #[test] + fn full_roundtrip_redact_restore_redact_output() { + let mut vault = SecretVault::new(); + let input = "my key is sk-abc123xyz789"; + let secrets = vec![make_match("OpenAI API Key", 10, 25, SecretType::ApiKey)]; + + // 1. Redact input + let redacted = vault.redact(&secrets, input); + assert_eq!(redacted, "my key is $SECRET_OPENAI_API_KEY"); + + // 2. AI generates command, restore for execution + let cmd = "echo $SECRET_OPENAI_API_KEY"; + let (restored_cmd, _) = vault.restore(cmd); + assert_eq!(restored_cmd, "echo sk-abc123xyz789"); + + // 3. Simulate command output, redact before sending to AI + let output = "sk-abc123xyz789"; + let (redacted_output, count) = vault.redact_output(output); + assert_eq!(redacted_output, "$SECRET_OPENAI_API_KEY"); + assert_eq!(count, 1); + } +} diff --git a/crates/aish-shell/src/app.rs b/crates/aish-shell/src/app.rs index 4eff430..51341ec 100644 --- a/crates/aish-shell/src/app.rs +++ b/crates/aish-shell/src/app.rs @@ -15,7 +15,7 @@ use aish_llm::{ CancellationToken, ChatMessage, LlmCallbackResult, LlmSession, }; use aish_memory::MemoryManager; -use aish_security::{SecurityManager, SecurityPolicy}; +use aish_security::{load_policy, SecurityManager}; use aish_session::SessionStore; use aish_skills::hotreload::SkillHotReloader; use aish_skills::SkillManager; @@ -235,6 +235,9 @@ pub struct AishShell { pub config: ConfigModel, pub ai_handler: AiHandler, pub security_manager: SecurityManager, + secret_check_closure: + std::sync::Arc Option + Send + Sync>, + secret_vault: std::sync::Arc>, pub session_store: Option, pub skill_manager: SkillManager, pub skill_hot_reloader: Option, @@ -305,7 +308,7 @@ impl AishShell { } // Initialize security manager (before tool registration) - let security_manager = SecurityManager::new(SecurityPolicy::default()); + let security_manager = SecurityManager::new(load_policy(None)); // Register tools let mut tool_registry = ToolRegistry::new(); @@ -314,6 +317,11 @@ impl AishShell { let mut bash_tool = aish_tools::bash::BashTool::new(); bash_tool.set_cancellation_token(llm_session.cancellation_token_arc()); bash_tool.set_pty_slot(pty_slot.clone()); + + // Shared secret vault slot — populated after AishShell construction. + let vault_slot: aish_tools::bash::VaultSlot = + std::sync::Arc::new(std::sync::Mutex::new(None)); + bash_tool.set_secret_vault(vault_slot.clone()); tool_registry.register(Box::new(bash_tool)); tool_registry.register(Box::new(aish_tools::fs::ReadFileTool::new())); tool_registry.register(Box::new(aish_tools::fs::WriteFileTool::new())); @@ -995,11 +1003,27 @@ impl AishShell { // inside AiHandler which needs mutable access during each turn. let shell_skill_manager = SkillManager::new(); + let secret_check_closure = + Self::build_secret_check_closure(security_manager.secret_scanner().clone()); + + let secret_vault = std::sync::Arc::new(std::sync::Mutex::new( + aish_security::secret::SecretVault::new(), + )); + + // Inject the vault into BashTool's slot so command execution can + // restore $SECRET_* placeholders. + { + let mut guard = vault_slot.lock().unwrap(); + *guard = Some(secret_vault.clone()); + } + Ok(Self { state, config, ai_handler, security_manager, + secret_check_closure, + secret_vault, session_store, skill_manager: shell_skill_manager, skill_hot_reloader, @@ -1174,7 +1198,7 @@ impl AishShell { match input::classify_input(input) { crate::types::InputIntent::Empty => {} crate::types::InputIntent::Ai => { - let question = input::extract_ai_question(input); + let mut question = input::extract_ai_question(input); // If just ";" with no question and there's a pending error, // trigger error correction instead of a normal AI query. @@ -1277,6 +1301,50 @@ impl AishShell { } } + // Security gate: detect secrets in AI input + let decision = self.security_manager.check_ai_input(&question); + if decision.require_confirmation { + let reasons = decision + .analysis + .detected_secrets + .as_ref() + .map(|s| { + s.iter() + .map(|m| m.format_reason()) + .collect::>() + .join("\n ") + }) + .unwrap_or_default(); + let mut args = std::collections::HashMap::new(); + args.insert("reasons".to_string(), reasons); + let title = t("shell.security.secret.title"); + let message = t_with_args("shell.security.secret.detected", &args); + let choice = crate::tui::show_secret_dialog(&title, &message); + match choice { + crate::tui::SecretDialogChoice::Abort => { + let aborted = t("shell.security.secret.aborted"); + println!("\x1b[33m{}\x1b[0m", aborted); + continue; + } + crate::tui::SecretDialogChoice::Redact => { + if let Some(secrets) = decision.analysis.detected_secrets { + let count = secrets.len(); + let redacted = self + .secret_vault + .lock() + .unwrap() + .redact(&secrets, &question); + let mut rargs = std::collections::HashMap::new(); + rargs.insert("count".to_string(), count.to_string()); + let msg = t_with_args("shell.security.secret.redacted", &rargs); + println!("\x1b[33m{}\x1b[0m", msg); + question = redacted; + } + } + crate::tui::SecretDialogChoice::Allow => {} + } + } + let old_sigint = self.install_ai_sigint_handler(); let token_ptr = self.ai_handler.cancellation_token() as *const CancellationToken; @@ -1805,7 +1873,13 @@ impl AishShell { let shared_host = Arc::new(Mutex::new(remote_host.clone())); let ai_cb = Self::build_session_ai_callback(&self.config, &self.animation, shared_host.clone()); - pty.send_command_interactive(command, ai_cb, Some(shared_host)) + pty.send_command_interactive( + command, + ai_cb, + Some(shared_host), + Some(self.secret_check_closure.clone()), + Some(self.secret_vault.clone()), + ) }; let (exit_code, cwd, output) = match result { Ok(result) => result, @@ -1852,9 +1926,18 @@ impl AishShell { } else { &output }; + + // Redact secrets from shell context before injecting into LLM context. + let (safe_output, _) = self + .secret_vault + .lock() + .unwrap() + .redact_output(output_preview); + let (safe_command, _) = self.secret_vault.lock().unwrap().redact_output(command); + let mut entry = format!( "[Shell] {}\n{}\n{}", - command, exit_code, output_preview + safe_command, exit_code, safe_output ); if let Some(ref new_cwd) = cwd_changed_to { entry.push_str(&format!("\n{}", new_cwd)); @@ -2090,7 +2173,7 @@ impl AishShell { .pty .lock() .unwrap() - .send_command_interactive(segment, None, None) + .send_command_interactive(segment, None, None, None, None) .unwrap_or((-1, self.state.cwd.clone(), String::new())); if !output.is_empty() { @@ -2732,6 +2815,33 @@ impl AishShell { ) } + /// Build a closure that checks AI input for secrets. + /// Returns Some(SshSecretCheckResult) if secrets found, None if clean. + fn build_secret_check_closure( + scanner: aish_security::secret::SecretScanner, + ) -> std::sync::Arc Option + Send + Sync> { + std::sync::Arc::new(move |input: &str| { + let secrets = scanner.scan(input); + if secrets.is_empty() { + return None; + } + let reasons = secrets + .iter() + .map(|s| s.format_reason()) + .collect::>() + .join("\n "); + let mut args = std::collections::HashMap::new(); + args.insert("reasons".to_string(), reasons); + let title = aish_i18n::t("shell.security.secret.title"); + let message = aish_i18n::t_with_args("shell.security.secret.detected", &args); + let warning = format!("\x1b[33m? {}:\x1b[0m {}", title, message); + Some(aish_pty::SshSecretCheckResult { + warning, + detected_secrets: secrets, + }) + }) + } + fn build_session_ai_callback( config: &aish_config::ConfigModel, animation: &Arc, diff --git a/crates/aish-shell/src/tui.rs b/crates/aish-shell/src/tui.rs index 88f3d4e..a64b4c8 100644 --- a/crates/aish-shell/src/tui.rs +++ b/crates/aish-shell/src/tui.rs @@ -75,11 +75,16 @@ pub fn show_selection_dialog( } } -/// Show a simple Yes/No confirmation dialog. -pub fn show_confirmation_dialog(title: &str, message: &str) -> bool { +/// Show a simple Yes/No confirmation dialog with custom labels. +pub fn show_confirmation_dialog( + title: &str, + message: &str, + yes_label: &str, + no_label: &str, +) -> bool { let options = vec![ - DialogOption::new("yes", "Yes"), - DialogOption::new("no", "No"), + DialogOption::new("yes", yes_label), + DialogOption::new("no", no_label), ]; matches!( show_selection_dialog(title, message, &options, false, true), @@ -87,6 +92,38 @@ pub fn show_confirmation_dialog(title: &str, message: &str) -> bool { ) } +/// Choice returned by the three-option secret detection dialog. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SecretDialogChoice { + /// Replace secrets with env var placeholders. + Redact, + /// Send original plaintext to AI. + Allow, + /// Cancel / abort. + Abort, +} + +/// Show a three-option dialog for secret detection (TUI path). +/// Default selection is "Redact" (safest option). +pub fn show_secret_dialog(title: &str, message: &str) -> SecretDialogChoice { + let redact_label = aish_i18n::t("shell.security.secret.redact"); + let allow_label = aish_i18n::t("shell.security.secret.allow"); + let abort_label = aish_i18n::t("shell.security.secret.abort"); + let options = vec![ + DialogOption::new("redact", redact_label), + DialogOption::new("allow", allow_label), + DialogOption::new("abort", abort_label), + ]; + match show_selection_dialog(title, message, &options, false, true) { + DialogResult::Selected(v) => match v.as_str() { + "redact" => SecretDialogChoice::Redact, + "allow" => SecretDialogChoice::Allow, + _ => SecretDialogChoice::Abort, + }, + _ => SecretDialogChoice::Abort, + } +} + // --------------------------------------------------------------------------- // inquire implementation // --------------------------------------------------------------------------- diff --git a/crates/aish-tools/src/bash.rs b/crates/aish-tools/src/bash.rs index 3ee0d90..3559a7d 100644 --- a/crates/aish-tools/src/bash.rs +++ b/crates/aish-tools/src/bash.rs @@ -9,7 +9,9 @@ use aish_llm::{ ToolResult, }; use aish_pty::{BashOffloadSettings, BashOutputOffload, CancelToken, PtyExecutor}; -use aish_security::{load_policy, SecurityDecision, SecurityManager, SecurityRequest}; +use aish_security::{ + load_policy, secret::SecretVault, SecurityDecision, SecurityManager, SecurityRequest, +}; /// Large keep_bytes for the silent PTY executor to capture full command output. /// The BashOutputOffload will handle threshold-based truncation and disk offload. @@ -80,6 +82,9 @@ impl Drop for InteractiveInputGuard { } } +/// Shared vault slot for secret placeholder restoration. +pub type VaultSlot = Arc>>>>; + /// Tool for executing bash commands via PTY. pub struct BashTool { /// Shared cancellation token from the AI handler. @@ -88,6 +93,8 @@ pub struct BashTool { cancellation_token: Option>, /// Shared slot for PersistentPty — set after PTY creation. pty_slot: PtySlot, + /// Shared vault for restoring $SECRET_* placeholders before command execution. + secret_vault: VaultSlot, } /// Cached translated description. @@ -230,6 +237,7 @@ impl BashTool { Self { cancellation_token: None, pty_slot: Arc::new(Mutex::new(None)), + secret_vault: Arc::new(Mutex::new(None)), } } @@ -244,6 +252,41 @@ impl BashTool { self.pty_slot = slot; } + /// Set the shared secret vault for placeholder restoration. + pub fn set_secret_vault(&mut self, vault: VaultSlot) { + self.secret_vault = vault; + } + + /// Restore `$SECRET_*` placeholders in the command to their actual values. + /// Returns `(restored_command, restore_count)`. + fn restore_secrets(&self, command: &str) -> (String, usize) { + let vault_arc = { + let guard = self.secret_vault.lock().unwrap(); + guard.clone() + }; + if let Some(vault) = vault_arc { + let vault = vault.lock().unwrap(); + vault.restore(command) + } else { + (command.to_string(), 0) + } + } + + /// Redact secret values from command output before sending to AI. + /// Returns `(redacted_text, count)`. + fn redact_output(&self, text: &str) -> (String, usize) { + let vault_arc = { + let guard = self.secret_vault.lock().unwrap(); + guard.clone() + }; + if let Some(vault) = vault_arc { + let vault = vault.lock().unwrap(); + vault.redact_output(text) + } else { + (text.to_string(), 0) + } + } + /// Execute via PersistentPty — supports Ctrl+Z/bg/fg job control. fn execute_via_persistent_pty( &self, @@ -435,12 +478,15 @@ impl Tool for BashTool { return PreflightResult::Allow; }; + // Restore secret placeholders so security checks run on the actual command. + let (command, _) = self.restore_secrets(command); + let cwd = std::env::current_dir().ok(); - security_preflight(command, cwd.as_deref(), None) + security_preflight(&command, cwd.as_deref(), None) } fn execute(&self, args: serde_json::Value) -> ToolResult { - let command = match args.get("command").and_then(|c| c.as_str()) { + let raw_command = match args.get("command").and_then(|c| c.as_str()) { Some(cmd) => cmd, None => return ToolResult::error(aish_i18n::t("tools.bash.missing_command")), }; @@ -449,17 +495,28 @@ impl Tool for BashTool { Err(error) => return error, }; + // Restore secret placeholders before execution. + let (command, _restore_count) = self.restore_secrets(raw_command); + // Try PersistentPty path first (supports Ctrl+Z/bg/fg). let pty_arc = { let guard = self.pty_slot.lock().unwrap(); guard.clone() }; - if let Some(pty_arc) = pty_arc { - return self.execute_via_persistent_pty(command, timeout_secs, pty_arc); + let mut result = if let Some(pty_arc) = pty_arc { + self.execute_via_persistent_pty(&command, timeout_secs, pty_arc) + } else { + // Fallback: one-shot PtyExecutor. + self.execute_via_pty_executor(&command, timeout_secs) + }; + + // Redact secret values from output before returning to AI. + if result.ok { + let (redacted, _) = self.redact_output(&result.output); + result.output = redacted; } - // Fallback: one-shot PtyExecutor. - self.execute_via_pty_executor(command, timeout_secs) + result } }