sandbox-agent/server/packages/agent-credentials/src/lib.rs

use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ProviderCredentials {
    pub api_key: String,
    pub source: String,
    pub auth_type: AuthType,
    pub provider: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AuthType {
    ApiKey,
    Oauth,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ExtractedCredentials {
    pub anthropic: Option<ProviderCredentials>,
    pub openai: Option<ProviderCredentials>,
    pub other: HashMap<String, ProviderCredentials>,
}

#[derive(Debug, Clone, Default)]
pub struct CredentialExtractionOptions {
    pub home_dir: Option<PathBuf>,
    pub include_oauth: bool,
}

impl CredentialExtractionOptions {
    pub fn new() -> Self {
        Self {
            home_dir: None,
            include_oauth: true,
        }
    }
}

pub fn extract_claude_credentials(options: &CredentialExtractionOptions) -> Option<ProviderCredentials> {
    let home_dir = options.home_dir.clone().unwrap_or_else(default_home_dir);
    let include_oauth = options.include_oauth;

    let config_paths = [
        home_dir.join(".claude.json.api"),
        home_dir.join(".claude.json"),
        home_dir.join(".claude.json.nathan"),
    ];

    let key_paths = [
        vec!["primaryApiKey"],
        vec!["apiKey"],
        vec!["anthropicApiKey"],
        vec!["customApiKey"],
    ];

    for path in config_paths {
        let data = read_json_file(&path)?;
        for key_path in &key_paths {
            if let Some(key) = read_string_field(&data, key_path) {
                if key.starts_with("sk-ant-") {
                    return Some(ProviderCredentials {
                        api_key: key,
                        source: "claude-code".to_string(),
                        auth_type: AuthType::ApiKey,
                        provider: "anthropic".to_string(),
                    });
                }
            }
        }
    }

    if include_oauth {
        let oauth_paths = [
            home_dir.join(".claude").join(".credentials.json"),
            home_dir.join(".claude-oauth-credentials.json"),
        ];
        for path in oauth_paths {
            let data = match read_json_file(&path) {
                Some(value) => value,
                None => continue,
            };
            let access = read_string_field(&data, &["claudeAiOauth", "accessToken"]);
            if let Some(token) = access {
                if let Some(expires_at) =
                    read_string_field(&data, &["claudeAiOauth", "expiresAt"])
                {
                    if is_expired_rfc3339(&expires_at) {
                        continue;
                    }
                }
                return Some(ProviderCredentials {
                    api_key: token,
                    source: "claude-code".to_string(),
                    auth_type: AuthType::Oauth,
                    provider: "anthropic".to_string(),
                });
            }
        }
    }

    None
}

pub fn extract_codex_credentials(options: &CredentialExtractionOptions) -> Option<ProviderCredentials> {
    let home_dir = options.home_dir.clone().unwrap_or_else(default_home_dir);
    let include_oauth = options.include_oauth;
    let path = home_dir.join(".codex").join("auth.json");
    let data = read_json_file(&path)?;

    if let Some(key) = data.get("OPENAI_API_KEY").and_then(Value::as_str) {
        if !key.is_empty() {
            return Some(ProviderCredentials {
                api_key: key.to_string(),
                source: "codex".to_string(),
                auth_type: AuthType::ApiKey,
                provider: "openai".to_string(),
            });
        }
    }

    if include_oauth {
        if let Some(token) = read_string_field(&data, &["tokens", "access_token"]) {
            return Some(ProviderCredentials {
                api_key: token,
                source: "codex".to_string(),
                auth_type: AuthType::Oauth,
                provider: "openai".to_string(),
            });
        }
    }

    None
}

pub fn extract_opencode_credentials(options: &CredentialExtractionOptions) -> ExtractedCredentials {
    let home_dir = options.home_dir.clone().unwrap_or_else(default_home_dir);
    let include_oauth = options.include_oauth;
    let path = home_dir
        .join(".local")
        .join("share")
        .join("opencode")
        .join("auth.json");

    let mut result = ExtractedCredentials::default();
    let data = match read_json_file(&path) {
        Some(value) => value,
        None => return result,
    };

    let obj = match data.as_object() {
        Some(obj) => obj,
        None => return result,
    };

    for (provider_name, value) in obj {
        let config = match value.as_object() {
            Some(config) => config,
            None => continue,
        };

        let auth_type = config
            .get("type")
            .and_then(Value::as_str)
            .unwrap_or("");

        let credentials = if auth_type == "api" {
            config.get("key").and_then(Value::as_str).map(|key| ProviderCredentials {
                api_key: key.to_string(),
                source: "opencode".to_string(),
                auth_type: AuthType::ApiKey,
                provider: provider_name.to_string(),
            })
        } else if auth_type == "oauth" && include_oauth {
            let expires = config.get("expires").and_then(Value::as_i64);
            if let Some(expires) = expires {
                if expires < current_epoch_millis() {
                    None
                } else {
                    config
                        .get("access")
                        .and_then(Value::as_str)
                        .map(|token| ProviderCredentials {
                            api_key: token.to_string(),
                            source: "opencode".to_string(),
                            auth_type: AuthType::Oauth,
                            provider: provider_name.to_string(),
                        })
                }
            } else {
                config
                    .get("access")
                    .and_then(Value::as_str)
                    .map(|token| ProviderCredentials {
                        api_key: token.to_string(),
                        source: "opencode".to_string(),
                        auth_type: AuthType::Oauth,
                        provider: provider_name.to_string(),
                    })
            }
        } else {
            None
        };

        if let Some(credentials) = credentials {
            if provider_name == "anthropic" {
                result.anthropic = Some(credentials.clone());
            } else if provider_name == "openai" {
                result.openai = Some(credentials.clone());
            } else {
                result.other.insert(provider_name.to_string(), credentials.clone());
            }
        }
    }

    result
}

pub fn extract_amp_credentials(options: &CredentialExtractionOptions) -> Option<ProviderCredentials> {
    let home_dir = options.home_dir.clone().unwrap_or_else(default_home_dir);
    let path = home_dir.join(".amp").join("config.json");
    let data = read_json_file(&path)?;

    let key_paths: Vec<Vec<&str>> = vec![
        vec!["anthropicApiKey"],
        vec!["anthropic_api_key"],
        vec!["apiKey"],
        vec!["api_key"],
        vec!["accessToken"],
        vec!["access_token"],
        vec!["token"],
        vec!["auth", "anthropicApiKey"],
        vec!["auth", "apiKey"],
        vec!["auth", "token"],
        vec!["anthropic", "apiKey"],
        vec!["anthropic", "token"],
    ];

    for key_path in key_paths {
        if let Some(key) = read_string_field(&data, &key_path) {
            if !key.is_empty() {
                return Some(ProviderCredentials {
                    api_key: key,
                    source: "amp".to_string(),
                    auth_type: AuthType::ApiKey,
                    provider: "anthropic".to_string(),
                });
            }
        }
    }

    None
}

pub fn extract_all_credentials(options: &CredentialExtractionOptions) -> ExtractedCredentials {
    let mut result = ExtractedCredentials::default();

    if let Ok(value) = std::env::var("ANTHROPIC_API_KEY") {
        result.anthropic = Some(ProviderCredentials {
            api_key: value,
            source: "environment".to_string(),
            auth_type: AuthType::ApiKey,
            provider: "anthropic".to_string(),
        });
    } else if let Ok(value) = std::env::var("CLAUDE_API_KEY") {
        result.anthropic = Some(ProviderCredentials {
            api_key: value,
            source: "environment".to_string(),
            auth_type: AuthType::ApiKey,
            provider: "anthropic".to_string(),
        });
    }

    if let Ok(value) = std::env::var("OPENAI_API_KEY") {
        result.openai = Some(ProviderCredentials {
            api_key: value,
            source: "environment".to_string(),
            auth_type: AuthType::ApiKey,
            provider: "openai".to_string(),
        });
    } else if let Ok(value) = std::env::var("CODEX_API_KEY") {
        result.openai = Some(ProviderCredentials {
            api_key: value,
            source: "environment".to_string(),
            auth_type: AuthType::ApiKey,
            provider: "openai".to_string(),
        });
    }

    if result.anthropic.is_none() {
        result.anthropic = extract_amp_credentials(options);
    }

    if result.anthropic.is_none() {
        result.anthropic = extract_claude_credentials(options);
    }

    if result.openai.is_none() {
        result.openai = extract_codex_credentials(options);
    }

    let opencode_credentials = extract_opencode_credentials(options);
    if result.anthropic.is_none() {
        result.anthropic = opencode_credentials.anthropic.clone();
    }
    if result.openai.is_none() {
        result.openai = opencode_credentials.openai.clone();
    }

    for (key, value) in opencode_credentials.other {
        result.other.entry(key).or_insert(value);
    }

    result
}

pub fn get_anthropic_api_key(options: &CredentialExtractionOptions) -> Option<String> {
    extract_all_credentials(options)
        .anthropic
        .map(|cred| cred.api_key)
}

pub fn get_openai_api_key(options: &CredentialExtractionOptions) -> Option<String> {
    extract_all_credentials(options)
        .openai
        .map(|cred| cred.api_key)
}

pub fn set_credentials_as_env_vars(credentials: &ExtractedCredentials) {
    if let Some(cred) = &credentials.anthropic {
        std::env::set_var("ANTHROPIC_API_KEY", &cred.api_key);
    }
    if let Some(cred) = &credentials.openai {
        std::env::set_var("OPENAI_API_KEY", &cred.api_key);
    }
}

fn read_json_file(path: &Path) -> Option<Value> {
    let contents = fs::read_to_string(path).ok()?;
    serde_json::from_str(&contents).ok()
}

fn read_string_field(value: &Value, path: &[&str]) -> Option<String> {
    let mut current = value;
    for key in path {
        current = current.get(*key)?;
    }
    current.as_str().map(|s| s.to_string())
}

fn default_home_dir() -> PathBuf {
    dirs::home_dir().unwrap_or_else(|| PathBuf::from("."))
}

fn current_epoch_millis() -> i64 {
    let now = OffsetDateTime::now_utc();
    (now.unix_timestamp() * 1000) + (now.millisecond() as i64)
}

fn is_expired_rfc3339(value: &str) -> bool {
    match OffsetDateTime::parse(value, &time::format_description::well_known::Rfc3339) {
        Ok(expiry) => expiry < OffsetDateTime::now_utc(),
        Err(_) => false,
    }
}