bamboo_config/

config_crypto.rs

//! Encryption, decryption, and hydration methods for [`Config`].
//!
//! These methods handle the in-memory hydration of encrypted credentials
//! (API keys, proxy auth, MCP secrets, env vars) and their re-encryption
//! before persisting to disk.

use anyhow::{Context, Result};

use super::{Config, ProxyAuth};

impl Config {
    // ── Proxy auth ─────────────────────────────────────────────────────

    /// Populate `proxy_auth` (plaintext) from `proxy_auth_encrypted` if present.
    ///
    /// Many parts of the code rely on `proxy_auth` being hydrated in-memory so
    /// we can re-encrypt deterministically on save without ever persisting
    /// plaintext credentials.
    pub fn hydrate_proxy_auth_from_encrypted(&mut self) {
        if self.proxy_auth.is_some() {
            return;
        }

        // Backward compatibility:
        // Older Bodhi/Tauri builds persisted proxy auth as per-scheme encrypted fields:
        // `http_proxy_auth_encrypted` / `https_proxy_auth_encrypted`.
        //
        // Those live under `extra` (flatten) in the unified config. Seed the new
        // `proxy_auth_encrypted` field so the rest of the code can stay uniform.
        if self
            .proxy_auth_encrypted
            .as_deref()
            .map(|s| s.trim().is_empty())
            .unwrap_or(true)
        {
            let legacy = self
                .extra
                .get("https_proxy_auth_encrypted")
                .and_then(|v| v.as_str())
                .or_else(|| {
                    self.extra
                        .get("http_proxy_auth_encrypted")
                        .and_then(|v| v.as_str())
                })
                .map(|s| s.trim())
                .filter(|s| !s.is_empty())
                .map(|s| s.to_string());

            if let Some(legacy) = legacy {
                self.proxy_auth_encrypted = Some(legacy);
            }
        }

        let Some(encrypted) = self.proxy_auth_encrypted.as_deref() else {
            return;
        };

        match crate::encryption::decrypt(encrypted) {
            Ok(decrypted) => match serde_json::from_str::<ProxyAuth>(&decrypted) {
                Ok(auth) => {
                    self.proxy_auth = Some(auth);
                    // Once hydrated successfully, drop legacy keys so a future save writes only
                    // the canonical `proxy_auth_encrypted` field.
                    self.extra.remove("http_proxy_auth_encrypted");
                    self.extra.remove("https_proxy_auth_encrypted");
                }
                Err(e) => tracing::warn!("Failed to parse decrypted proxy auth JSON: {}", e),
            },
            Err(e) => tracing::warn!("Failed to decrypt proxy auth: {}", e),
        }
    }

    /// Refresh `proxy_auth_encrypted` from the current in-memory `proxy_auth`.
    ///
    /// This is used both when persisting the config to disk and when generating
    /// API responses that should never include plaintext proxy credentials.
    pub fn refresh_proxy_auth_encrypted(&mut self) -> Result<()> {
        // Keep on-disk representation fully derived from the in-memory plaintext:
        // - Some(auth)  => always (re-)encrypt and store `proxy_auth_encrypted`
        // - None        => remove `proxy_auth_encrypted`
        let Some(auth) = self.proxy_auth.as_ref() else {
            self.proxy_auth_encrypted = None;
            return Ok(());
        };

        let auth_str = serde_json::to_string(auth).context("Failed to serialize proxy auth")?;
        let encrypted =
            crate::encryption::encrypt(&auth_str).context("Failed to encrypt proxy auth")?;
        self.proxy_auth_encrypted = Some(encrypted);
        Ok(())
    }

    // ── Provider API keys ──────────────────────────────────────────────

    pub fn hydrate_provider_api_keys_from_encrypted(&mut self) {
        if let Some(openai) = self.providers.openai.as_mut() {
            if openai.api_key.trim().is_empty() {
                if let Some(encrypted) = openai.api_key_encrypted.as_deref() {
                    match crate::encryption::decrypt(encrypted) {
                        Ok(value) => openai.api_key = value,
                        Err(e) => tracing::warn!("Failed to decrypt OpenAI api_key: {}", e),
                    }
                }
            }
        }

        if let Some(anthropic) = self.providers.anthropic.as_mut() {
            if anthropic.api_key.trim().is_empty() {
                if let Some(encrypted) = anthropic.api_key_encrypted.as_deref() {
                    match crate::encryption::decrypt(encrypted) {
                        Ok(value) => anthropic.api_key = value,
                        Err(e) => tracing::warn!("Failed to decrypt Anthropic api_key: {}", e),
                    }
                }
            }
        }

        if let Some(gemini) = self.providers.gemini.as_mut() {
            if gemini.api_key.trim().is_empty() {
                if let Some(encrypted) = gemini.api_key_encrypted.as_deref() {
                    match crate::encryption::decrypt(encrypted) {
                        Ok(value) => gemini.api_key = value,
                        Err(e) => tracing::warn!("Failed to decrypt Gemini api_key: {}", e),
                    }
                }
            }
        }

        if let Some(bodhi) = self.providers.bodhi.as_mut() {
            if bodhi.api_key.trim().is_empty() {
                if let Some(encrypted) = bodhi.api_key_encrypted.as_deref() {
                    match crate::encryption::decrypt(encrypted) {
                        Ok(value) => bodhi.api_key = value,
                        Err(e) => tracing::warn!("Failed to decrypt Bodhi api_key: {}", e),
                    }
                }
            }
        }
    }

    pub fn refresh_provider_api_keys_encrypted(&mut self) -> Result<()> {
        if let Some(openai) = self.providers.openai.as_mut() {
            let api_key = openai.api_key.trim();
            openai.api_key_encrypted = if api_key.is_empty() {
                None
            } else {
                Some(
                    crate::encryption::encrypt(api_key)
                        .context("Failed to encrypt OpenAI api_key")?,
                )
            };
        }

        if let Some(anthropic) = self.providers.anthropic.as_mut() {
            let api_key = anthropic.api_key.trim();
            anthropic.api_key_encrypted = if api_key.is_empty() {
                None
            } else {
                Some(
                    crate::encryption::encrypt(api_key)
                        .context("Failed to encrypt Anthropic api_key")?,
                )
            };
        }

        if let Some(gemini) = self.providers.gemini.as_mut() {
            let api_key = gemini.api_key.trim();
            gemini.api_key_encrypted = if api_key.is_empty() {
                None
            } else {
                Some(
                    crate::encryption::encrypt(api_key)
                        .context("Failed to encrypt Gemini api_key")?,
                )
            };
        }

        if let Some(bodhi) = self.providers.bodhi.as_mut() {
            let api_key = bodhi.api_key.trim();
            bodhi.api_key_encrypted = if api_key.is_empty() {
                None
            } else {
                Some(
                    crate::encryption::encrypt(api_key)
                        .context("Failed to encrypt Bodhi api_key")?,
                )
            };
        }

        Ok(())
    }

    // ── Provider instance API keys ─────────────────────────────────────

    /// Hydrate plaintext `api_key` fields on provider instances from their
    /// encrypted counterparts.
    pub fn hydrate_provider_instance_api_keys_from_encrypted(&mut self) {
        for (id, instance) in self.provider_instances.iter_mut() {
            if instance.api_key.trim().is_empty() {
                if let Some(encrypted) = instance.api_key_encrypted.as_deref() {
                    match crate::encryption::decrypt(encrypted) {
                        Ok(value) => instance.api_key = value,
                        Err(e) => {
                            tracing::warn!(instance_id = id, "Failed to decrypt api_key: {}", e)
                        }
                    }
                }
            }
        }
    }

    /// Re-encrypt all provider instance API keys and write back to
    /// `api_key_encrypted`. Used before persisting to disk.
    pub fn refresh_provider_instance_api_keys_encrypted(&mut self) -> Result<()> {
        for (id, instance) in self.provider_instances.iter_mut() {
            let api_key = instance.api_key.trim();
            instance.api_key_encrypted = if api_key.is_empty() {
                None
            } else {
                Some(crate::encryption::encrypt(api_key).context(format!(
                    "Failed to encrypt api_key for provider instance '{}'",
                    id
                ))?)
            };
        }
        Ok(())
    }

    // ── MCP secrets ────────────────────────────────────────────────────

    pub fn hydrate_mcp_secrets_from_encrypted(&mut self) {
        for server in self.mcp.servers.iter_mut() {
            match &mut server.transport {
                bamboo_domain::mcp_config::TransportConfig::Stdio(stdio) => {
                    if stdio.env_encrypted.is_empty() {
                        continue;
                    }

                    // Avoid borrow-checker gymnastics by iterating a cloned map.
                    for (key, encrypted) in stdio.env_encrypted.clone() {
                        let should_hydrate = stdio
                            .env
                            .get(&key)
                            .map(|v| v.trim().is_empty())
                            .unwrap_or(true);
                        if !should_hydrate {
                            continue;
                        }

                        match crate::encryption::decrypt(&encrypted) {
                            Ok(value) => {
                                stdio.env.insert(key, value);
                            }
                            Err(e) => tracing::warn!("Failed to decrypt MCP stdio env var: {}", e),
                        }
                    }
                }
                bamboo_domain::mcp_config::TransportConfig::Sse(sse) => {
                    for header in sse.headers.iter_mut() {
                        if !header.value.trim().is_empty() {
                            continue;
                        }
                        let Some(encrypted) = header.value_encrypted.as_deref() else {
                            continue;
                        };
                        match crate::encryption::decrypt(encrypted) {
                            Ok(value) => header.value = value,
                            Err(e) => {
                                tracing::warn!("Failed to decrypt MCP SSE header value: {}", e)
                            }
                        }
                    }
                }
                bamboo_domain::mcp_config::TransportConfig::StreamableHttp(sh) => {
                    for header in sh.headers.iter_mut() {
                        if !header.value.trim().is_empty() {
                            continue;
                        }
                        let Some(encrypted) = header.value_encrypted.as_deref() else {
                            continue;
                        };
                        match crate::encryption::decrypt(encrypted) {
                            Ok(value) => header.value = value,
                            Err(e) => {
                                tracing::warn!(
                                    "Failed to decrypt MCP StreamableHTTP header value: {}",
                                    e
                                )
                            }
                        }
                    }
                }
            }
        }
    }

    pub fn refresh_mcp_secrets_encrypted(&mut self) -> Result<()> {
        for server in self.mcp.servers.iter_mut() {
            match &mut server.transport {
                bamboo_domain::mcp_config::TransportConfig::Stdio(stdio) => {
                    stdio.env_encrypted.clear();
                    for (key, value) in &stdio.env {
                        let encrypted = crate::encryption::encrypt(value).with_context(|| {
                            format!("Failed to encrypt MCP stdio env var '{key}'")
                        })?;
                        stdio.env_encrypted.insert(key.clone(), encrypted);
                    }
                }
                bamboo_domain::mcp_config::TransportConfig::Sse(sse) => {
                    for header in sse.headers.iter_mut() {
                        let configured = !header.value.trim().is_empty();
                        header.value_encrypted = if !configured {
                            None
                        } else {
                            Some(crate::encryption::encrypt(&header.value).with_context(|| {
                                format!("Failed to encrypt MCP SSE header '{}'", header.name)
                            })?)
                        };
                    }
                }
                bamboo_domain::mcp_config::TransportConfig::StreamableHttp(sh) => {
                    for header in sh.headers.iter_mut() {
                        let configured = !header.value.trim().is_empty();
                        header.value_encrypted = if !configured {
                            None
                        } else {
                            Some(crate::encryption::encrypt(&header.value).with_context(|| {
                                format!(
                                    "Failed to encrypt MCP StreamableHTTP header '{}'",
                                    header.name
                                )
                            })?)
                        };
                    }
                }
            }
        }

        Ok(())
    }

    // ── Env vars encryption ────────────────────────────────────────────

    /// Decrypt secret env vars into in-memory plaintext after loading config.
    pub fn hydrate_env_vars_from_encrypted(&mut self) {
        for entry in &mut self.env_vars {
            if !entry.secret {
                continue;
            }
            if !entry.value.trim().is_empty() {
                // Already has plaintext (e.g. in-memory update).
                continue;
            }
            let Some(encrypted) = &entry.value_encrypted else {
                continue;
            };
            match crate::encryption::decrypt(encrypted) {
                Ok(value) => entry.value = value,
                Err(e) => tracing::warn!("Failed to decrypt env var '{}': {}", entry.name, e),
            }
        }
    }

    /// Re-encrypt secret env vars before persisting to disk.
    pub fn refresh_env_vars_encrypted(&mut self) -> Result<()> {
        for entry in &mut self.env_vars {
            if entry.secret && !entry.value.trim().is_empty() {
                entry.value_encrypted = Some(
                    crate::encryption::encrypt(&entry.value)
                        .with_context(|| format!("Failed to encrypt env var '{}'", entry.name))?,
                );
            } else if !entry.secret {
                entry.value_encrypted = None;
            }
        }
        Ok(())
    }

    /// Clear plaintext values for secrets before serialization to disk.
    pub fn sanitize_env_vars_for_disk(&mut self) {
        for entry in &mut self.env_vars {
            if entry.secret {
                entry.value = String::new();
            }
        }
    }
}