mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
Add hosted MCP server support with OAuth 2.1 and token refresh
Enables connecting to official MCP servers (like Notion) instead of building custom WASM tools. Uses OAuth 2.1 with PKCE and supports Dynamic Client Registration for zero-config authentication. Key features: - OAuth 2.1 flow with PKCE for secure browser-based auth - Dynamic Client Registration (DCR) for servers without pre-configured clients - Automatic token refresh on 401 responses - Session management with Mcp-Session-Id headers - SSE streaming response handling New CLI commands: - `mcp add <name> <url>` - Add an MCP server - `mcp remove <name>` - Remove an MCP server - `mcp list` - List configured servers - `mcp auth <name>` - Authenticate with a server - `mcp test <name>` - Test connection Also removes the Notion WASM tool example since it's superseded by the Notion MCP server which provides 13 official tools. Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.5
parent
5992e27507
commit
974bc8d407
+484
@@ -0,0 +1,484 @@
|
||||
//! MCP server management CLI commands.
|
||||
//!
|
||||
//! Commands for adding, removing, authenticating, and testing MCP servers.
|
||||
|
||||
use std::io::Write;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::mcp::{
|
||||
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
|
||||
auth::{authorize_mcp_server, is_authenticated},
|
||||
config::{
|
||||
add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, save_mcp_servers,
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum McpCommand {
|
||||
/// Add an MCP server
|
||||
Add {
|
||||
/// Server name (e.g., "notion", "github")
|
||||
name: String,
|
||||
|
||||
/// Server URL (e.g., "https://mcp.notion.com")
|
||||
url: String,
|
||||
|
||||
/// OAuth client ID (if authentication is required)
|
||||
#[arg(long)]
|
||||
client_id: Option<String>,
|
||||
|
||||
/// OAuth authorization URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
auth_url: Option<String>,
|
||||
|
||||
/// OAuth token URL (optional, can be discovered)
|
||||
#[arg(long)]
|
||||
token_url: Option<String>,
|
||||
|
||||
/// Scopes to request (comma-separated)
|
||||
#[arg(long)]
|
||||
scopes: Option<String>,
|
||||
|
||||
/// Server description
|
||||
#[arg(long)]
|
||||
description: Option<String>,
|
||||
},
|
||||
|
||||
/// Remove an MCP server
|
||||
Remove {
|
||||
/// Server name to remove
|
||||
name: String,
|
||||
},
|
||||
|
||||
/// List configured MCP servers
|
||||
List {
|
||||
/// Show detailed information
|
||||
#[arg(short, long)]
|
||||
verbose: bool,
|
||||
},
|
||||
|
||||
/// Authenticate with an MCP server (OAuth flow)
|
||||
Auth {
|
||||
/// Server name to authenticate
|
||||
name: String,
|
||||
|
||||
/// User ID for storing the token (default: "default")
|
||||
#[arg(short, long, default_value = "default")]
|
||||
user: String,
|
||||
},
|
||||
|
||||
/// Test connection to an MCP server
|
||||
Test {
|
||||
/// Server name to test
|
||||
name: String,
|
||||
|
||||
/// User ID for authentication (default: "default")
|
||||
#[arg(short, long, default_value = "default")]
|
||||
user: String,
|
||||
},
|
||||
|
||||
/// Enable or disable an MCP server
|
||||
Toggle {
|
||||
/// Server name
|
||||
name: String,
|
||||
|
||||
/// Enable the server
|
||||
#[arg(long, conflicts_with = "disable")]
|
||||
enable: bool,
|
||||
|
||||
/// Disable the server
|
||||
#[arg(long, conflicts_with = "enable")]
|
||||
disable: bool,
|
||||
},
|
||||
}
|
||||
|
||||
/// Run an MCP command.
|
||||
pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
|
||||
match cmd {
|
||||
McpCommand::Add {
|
||||
name,
|
||||
url,
|
||||
client_id,
|
||||
auth_url,
|
||||
token_url,
|
||||
scopes,
|
||||
description,
|
||||
} => {
|
||||
add_server(
|
||||
name,
|
||||
url,
|
||||
client_id,
|
||||
auth_url,
|
||||
token_url,
|
||||
scopes,
|
||||
description,
|
||||
)
|
||||
.await
|
||||
}
|
||||
McpCommand::Remove { name } => remove_server(name).await,
|
||||
McpCommand::List { verbose } => list_servers(verbose).await,
|
||||
McpCommand::Auth { name, user } => auth_server(name, user).await,
|
||||
McpCommand::Test { name, user } => test_server(name, user).await,
|
||||
McpCommand::Toggle {
|
||||
name,
|
||||
enable,
|
||||
disable,
|
||||
} => toggle_server(name, enable, disable).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a new MCP server.
|
||||
async fn add_server(
|
||||
name: String,
|
||||
url: String,
|
||||
client_id: Option<String>,
|
||||
auth_url: Option<String>,
|
||||
token_url: Option<String>,
|
||||
scopes: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut config = McpServerConfig::new(&name, &url);
|
||||
|
||||
if let Some(desc) = description {
|
||||
config = config.with_description(desc);
|
||||
}
|
||||
|
||||
// Track if auth is required
|
||||
let requires_auth = client_id.is_some();
|
||||
|
||||
// Set up OAuth if client_id is provided
|
||||
if let Some(client_id) = client_id {
|
||||
let mut oauth = OAuthConfig::new(client_id);
|
||||
|
||||
if let (Some(auth), Some(token)) = (auth_url, token_url) {
|
||||
oauth = oauth.with_endpoints(auth, token);
|
||||
}
|
||||
|
||||
if let Some(scopes_str) = scopes {
|
||||
let scope_list: Vec<String> = scopes_str
|
||||
.split(',')
|
||||
.map(|s| s.trim().to_string())
|
||||
.collect();
|
||||
oauth = oauth.with_scopes(scope_list);
|
||||
}
|
||||
|
||||
config = config.with_oauth(oauth);
|
||||
}
|
||||
|
||||
// Validate
|
||||
config.validate()?;
|
||||
|
||||
// Save
|
||||
add_mcp_server(config).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Added MCP server '{}'", name);
|
||||
println!(" URL: {}", url);
|
||||
|
||||
if requires_auth {
|
||||
println!();
|
||||
println!(" Run 'ironclaw mcp auth {}' to authenticate.", name);
|
||||
}
|
||||
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove an MCP server.
|
||||
async fn remove_server(name: String) -> anyhow::Result<()> {
|
||||
remove_mcp_server(&name).await?;
|
||||
|
||||
println!();
|
||||
println!(" ✓ Removed MCP server '{}'", name);
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List configured MCP servers.
|
||||
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
|
||||
let servers = load_mcp_servers().await?;
|
||||
|
||||
if servers.servers.is_empty() {
|
||||
println!();
|
||||
println!(" No MCP servers configured.");
|
||||
println!();
|
||||
println!(" Add a server with:");
|
||||
println!(" ironclaw mcp add <name> <url> [--client-id <id>]");
|
||||
println!();
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!();
|
||||
println!(" Configured MCP servers:");
|
||||
println!();
|
||||
|
||||
for server in &servers.servers {
|
||||
let status = if server.enabled { "●" } else { "○" };
|
||||
let auth_status = if server.requires_auth() {
|
||||
" (auth required)"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
if verbose {
|
||||
println!(" {} {}{}", status, server.name, auth_status);
|
||||
println!(" URL: {}", server.url);
|
||||
if let Some(ref desc) = server.description {
|
||||
println!(" Description: {}", desc);
|
||||
}
|
||||
if let Some(ref oauth) = server.oauth {
|
||||
println!(" OAuth Client ID: {}", oauth.client_id);
|
||||
if !oauth.scopes.is_empty() {
|
||||
println!(" Scopes: {}", oauth.scopes.join(", "));
|
||||
}
|
||||
}
|
||||
println!();
|
||||
} else {
|
||||
println!(
|
||||
" {} {} - {}{}",
|
||||
status, server.name, server.url, auth_status
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if !verbose {
|
||||
println!();
|
||||
println!(" Use --verbose for more details.");
|
||||
}
|
||||
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Authenticate with an MCP server.
|
||||
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let server = get_mcp_server(&name).await?;
|
||||
|
||||
// Initialize secrets store
|
||||
let secrets = get_secrets_store().await?;
|
||||
|
||||
// Check if already authenticated
|
||||
if is_authenticated(&server, &secrets, &user_id).await {
|
||||
println!();
|
||||
println!(" Server '{}' is already authenticated.", name);
|
||||
println!();
|
||||
print!(" Re-authenticate? [y/N]: ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
let mut input = String::new();
|
||||
std::io::stdin().read_line(&mut input)?;
|
||||
|
||||
if !input.trim().eq_ignore_ascii_case("y") {
|
||||
return Ok(());
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
println!();
|
||||
println!("╔════════════════════════════════════════════════════════════════╗");
|
||||
println!(
|
||||
"║ {:^62}║",
|
||||
format!("{} Authentication", name.to_uppercase())
|
||||
);
|
||||
println!("╚════════════════════════════════════════════════════════════════╝");
|
||||
println!();
|
||||
|
||||
// Perform OAuth flow (supports both pre-configured OAuth and DCR)
|
||||
match authorize_mcp_server(&server, &secrets, &user_id).await {
|
||||
Ok(_token) => {
|
||||
println!();
|
||||
println!(" ✓ Successfully authenticated with '{}'!", name);
|
||||
println!();
|
||||
println!(" You can now use tools from this server.");
|
||||
println!();
|
||||
}
|
||||
Err(crate::tools::mcp::auth::AuthError::NotSupported) => {
|
||||
println!();
|
||||
println!(" ✗ Server does not support OAuth authentication.");
|
||||
println!();
|
||||
println!(" The server may require a different authentication method,");
|
||||
println!(" or you may need to configure OAuth manually:");
|
||||
println!();
|
||||
println!(" ironclaw mcp remove {}", name);
|
||||
println!(
|
||||
" ironclaw mcp add {} {} --client-id YOUR_CLIENT_ID",
|
||||
name, server.url
|
||||
);
|
||||
println!();
|
||||
}
|
||||
Err(e) => {
|
||||
println!();
|
||||
println!(" ✗ Authentication failed: {}", e);
|
||||
println!();
|
||||
return Err(e.into());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test connection to an MCP server.
|
||||
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
|
||||
// Get server config
|
||||
let server = get_mcp_server(&name).await?;
|
||||
|
||||
println!();
|
||||
println!(" Testing connection to '{}'...", name);
|
||||
|
||||
// Create client
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
|
||||
// Always check for stored tokens (from either pre-configured OAuth or DCR)
|
||||
let secrets = get_secrets_store().await?;
|
||||
let has_tokens = is_authenticated(&server, &secrets, &user_id).await;
|
||||
|
||||
let client = if has_tokens {
|
||||
// We have stored tokens, use authenticated client
|
||||
McpClient::new_authenticated(server.clone(), session_manager, secrets, user_id)
|
||||
} else if server.requires_auth() {
|
||||
// OAuth configured but no tokens - need to authenticate
|
||||
println!();
|
||||
println!(
|
||||
" ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.",
|
||||
name
|
||||
);
|
||||
println!();
|
||||
return Ok(());
|
||||
} else {
|
||||
// No OAuth and no tokens - try unauthenticated
|
||||
McpClient::new_with_name(&server.name, &server.url)
|
||||
};
|
||||
|
||||
// Test connection
|
||||
match client.test_connection().await {
|
||||
Ok(()) => {
|
||||
println!(" ✓ Connection successful!");
|
||||
println!();
|
||||
|
||||
// List tools
|
||||
match client.list_tools().await {
|
||||
Ok(tools) => {
|
||||
println!(" Available tools ({}):", tools.len());
|
||||
for tool in tools {
|
||||
let approval = if tool.requires_approval() {
|
||||
" [approval required]"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
println!(" • {}{}", tool.name, approval);
|
||||
if !tool.description.is_empty() {
|
||||
// Truncate long descriptions
|
||||
let desc = if tool.description.len() > 60 {
|
||||
format!("{}...", &tool.description[..57])
|
||||
} else {
|
||||
tool.description.clone()
|
||||
};
|
||||
println!(" {}", desc);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" ✗ Failed to list tools: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
// Check if server requires auth but we don't have valid tokens
|
||||
if err_str.contains("401") || err_str.contains("requires authentication") {
|
||||
if has_tokens {
|
||||
// We had tokens but they failed - need to re-authenticate
|
||||
println!(
|
||||
" ✗ Authentication failed (token may be expired). Try re-authenticating:"
|
||||
);
|
||||
println!(" ironclaw mcp auth {}", name);
|
||||
} else {
|
||||
// No tokens - server requires auth
|
||||
println!(" ✗ Server requires authentication.");
|
||||
println!();
|
||||
println!(" Run 'ironclaw mcp auth {}' to authenticate.", name);
|
||||
}
|
||||
} else {
|
||||
println!(" ✗ Connection failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Toggle server enabled/disabled state.
|
||||
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
|
||||
let mut servers = load_mcp_servers().await?;
|
||||
|
||||
let server = servers
|
||||
.get_mut(&name)
|
||||
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
|
||||
|
||||
let new_state = if enable {
|
||||
true
|
||||
} else if disable {
|
||||
false
|
||||
} else {
|
||||
!server.enabled // Toggle if neither specified
|
||||
};
|
||||
|
||||
server.enabled = new_state;
|
||||
save_mcp_servers(&servers).await?;
|
||||
|
||||
let status = if new_state { "enabled" } else { "disabled" };
|
||||
println!();
|
||||
println!(" ✓ Server '{}' is now {}.", name, status);
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Initialize and return the secrets store.
|
||||
async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sync>> {
|
||||
let config = Config::from_env()?;
|
||||
|
||||
let master_key = config.secrets.master_key().ok_or_else(|| {
|
||||
anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env")
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
Ok(Arc::new(PostgresSecretsStore::new(
|
||||
store.pool(),
|
||||
Arc::new(crypto),
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_mcp_command_parsing() {
|
||||
// Just verify the command structure is valid
|
||||
use clap::CommandFactory;
|
||||
|
||||
// Create a dummy parent command to test subcommand parsing
|
||||
#[derive(clap::Parser)]
|
||||
struct TestCli {
|
||||
#[command(subcommand)]
|
||||
cmd: McpCommand,
|
||||
}
|
||||
|
||||
TestCli::command().debug_assert();
|
||||
}
|
||||
}
|
||||
@@ -5,11 +5,14 @@
|
||||
//! - Interactive setup wizard (`setup`)
|
||||
//! - Managing configuration (`config list`, `config get`, `config set`)
|
||||
//! - Managing WASM tools (`tool install`, `tool list`, `tool remove`)
|
||||
//! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`)
|
||||
|
||||
mod config;
|
||||
mod mcp;
|
||||
mod tool;
|
||||
|
||||
pub use config::{ConfigCommand, run_config_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use tool::{ToolCommand, run_tool_command};
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
@@ -72,6 +75,10 @@ pub enum Command {
|
||||
/// Manage WASM tools
|
||||
#[command(subcommand)]
|
||||
Tool(ToolCommand),
|
||||
|
||||
/// Manage MCP servers (hosted tool providers)
|
||||
#[command(subcommand)]
|
||||
Mcp(McpCommand),
|
||||
}
|
||||
|
||||
impl Cli {
|
||||
|
||||
+586
-1
@@ -1,13 +1,18 @@
|
||||
//! Tool management CLI commands.
|
||||
//!
|
||||
//! Commands for installing, listing, and removing WASM tools.
|
||||
//! Commands for installing, listing, removing, and authenticating WASM tools.
|
||||
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Command as ProcessCommand;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
use tokio::fs;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
|
||||
|
||||
/// Default tools directory.
|
||||
@@ -79,6 +84,20 @@ pub enum ToolCommand {
|
||||
#[arg(short, long)]
|
||||
dir: Option<PathBuf>,
|
||||
},
|
||||
|
||||
/// Configure authentication for a tool
|
||||
Auth {
|
||||
/// Name of the tool
|
||||
name: String,
|
||||
|
||||
/// Directory to look for tool (default: ~/.ironclaw/tools/)
|
||||
#[arg(short, long)]
|
||||
dir: Option<PathBuf>,
|
||||
|
||||
/// User ID for storing the secret (default: "default")
|
||||
#[arg(short, long, default_value = "default")]
|
||||
user: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Run a tool command.
|
||||
@@ -96,6 +115,7 @@ pub async fn run_tool_command(cmd: ToolCommand) -> anyhow::Result<()> {
|
||||
ToolCommand::List { dir, verbose } => list_tools(dir, verbose).await,
|
||||
ToolCommand::Remove { name, dir } => remove_tool(name, dir).await,
|
||||
ToolCommand::Info { name_or_path, dir } => show_tool_info(name_or_path, dir).await,
|
||||
ToolCommand::Auth { name, dir, user } => auth_tool(name, dir, user).await,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -658,6 +678,571 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
|
||||
}
|
||||
}
|
||||
|
||||
/// Configure authentication for a tool.
|
||||
async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyhow::Result<()> {
|
||||
let tools_dir = dir.unwrap_or_else(default_tools_dir);
|
||||
let caps_path = tools_dir.join(format!("{}.capabilities.json", name));
|
||||
|
||||
if !caps_path.exists() {
|
||||
anyhow::bail!(
|
||||
"Tool '{}' not found or has no capabilities file at {}",
|
||||
name,
|
||||
caps_path.display()
|
||||
);
|
||||
}
|
||||
|
||||
// Parse capabilities
|
||||
let content = fs::read_to_string(&caps_path).await?;
|
||||
let caps = CapabilitiesFile::from_json(&content)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid capabilities file: {}", e))?;
|
||||
|
||||
// Check for auth section
|
||||
let auth = caps.auth.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"Tool '{}' has no auth configuration.\n\
|
||||
The tool may not require authentication, or auth setup is not defined.",
|
||||
name
|
||||
)
|
||||
})?;
|
||||
|
||||
let display_name = auth.display_name.as_deref().unwrap_or(&name);
|
||||
|
||||
let header = format!("{} Authentication", display_name);
|
||||
println!();
|
||||
println!("╔════════════════════════════════════════════════════════════════╗");
|
||||
println!("║ {:^62}║", header);
|
||||
println!("╚════════════════════════════════════════════════════════════════╝");
|
||||
println!();
|
||||
|
||||
// Initialize secrets store
|
||||
let config = Config::from_env()?;
|
||||
let master_key = config.secrets.master_key().ok_or_else(|| {
|
||||
anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env")
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
|
||||
|
||||
// Check if already configured
|
||||
let already_configured = secrets_store
|
||||
.exists(&user_id, &auth.secret_name)
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if already_configured {
|
||||
println!(" {} is already configured.", display_name);
|
||||
println!();
|
||||
print!(" Replace existing credentials? [y/N]: ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
let mut input = String::new();
|
||||
std::io::stdin().read_line(&mut input)?;
|
||||
|
||||
if !input.trim().eq_ignore_ascii_case("y") {
|
||||
println!();
|
||||
println!(" Keeping existing credentials.");
|
||||
return Ok(());
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
// Check for environment variable
|
||||
if let Some(ref env_var) = auth.env_var {
|
||||
if let Ok(token) = std::env::var(env_var) {
|
||||
if !token.is_empty() {
|
||||
println!(" Found {} in environment.", env_var);
|
||||
println!();
|
||||
|
||||
// Validate if endpoint is provided
|
||||
if let Some(ref validation) = auth.validation_endpoint {
|
||||
print!(" Validating token...");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
match validate_token(&token, validation, &auth.secret_name).await {
|
||||
Ok(()) => {
|
||||
println!(" ✓");
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" ✗");
|
||||
println!(" Validation failed: {}", e);
|
||||
println!();
|
||||
println!(" Falling back to manual entry...");
|
||||
return auth_tool_manual(&secrets_store, &user_id, &auth).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(&secrets_store, &user_id, &auth, &token).await?;
|
||||
print_success(display_name);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for OAuth configuration
|
||||
if let Some(ref oauth) = auth.oauth {
|
||||
return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await;
|
||||
}
|
||||
|
||||
// Fall back to manual entry
|
||||
auth_tool_manual(&secrets_store, &user_id, &auth).await
|
||||
}
|
||||
|
||||
/// OAuth browser-based login flow.
|
||||
async fn auth_tool_oauth(
|
||||
store: &PostgresSecretsStore,
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
oauth: &crate::tools::wasm::OAuthConfigSchema,
|
||||
) -> anyhow::Result<()> {
|
||||
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use rand::RngCore;
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
|
||||
|
||||
// Get client_id from config or env
|
||||
let client_id = oauth
|
||||
.client_id
|
||||
.clone()
|
||||
.or_else(|| {
|
||||
oauth
|
||||
.client_id_env
|
||||
.as_ref()
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"OAuth client_id not configured.\n\
|
||||
Set it in the capabilities file or via environment variable."
|
||||
)
|
||||
})?;
|
||||
|
||||
// Get client_secret if provided
|
||||
let client_secret = oauth.client_secret.clone().or_else(|| {
|
||||
oauth
|
||||
.client_secret_env
|
||||
.as_ref()
|
||||
.and_then(|env| std::env::var(env).ok())
|
||||
});
|
||||
|
||||
println!(" Starting OAuth authentication...");
|
||||
println!();
|
||||
|
||||
// Find an available port for the callback
|
||||
let mut listener = None;
|
||||
let mut port = 0;
|
||||
|
||||
for p in 9876..=9886 {
|
||||
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
|
||||
Ok(l) => {
|
||||
listener = Some(l);
|
||||
port = p;
|
||||
break;
|
||||
}
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let listener = listener.ok_or_else(|| anyhow::anyhow!("Could not find available port"))?;
|
||||
let redirect_uri = format!("http://localhost:{}/callback", port);
|
||||
|
||||
// Generate PKCE verifier and challenge
|
||||
let (code_verifier, code_challenge) = if oauth.use_pkce {
|
||||
let mut verifier_bytes = [0u8; 32];
|
||||
rand::thread_rng().fill_bytes(&mut verifier_bytes);
|
||||
let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(verifier.as_bytes());
|
||||
let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize());
|
||||
|
||||
(Some(verifier), Some(challenge))
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
// Build authorization URL
|
||||
let mut auth_url = format!(
|
||||
"{}?client_id={}&response_type=code&redirect_uri={}",
|
||||
oauth.authorization_url,
|
||||
urlencoding::encode(&client_id),
|
||||
urlencoding::encode(&redirect_uri)
|
||||
);
|
||||
|
||||
if !oauth.scopes.is_empty() {
|
||||
auth_url.push_str(&format!(
|
||||
"&scope={}",
|
||||
urlencoding::encode(&oauth.scopes.join(" "))
|
||||
));
|
||||
}
|
||||
|
||||
if let Some(ref challenge) = code_challenge {
|
||||
auth_url.push_str(&format!(
|
||||
"&code_challenge={}&code_challenge_method=S256",
|
||||
challenge
|
||||
));
|
||||
}
|
||||
|
||||
// Add extra params
|
||||
for (key, value) in &oauth.extra_params {
|
||||
auth_url.push_str(&format!(
|
||||
"&{}={}",
|
||||
urlencoding::encode(key),
|
||||
urlencoding::encode(value)
|
||||
));
|
||||
}
|
||||
|
||||
println!(" Opening browser for {} login...", display_name);
|
||||
println!();
|
||||
|
||||
if let Err(e) = open::that(&auth_url) {
|
||||
println!(" Could not open browser: {}", e);
|
||||
println!(" Please open this URL manually:");
|
||||
println!(" {}", auth_url);
|
||||
}
|
||||
|
||||
println!(" Waiting for authorization...");
|
||||
|
||||
// Wait for callback with timeout
|
||||
let timeout = std::time::Duration::from_secs(300);
|
||||
let code = tokio::time::timeout(timeout, async {
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await?;
|
||||
|
||||
let mut reader = BufReader::new(&mut socket);
|
||||
let mut request_line = String::new();
|
||||
reader.read_line(&mut request_line).await?;
|
||||
|
||||
// Parse GET /callback?code=xxx HTTP/1.1
|
||||
if let Some(path) = request_line.split_whitespace().nth(1) {
|
||||
if path.starts_with("/callback") {
|
||||
if let Some(query) = path.split('?').nth(1) {
|
||||
for param in query.split('&') {
|
||||
let parts: Vec<&str> = param.splitn(2, '=').collect();
|
||||
if parts.len() == 2 && parts[0] == "code" {
|
||||
let code = urlencoding::decode(parts[1])
|
||||
.unwrap_or_else(|_| parts[1].into())
|
||||
.into_owned();
|
||||
|
||||
// Send success response
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\n\
|
||||
Content-Type: text/html\r\n\
|
||||
\r\n\
|
||||
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
|
||||
display: flex; justify-content: center; align-items: center; \
|
||||
height: 100vh; margin: 0; background: #191919; color: white;\">\
|
||||
<div style=\"text-align: center;\">\
|
||||
<h1>✓ {} Connected!</h1>\
|
||||
<p>You can close this window.</p>\
|
||||
</div></body></html>",
|
||||
display_name
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
let _ = socket.shutdown().await;
|
||||
|
||||
return Ok::<_, anyhow::Error>(code);
|
||||
}
|
||||
}
|
||||
|
||||
// Check for error
|
||||
if query.contains("error=") {
|
||||
let response =
|
||||
"HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
return Err(anyhow::anyhow!("Authorization denied by user"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let response = "HTTP/1.1 404 Not Found\r\n\r\n";
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Timed out waiting for authorization"))??;
|
||||
|
||||
println!();
|
||||
println!(" Exchanging code for token...");
|
||||
|
||||
// Exchange code for token
|
||||
let client = reqwest::Client::new();
|
||||
let mut token_params = vec![
|
||||
("grant_type", "authorization_code".to_string()),
|
||||
("code", code),
|
||||
("redirect_uri", redirect_uri),
|
||||
];
|
||||
|
||||
if let Some(ref verifier) = code_verifier {
|
||||
token_params.push(("code_verifier", verifier.to_string()));
|
||||
}
|
||||
|
||||
// Build token request
|
||||
let mut request = client.post(&oauth.token_url);
|
||||
|
||||
// Use Basic auth if client_secret is provided, otherwise include client_id in body
|
||||
if let Some(ref secret) = client_secret {
|
||||
request = request.basic_auth(&client_id, Some(secret));
|
||||
} else {
|
||||
token_params.push(("client_id", client_id));
|
||||
}
|
||||
|
||||
let token_response = request.form(&token_params).send().await?;
|
||||
|
||||
if !token_response.status().is_success() {
|
||||
let status = token_response.status();
|
||||
let body = token_response.text().await.unwrap_or_default();
|
||||
return Err(anyhow::anyhow!(
|
||||
"Token exchange failed: {} - {}",
|
||||
status,
|
||||
body
|
||||
));
|
||||
}
|
||||
|
||||
let token_data: serde_json::Value = token_response.json().await?;
|
||||
let access_token = token_data
|
||||
.get(&oauth.access_token_field)
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"No {} in token response: {:?}",
|
||||
oauth.access_token_field,
|
||||
token_data
|
||||
)
|
||||
})?;
|
||||
|
||||
// Save the token
|
||||
save_token(store, user_id, auth, access_token).await?;
|
||||
|
||||
// Extract any additional info for display
|
||||
let workspace_name = token_data
|
||||
.get("workspace_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| token_data.get("team_name").and_then(|v| v.as_str()));
|
||||
|
||||
println!();
|
||||
println!(" ✓ {} connected!", display_name);
|
||||
if let Some(workspace) = workspace_name {
|
||||
println!(" Workspace: {}", workspace);
|
||||
}
|
||||
println!();
|
||||
println!(" The tool can now access the API.");
|
||||
println!();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Manual token entry flow.
|
||||
async fn auth_tool_manual(
|
||||
store: &PostgresSecretsStore,
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
) -> anyhow::Result<()> {
|
||||
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
|
||||
|
||||
// Show instructions
|
||||
if let Some(ref instructions) = auth.instructions {
|
||||
println!(" Setup instructions:");
|
||||
println!();
|
||||
for line in instructions.lines() {
|
||||
println!(" {}", line);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
// Offer to open setup URL
|
||||
if let Some(ref url) = auth.setup_url {
|
||||
print!(" Press Enter to open setup page (or 's' to skip): ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
let mut input = String::new();
|
||||
std::io::stdin().read_line(&mut input)?;
|
||||
|
||||
if !input.trim().eq_ignore_ascii_case("s") {
|
||||
if let Err(e) = open::that(url) {
|
||||
println!(" Could not open browser: {}", e);
|
||||
println!(" Please open manually: {}", url);
|
||||
} else {
|
||||
println!(" Opening browser...");
|
||||
}
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
// Show token hint
|
||||
if let Some(ref hint) = auth.token_hint {
|
||||
println!(" Token format: {}", hint);
|
||||
println!();
|
||||
}
|
||||
|
||||
// Prompt for token
|
||||
print!(" Paste your token: ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
let token = read_hidden_input()?;
|
||||
println!();
|
||||
|
||||
if token.is_empty() {
|
||||
println!(" No token provided. Aborting.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Validate if endpoint is provided
|
||||
if let Some(ref validation) = auth.validation_endpoint {
|
||||
print!(" Validating token...");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
match validate_token(&token, validation, &auth.secret_name).await {
|
||||
Ok(()) => {
|
||||
println!(" ✓");
|
||||
}
|
||||
Err(e) => {
|
||||
println!(" ✗");
|
||||
println!(" Validation failed: {}", e);
|
||||
println!();
|
||||
print!(" Save anyway? [y/N]: ");
|
||||
std::io::stdout().flush()?;
|
||||
|
||||
let mut confirm = String::new();
|
||||
std::io::stdin().read_line(&mut confirm)?;
|
||||
|
||||
if !confirm.trim().eq_ignore_ascii_case("y") {
|
||||
println!(" Aborting.");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save the token
|
||||
save_token(store, user_id, auth, &token).await?;
|
||||
print_success(display_name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read input with hidden characters.
|
||||
fn read_hidden_input() -> anyhow::Result<String> {
|
||||
use crossterm::{
|
||||
event::{self, Event, KeyCode, KeyModifiers},
|
||||
terminal,
|
||||
};
|
||||
|
||||
let mut input = String::new();
|
||||
|
||||
terminal::enable_raw_mode()?;
|
||||
|
||||
loop {
|
||||
if let Event::Key(key_event) = event::read()? {
|
||||
match key_event.code {
|
||||
KeyCode::Enter => {
|
||||
break;
|
||||
}
|
||||
KeyCode::Backspace => {
|
||||
if !input.is_empty() {
|
||||
input.pop();
|
||||
print!("\x08 \x08");
|
||||
std::io::stdout().flush()?;
|
||||
}
|
||||
}
|
||||
KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
|
||||
terminal::disable_raw_mode()?;
|
||||
return Err(anyhow::anyhow!("Interrupted"));
|
||||
}
|
||||
KeyCode::Char(c) => {
|
||||
input.push(c);
|
||||
print!("*");
|
||||
std::io::stdout().flush()?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
terminal::disable_raw_mode()?;
|
||||
|
||||
Ok(input)
|
||||
}
|
||||
|
||||
/// Validate a token against the validation endpoint.
|
||||
async fn validate_token(
|
||||
token: &str,
|
||||
validation: &crate::tools::wasm::ValidationEndpointSchema,
|
||||
_secret_name: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()?;
|
||||
|
||||
// Build request based on method
|
||||
let request = match validation.method.to_uppercase().as_str() {
|
||||
"GET" => client.get(&validation.url),
|
||||
"POST" => client.post(&validation.url),
|
||||
_ => client.get(&validation.url),
|
||||
};
|
||||
|
||||
// Add authorization header (assume Bearer for now, could be extended)
|
||||
let response = request
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.header("Notion-Version", "2022-06-28") // Notion-specific, but harmless for others
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if response.status().as_u16() == validation.success_status {
|
||||
Ok(())
|
||||
} else {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
Err(anyhow::anyhow!(
|
||||
"HTTP {} (expected {}): {}",
|
||||
status,
|
||||
validation.success_status,
|
||||
if body.len() > 100 {
|
||||
format!("{}...", &body[..100])
|
||||
} else {
|
||||
body
|
||||
}
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Save token to secrets store.
|
||||
async fn save_token(
|
||||
store: &PostgresSecretsStore,
|
||||
user_id: &str,
|
||||
auth: &crate::tools::wasm::AuthCapabilitySchema,
|
||||
token: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut params = CreateSecretParams::new(&auth.secret_name, token);
|
||||
|
||||
if let Some(ref provider) = auth.provider {
|
||||
params = params.with_provider(provider);
|
||||
}
|
||||
|
||||
store
|
||||
.create(user_id, params)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Print success message.
|
||||
fn print_success(display_name: &str) {
|
||||
println!();
|
||||
println!(" ✓ {} connected!", display_name);
|
||||
println!();
|
||||
println!(" The tool can now access the API.");
|
||||
println!();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
Reference in New Issue
Block a user