newton_cli/commands/

newton_dashboard.rs

//! Newton Dashboard management module for authentication and policy secret management
//!
//! This module provides functionality to:
//! - Authenticate with Newton Dashboard via browser-based email + OTP login
//! - Store bearer tokens for subsequent API calls
//! - Encrypt secrets using RSA-OAEP with SHA-256 (compatible with AWS KMS)
//! - Upload encrypted secrets to the Newton Dashboard policy management API
//!
//! # Authentication Flow
//!
//! ## Step 1: Login to Dashboard
//!
//! First, authenticate with the Newton Dashboard to obtain and store a bearer token:
//!
//! ```bash
//! cargo run --bin newton-cli -- --chain-id 11155111 newton-dashboard login \
//!     --dashboard-url http://localhost:3000 \
//!     --api-base-url http://localhost:8000 \
//!     --poll-interval 3 \
//!     --max-poll-time 300
//! ```
//!
//! This command will:
//! 1. Generate a unique session ID
//! 2. Open your browser to the dashboard login page at `{dashboard-url}/cli-login?session_id={id}`
//! 3. Wait for you to complete email + OTP authentication in the browser
//! 4. Poll the backend API endpoint `/v1/authn/cli-token/{session_id}` every 3 seconds
//! 5. Once authenticated, save the bearer token to `~/.newton/newton-cli.toml`
//!
//! The bearer token will be automatically loaded for subsequent commands.
//!
//! ## Step 2: Upload Policy Secrets
//!
//! After logging in, you can upload encrypted secrets without providing the bearer token:
//!
//! ```bash
//! # Bearer token loaded automatically from ~/.newton/newton-cli.toml
//! cargo run --bin newton-cli -- --chain-id 11155111 newton-dashboard upload-policy-and-secret \
//!     --policy-name my-policy \
//!     --policy-address 0x1234567890abcdef1234567890abcdef12345678 \
//!     --policy-data-address 0xabcdef1234567890abcdef1234567890abcdef12 \
//!     --secret-key my-secret-key \
//!     --secret-value "my-secret-value"
//! ```
//!
//! Or provide the bearer token explicitly:
//!
//! ```bash
//! cargo run --bin newton-cli -- --chain-id 11155111 newton-dashboard upload-policy-and-secret \
//!     --api-base-url http://localhost:8000 \
//!     --policy-name my-policy \
//!     --policy-address 0x1234567890abcdef1234567890abcdef12345678 \
//!     --policy-data-address 0xabcdef1234567890abcdef1234567890abcdef12 \
//!     --secret-key my-secret-key \
//!     --secret-value "my-secret-value" \
//!     --bearer-token "your-bearer-token" \
//!     --skip-tls-verify
//! ```
//!
//! # Notes
//!
//! - The public key for secret encryption is hardcoded in the `KMS_PUBLIC_KEY_PEM` constant.
//!   Update it with your actual KMS public key before use.
//! - The bearer token and API base URL are stored in `~/.newton/newton-cli.toml`
//! - Use `--skip-tls-verify` flag if connecting to localhost with self-signed certificates.
//! - The `--bearer-token` and `--api-base-url` flags are optional if already set via login.

use base64::{engine::general_purpose, Engine as _};
use clap::{Parser, Subcommand};
use eyre::{Context, Result};
use newton_prover_core::config::NewtonAvsConfig;
use pkcs8::DecodePublicKey;
use rsa::{rand_core::OsRng, Oaep, RsaPublicKey};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use tracing::{self, info};

use crate::config::NewtonCliConfig;

/// Hardcoded KMS public key for encryption
/// Replace this with your actual KMS public key
const KMS_PUBLIC_KEY_PEM: &str = r#"-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA5aOjaOmoIiITeXSIe5Yt
9vmugRSahtf+MyMXEKB/kY96/UFIKCPGHXO7XF7RsccTgCYSi6JgTH2DiD1saazd
mJaHWMnz/UDY7/omFnW6RxGmGQxxXVI2bHgUhvf50JEculJvuUiljBXTUnofrfJI
LhZ2bAls3BBDDCxW8OFeJ3zQhJdOL9JP4AaYyiXlZZ1xpmDB3IEW+s9ma4524lqg
j/dmSVk40IB2K1EhaPj/zKs37nH2uOnLBkE2OWzQoEAVTdd6B1TLcFiGFrR9d4Xk
9RqHifqyaIlDWjUZ1lO9GX++2x6YIq9q+jDzidOzi43QLfBKjHELmM6xIXQ+Vwqa
WwIDAQAB
-----END PUBLIC KEY-----"#;

/// Encrypt data using RSA-OAEP with SHA-256
///
/// This function encrypts the provided data using RSA-OAEP padding with SHA-256,
/// which is compatible with AWS KMS RSA_2048 ENCRYPT_DECRYPT keys.
///
/// # Arguments
/// * `data` - The data to encrypt (as bytes)
/// * `public_key_pem` - The RSA public key in PEM format
///
/// # Returns
/// Base64-encoded ciphertext
pub fn encrypt_data(data: &[u8], public_key_pem: &str) -> Result<String> {
    // Parse the RSA public key from PEM (supports both PKCS#1 and PKCS#8 formats)
    let public_key =
        RsaPublicKey::from_public_key_pem(public_key_pem).context("Failed to parse RSA public key from PEM")?;

    // Encrypt using RSAES-OAEP with SHA-256
    let mut rng = OsRng;
    let padding = Oaep::new::<Sha256>();
    let ciphertext = public_key
        .encrypt(&mut rng, padding, data)
        .context("Failed to encrypt data with RSA-OAEP")?;

    // Encode to base64 for easy transport
    let b64 = general_purpose::STANDARD.encode(ciphertext);
    Ok(b64)
}

/// Policy creation response from the API
#[derive(Debug, Serialize, Deserialize)]
pub struct PolicyResponse {
    pub id: String,
}

/// Secret upload request
#[derive(Debug, Serialize, Deserialize)]
pub struct SecretRequest {
    pub key: String,
    pub value: String,
}

/// CLI token response from polling API
#[derive(Debug, Serialize, Deserialize)]
pub struct CliTokenResponse {
    pub token: Option<String>,
    pub status: String, // "pending", "completed", "expired"
}

/// Create a policy data via API
///
/// # Arguments
/// * `api_base_url` - Base URL of the API (e.g., "http://localhost:8000")
/// * `policy_name` - Name of the policy to create
/// * `policy_data_address` - Address of the policy data
/// * `bearer_token` - Bearer token for API authentication
/// * `verify_tls` - Whether to verify TLS certificates (set to false for localhost with self-signed certs)
///
/// # Returns
/// The policy ID
pub async fn create_policy_data(
    api_base_url: &str,
    policy_name: &str,
    policy_data_address: &str,
    bearer_token: &str,
    verify_tls: bool,
) -> Result<String> {
    let url = format!("{}/v1/policy-data", api_base_url.trim_end_matches('/'));

    let client = if verify_tls {
        reqwest::Client::new()
    } else {
        reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .context("Failed to create HTTP client")?
    };

    let response = client
        .post(&url)
        .header("Authorization", format!("Bearer {}", bearer_token))
        .json(&serde_json::json!({
            "name": policy_name,
            "address": policy_data_address
        }))
        .send()
        .await
        .with_context(|| format!("Failed to send request to {}", url))?;

    let status = response.status();
    if !status.is_success() {
        let error_text = response.text().await.unwrap_or_default();
        return Err(eyre::eyre!("API request failed with status {}: {}", status, error_text));
    }

    let policy_response: PolicyResponse = response.json().await.context("Failed to parse policy response")?;

    Ok(policy_response.id)
}

/// Create a policy call via API
///
/// # Arguments
/// * `api_base_url` - Base URL of the API (e.g., "http://localhost:8000")
/// * `policy_name` - Name of the policy to create
/// * `policy_address` - Address of the policy
/// * `bearer_token` - Bearer token for API authentication
/// * `verify_tls` - Whether to verify TLS certificates (set to false for localhost with self-signed certs)
///
/// # Returns
/// The policy ID
pub async fn create_policy_call(
    api_base_url: &str,
    policy_name: &str,
    policy_address: &str,
    bearer_token: &str,
    verify_tls: bool,
) -> Result<String> {
    let url = format!("{}/v1/policy", api_base_url.trim_end_matches('/'));

    let client = if verify_tls {
        reqwest::Client::new()
    } else {
        reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .context("Failed to create HTTP client")?
    };

    let response = client
        .post(&url)
        .header("Authorization", format!("Bearer {}", bearer_token))
        .json(&serde_json::json!({
            "name": policy_name,
            "address": policy_address
        }))
        .send()
        .await
        .with_context(|| format!("Failed to send request to {}", url))?;

    let status = response.status();
    if !status.is_success() {
        let error_text = response.text().await.unwrap_or_default();
        return Err(eyre::eyre!("API request failed with status {}: {}", status, error_text));
    }

    let policy_response: PolicyResponse = response.json().await.context("Failed to parse policy response")?;

    Ok(policy_response.id)
}

/// Upload a secret to a policy data via API
///
/// # Arguments
/// * `api_base_url` - Base URL of the API (e.g., "http://localhost:8000")
/// * `policy_data_id` - The policy data ID to upload the secret to
/// * `key` - The secret key
/// * `encrypted_value` - The encrypted secret value (base64-encoded)
/// * `bearer_token` - Bearer token for API authentication
/// * `verify_tls` - Whether to verify TLS certificates (set to false for localhost with self-signed certs)
///
/// # Returns
/// Success result
pub async fn upload_policy_data_secret(
    api_base_url: &str,
    policy_data_id: &str,
    key: &str,
    encrypted_value: &str,
    bearer_token: &str,
    verify_tls: bool,
) -> Result<()> {
    let url = format!(
        "{}/v1/policy-data/{}/secret",
        api_base_url.trim_end_matches('/'),
        policy_data_id
    );

    let client = if verify_tls {
        reqwest::Client::new()
    } else {
        reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .context("Failed to create HTTP client")?
    };

    let secret_request = SecretRequest {
        key: key.to_string(),
        value: encrypted_value.to_string(),
    };

    let response = client
        .post(&url)
        .header("Authorization", format!("Bearer {}", bearer_token))
        .json(&secret_request)
        .send()
        .await
        .with_context(|| format!("Failed to send request to {}", url))?;

    let status = response.status();
    if !status.is_success() {
        let error_text = response.text().await.unwrap_or_default();
        return Err(eyre::eyre!("API request failed with status {}: {}", status, error_text));
    }

    Ok(())
}

/// Newton Dashboard commands
#[derive(Debug, Parser)]
#[command(name = "newton-dashboard")]
pub struct NewtonDashboardCommand {
    #[command(subcommand)]
    pub subcommand: NewtonDashboardSubcommand,
}

#[derive(Debug, Subcommand)]
pub enum NewtonDashboardSubcommand {
    /// Login to Newton Dashboard and store bearer token
    Login(LoginCommand),
    /// Upload policy and secret to Newton Dashboard
    #[command(name = "upload-policy-and-secret")]
    UploadPolicyAndSecret(UploadPolicyAndSecretCommand),
}

/// Login command
#[derive(Debug, Parser)]
pub struct LoginCommand {
    /// Dashboard URL (e.g., "http://localhost:3000")
    #[arg(long, default_value = "http://localhost:3000")]
    dashboard_url: String,

    /// API base URL (e.g., "http://localhost:8000")
    #[arg(long, default_value = "http://localhost:8000")]
    api_base_url: String,

    /// Polling interval in seconds
    #[arg(long, default_value = "3")]
    poll_interval: u64,

    /// Maximum polling time in seconds
    #[arg(long, default_value = "300")]
    max_poll_time: u64,
}

/// Upload policy and secret command
#[derive(Debug, Parser)]
pub struct UploadPolicyAndSecretCommand {
    /// API base URL (e.g., "http://localhost:8000")
    #[arg(long)]
    api_base_url: Option<String>,

    /// Policy name to create
    #[arg(long)]
    policy_name: String,

    /// Policy address
    #[arg(long)]
    policy_address: String,

    /// Policy data address
    #[arg(long)]
    policy_data_address: String,

    /// Secret key name
    #[arg(long)]
    secret_key: String,

    /// Secret value to encrypt and upload
    #[arg(long)]
    secret_value: String,

    /// Bearer token for API authentication (optional if already logged in)
    #[arg(long)]
    bearer_token: Option<String>,

    /// Skip TLS certificate verification (useful for localhost with self-signed certs)
    #[arg(long, default_value = "false")]
    skip_tls_verify: bool,
}

/// Poll for CLI token from the backend API
async fn poll_for_token(
    api_base_url: &str,
    session_id: &str,
    poll_interval: u64,
    max_poll_time: u64,
) -> eyre::Result<String> {
    let url = format!(
        "{}/v1/authn/cli-token/{}",
        api_base_url.trim_end_matches('/'),
        session_id
    );

    let client = reqwest::Client::new();
    let start_time = std::time::Instant::now();
    let max_duration = std::time::Duration::from_secs(max_poll_time);
    let poll_interval_duration = std::time::Duration::from_secs(poll_interval);

    info!("Polling for token (timeout: {}s)...", max_poll_time);

    loop {
        if start_time.elapsed() > max_duration {
            return Err(eyre::eyre!(
                "Login timeout: no token received within {} seconds. Please try again.",
                max_poll_time
            ));
        }

        let response = client.get(&url).send().await;

        match response {
            Ok(resp) if resp.status().is_success() => {
                match resp.json::<CliTokenResponse>().await {
                    Ok(token_response) => {
                        match token_response.status.as_str() {
                            "completed" => {
                                if let Some(token) = token_response.token {
                                    return Ok(token);
                                } else {
                                    return Err(eyre::eyre!("Token marked as completed but no token provided"));
                                }
                            }
                            "expired" => {
                                return Err(eyre::eyre!("Login session expired. Please try again."));
                            }
                            "pending" => {
                                // Continue polling
                                info!("Waiting for login... ({:.0}s elapsed)", start_time.elapsed().as_secs());
                            }
                            status => {
                                tracing::warn!("Unknown status: {}", status);
                            }
                        }
                    }
                    Err(e) => {
                        tracing::warn!("Failed to parse token response: {}", e);
                    }
                }
            }
            Ok(resp) => {
                tracing::warn!("Polling failed with status: {}", resp.status());
            }
            Err(e) => {
                tracing::warn!("Polling request failed: {}", e);
            }
        }

        tokio::time::sleep(poll_interval_duration).await;
    }
}

/// Save the configuration to a TOML file
fn save_config(config: &NewtonCliConfig, config_path: &std::path::Path) -> eyre::Result<()> {
    let toml_string = toml::to_string_pretty(config).context("Failed to serialize config to TOML")?;

    // Create parent directories if they don't exist
    if let Some(parent) = config_path.parent() {
        std::fs::create_dir_all(parent).with_context(|| format!("Failed to create config directory: {:?}", parent))?;
    }

    std::fs::write(config_path, toml_string)
        .with_context(|| format!("Failed to write config file: {:?}", config_path))?;

    Ok(())
}

impl NewtonDashboardCommand {
    /// Execute the newton dashboard command
    pub async fn execute(self: Box<Self>, mut config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        match self.subcommand {
            NewtonDashboardSubcommand::Login(cmd) => {
                info!("Starting login flow for Newton Dashboard...");

                // Generate a session ID
                let session_id = uuid::Uuid::new_v4().to_string();

                // Construct the dashboard URL with session ID
                let dashboard_login_url = format!(
                    "{}/cli-login?session_id={}",
                    cmd.dashboard_url.trim_end_matches('/'),
                    session_id
                );

                info!("Opening browser to: {}", dashboard_login_url);
                info!("Please log in via the browser using your email and OTP...");

                // Open browser
                if let Err(e) = open::that(&dashboard_login_url) {
                    tracing::warn!("Failed to open browser automatically: {}", e);
                    info!("Please manually open this URL in your browser:");
                    info!("{}", dashboard_login_url);
                }

                // Poll for token
                let token =
                    poll_for_token(&cmd.api_base_url, &session_id, cmd.poll_interval, cmd.max_poll_time).await?;

                info!("Successfully received bearer token!");

                // Update config
                config.service.set_dashboard_bearer_token(Some(token));
                config
                    .service
                    .set_dashboard_api_base_url(Some(cmd.api_base_url.clone()));

                // Save config to default location
                let config_path = {
                    let home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string());
                    std::path::PathBuf::from(home).join(".newton").join("newton-cli.toml")
                };

                save_config(&config.service, &config_path)?;

                info!("Bearer token saved to config file: {:?}", config_path);
                info!("You can now use the 'upload-policy-and-secret' command without providing a bearer token.");

                Ok(())
            }
            NewtonDashboardSubcommand::UploadPolicyAndSecret(cmd) => {
                info!("Uploading policy and secret to Newton Dashboard...");

                // Get bearer token: from command line arg, or from config, or fail
                let bearer_token = cmd
                    .bearer_token
                    .or_else(|| config.service.dashboard_bearer_token.clone())
                    .ok_or_else(|| {
                        eyre::eyre!(
                        "Bearer token not found. Please provide --bearer-token or run 'newton-dashboard login' first."
                    )
                    })?;

                // Get API base URL: from command line arg, or from config, or use default
                let api_base_url = cmd
                    .api_base_url
                    .or_else(|| config.service.dashboard_api_base_url.clone())
                    .unwrap_or_else(|| "http://localhost:8000".to_string());

                // Encrypt the secret value using hardcoded public key
                info!("Encrypting secret value");
                let encrypted_value = encrypt_data(cmd.secret_value.as_bytes(), KMS_PUBLIC_KEY_PEM)
                    .with_context(|| "Failed to encrypt secret data")?;

                info!("Secret encrypted successfully");

                // Create policy call
                info!("Creating policy call: {}", cmd.policy_name);
                let policy_id = create_policy_call(
                    &api_base_url,
                    &cmd.policy_name,
                    &cmd.policy_address,
                    &bearer_token,
                    !cmd.skip_tls_verify,
                )
                .await
                .with_context(|| "Failed to create policy call")?;

                info!("Policy call created with ID: {}", policy_id);

                // Create policy data
                info!("Creating policy data: {}", cmd.policy_name);
                let policy_data_id = create_policy_data(
                    &api_base_url,
                    &cmd.policy_name,
                    &cmd.policy_data_address,
                    &bearer_token,
                    !cmd.skip_tls_verify,
                )
                .await
                .with_context(|| "Failed to create policy data")?;

                info!("Policy data created with ID: {}", policy_data_id);

                // Upload secret
                info!("Uploading secret with key: {}", cmd.secret_key);
                upload_policy_data_secret(
                    &api_base_url,
                    &policy_data_id,
                    &cmd.secret_key,
                    &encrypted_value,
                    &bearer_token,
                    !cmd.skip_tls_verify,
                )
                .await
                .with_context(|| "Failed to upload secret")?;

                info!("Secret uploaded successfully!");
                info!("Policy ID: {}", policy_id);
                info!("Policy data ID: {}", policy_data_id);
                info!("Secret Key: {}", cmd.secret_key);

                Ok(())
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // Test RSA public key (2048-bit) - this is a test key, not for production use
    const TEST_PUBLIC_KEY_PEM: &str = r#"-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAy8Dbv8prpJ/0kKhlGeJY
ozo2t60EG8L0561g13R29oVDmMytRDtpYRZxlYQbzhhJTV1bYgHyUwYj90fd6y0m
UHHJ9ZVKyj6HCdbO0nbt9Z4Khd9m5niJ2STgGUF0Q4Kq0lz4TgFYq9czvHBNc2Hq
dxFyPtkMDVDN7+auviBN+iHXj3pVKcQYBSyEbUGq+74Ceb5pdNCjBiWQrt40etnx
yU4aiUxE8SMBcgZnutK5UTjOSawafuqfBc9wfYTxEXk8qZ9s2qHI0wMDQN/POIoc
l9X5BWFX3ONjT9hCsBCdG7hm2XxKpnUbEUVyHNbVyHwIdJTCHNbRWwq7wym2MxYv
VwIDAQAB
-----END PUBLIC KEY-----"#;

    #[test]
    fn test_encrypt_data() {
        let data = b"test data";
        let result = encrypt_data(data, TEST_PUBLIC_KEY_PEM);
        assert!(result.is_ok());
        let encrypted = result.unwrap();
        // Base64 encoded ciphertext should be non-empty
        assert!(!encrypted.is_empty());
        // Base64 should be valid
        assert!(general_purpose::STANDARD.decode(&encrypted).is_ok());
    }

    #[test]
    fn test_encrypt_data_empty() {
        let data = b"";
        let result = encrypt_data(data, TEST_PUBLIC_KEY_PEM);
        assert!(result.is_ok());
    }

    #[test]
    fn test_encrypt_data_large() {
        // RSA-2048 with OAEP-SHA256 can encrypt up to ~190 bytes
        // Use 150 bytes to be safe
        let data = "x".repeat(150).into_bytes();
        let result = encrypt_data(&data, TEST_PUBLIC_KEY_PEM);
        assert!(result.is_ok());
    }

    #[test]
    fn test_encrypt_data_invalid_key() {
        let data = b"test data";
        let invalid_key = "invalid key";
        let result = encrypt_data(data, invalid_key);
        assert!(result.is_err());
    }

    #[test]
    fn test_secret_request_serialization() {
        let request = SecretRequest {
            key: "test_key".to_string(),
            value: "encrypted_value".to_string(),
        };
        let json = serde_json::to_string(&request).unwrap();
        assert!(json.contains("test_key"));
        assert!(json.contains("encrypted_value"));
    }

    #[test]
    fn test_policy_response_deserialization() {
        let json = r#"{"id": "policy-123"}"#;
        let response: PolicyResponse = serde_json::from_str(json).unwrap();
        assert_eq!(response.id, "policy-123");
    }
}