nushell-plugins/nu_plugin_kms/src/main.rs

//! Nushell plugin for KMS operations.
//!
//! This plugin provides Key Management System commands supporting multiple backends:
//! - RustyVault (primary - local Vault-compatible)
//! - Age (local file-based encryption)
//! - Cosmian (privacy-preserving)
//! - AWS KMS (production fallback)
//! - HashiCorp Vault
//!
//! Commands:
//! - `kms encrypt` - Encrypt data with selected backend
//! - `kms decrypt` - Decrypt data
//! - `kms generate-key` - Generate encryption keys
//! - `kms status` - Show backend status
//! - `kms list-backends` - List supported backends

use nu_plugin::{
    serve_plugin, EngineInterface, EvaluatedCall, MsgPackSerializer, Plugin, PluginCommand,
    SimplePluginCommand,
};
use nu_protocol::{record, Category, Example, LabeledError, Signature, SyntaxShape, Type, Value};

pub mod error;
mod helpers;

#[cfg(test)]
mod tests;

/// Nushell plugin for KMS operations.
#[derive(Debug)]
pub struct KmsPlugin;

impl Plugin for KmsPlugin {
    fn version(&self) -> String {
        env!("CARGO_PKG_VERSION").into()
    }

    fn commands(&self) -> Vec<Box<dyn PluginCommand<Plugin = Self>>> {
        vec![
            Box::new(KmsEncrypt),
            Box::new(KmsDecrypt),
            Box::new(KmsGenerateKey),
            Box::new(KmsStatus),
            Box::new(KmsListBackends),
        ]
    }
}

// =============================================================================
// Encrypt Command
// =============================================================================

/// Encrypt command - Encrypt data with selected backend.
#[derive(Debug)]
pub struct KmsEncrypt;

impl SimplePluginCommand for KmsEncrypt {
    type Plugin = KmsPlugin;

    fn name(&self) -> &str {
        "kms encrypt"
    }

    fn signature(&self) -> Signature {
        Signature::build(PluginCommand::name(self))
            .input_output_type(Type::String, Type::String)
            .required("data", SyntaxShape::String, "Data to encrypt")
            .named(
                "backend",
                SyntaxShape::String,
                "Backend: rustyvault, age, cosmian, aws, vault",
                Some('b'),
            )
            .named("key", SyntaxShape::String, "Key ID or recipient", Some('k'))
            .category(Category::Custom("provisioning".into()))
    }

    fn description(&self) -> &str {
        "Encrypt data using KMS backend"
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![
            Example {
                example: "kms encrypt \"secret data\" --backend rustyvault",
                description: "Encrypt with RustyVault",
                result: None,
            },
            Example {
                example: "kms encrypt \"data\" --backend age --key age1...",
                description: "Encrypt with Age recipient",
                result: None,
            },
            Example {
                example: "kms encrypt \"data\" --backend aws --key alias/my-key",
                description: "Encrypt with AWS KMS",
                result: None,
            },
        ]
    }

    fn run(
        &self,
        _plugin: &KmsPlugin,
        _engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        let data: String = call.req(0)?;
        let backend_name: Option<String> = call.get_flag("backend")?;
        let key: Option<String> = call.get_flag("key")?;

        // Create tokio runtime for async operations
        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| LabeledError::new(format!("Failed to create runtime: {}", e)))?;

        // Detect or use specified backend
        let backend = runtime.block_on(async {
            if let Some(name) = backend_name {
                match name.as_str() {
                    "rustyvault" => {
                        let addr = std::env::var("RUSTYVAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("RUSTYVAULT_TOKEN")
                            .map_err(|_| LabeledError::new("RUSTYVAULT_TOKEN not set"))?;
                        helpers::Backend::new_rustyvault(&addr, &token)
                            .await
                            .map_err(|e| LabeledError::new(e))
                    }
                    "age" => {
                        let recipient = key
                            .clone()
                            .ok_or_else(|| LabeledError::new("Age requires --key recipient"))?;
                        helpers::Backend::new_age(&recipient, None)
                            .map_err(|e| LabeledError::new(e))
                    }
                    "aws" => {
                        let key_id = key
                            .clone()
                            .ok_or_else(|| LabeledError::new("AWS KMS requires --key key-id"))?;
                        Ok(helpers::Backend::new_aws_kms(&key_id))
                    }
                    "vault" => {
                        let addr = std::env::var("VAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("VAULT_TOKEN")
                            .map_err(|_| LabeledError::new("VAULT_TOKEN not set"))?;
                        Ok(helpers::Backend::new_vault(&addr, &token))
                    }
                    backend @ _ => {
                        let url = std::env::var("KMS_HTTP_URL")
                            .unwrap_or_else(|_| "http://localhost:8081".to_string());
                        Ok(helpers::Backend::new_http_fallback(backend, &url))
                    }
                }
            } else {
                Ok(helpers::detect_backend().await)
            }
        })?;

        // Encrypt based on backend
        let encrypted = match backend {
            helpers::Backend::RustyVault { ref client } => {
                let key_name = key.unwrap_or_else(|| "provisioning-main".to_string());
                helpers::encrypt_rustyvault(client, &key_name, data.as_bytes())
                    .map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::Age { ref recipient, .. } => {
                helpers::encrypt_age(data.as_bytes(), recipient)
                    .map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::AwsKms { ref key_id } => runtime.block_on(async {
                helpers::encrypt_aws_kms(key_id, data.as_bytes())
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
            helpers::Backend::Vault {
                ref addr,
                ref token,
            } => {
                let key_name = key.unwrap_or_else(|| "provisioning-main".to_string());
                runtime.block_on(async {
                    helpers::encrypt_vault(addr, token, &key_name, data.as_bytes())
                        .await
                        .map_err(|e| LabeledError::new(e))
                })?
            }
            helpers::Backend::HttpFallback {
                ref backend_name,
                ref url,
            } => runtime.block_on(async {
                helpers::encrypt_http(url, backend_name, data.as_bytes())
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
        };

        Ok(Value::string(encrypted, call.head))
    }
}

// =============================================================================
// Decrypt Command
// =============================================================================

/// Decrypt command - Decrypt data using KMS backend.
#[derive(Debug)]
pub struct KmsDecrypt;

impl SimplePluginCommand for KmsDecrypt {
    type Plugin = KmsPlugin;

    fn name(&self) -> &str {
        "kms decrypt"
    }

    fn signature(&self) -> Signature {
        Signature::build(PluginCommand::name(self))
            .input_output_type(Type::String, Type::String)
            .required("encrypted", SyntaxShape::String, "Encrypted data")
            .named(
                "backend",
                SyntaxShape::String,
                "Backend: rustyvault, age, cosmian, aws, vault",
                Some('b'),
            )
            .named(
                "key",
                SyntaxShape::String,
                "Key ID or private key path",
                Some('k'),
            )
            .category(Category::Custom("provisioning".into()))
    }

    fn description(&self) -> &str {
        "Decrypt data using KMS backend"
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![
            Example {
                example: "kms decrypt $encrypted --backend rustyvault",
                description: "Decrypt with RustyVault",
                result: None,
            },
            Example {
                example: "kms decrypt $encrypted --backend age --key ~/.age/key.txt",
                description: "Decrypt with Age identity file",
                result: None,
            },
        ]
    }

    fn run(
        &self,
        _plugin: &KmsPlugin,
        _engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        let encrypted: String = call.req(0)?;
        let backend_name: Option<String> = call.get_flag("backend")?;
        let key: Option<String> = call.get_flag("key")?;

        // Create tokio runtime for async operations
        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| LabeledError::new(format!("Failed to create runtime: {}", e)))?;

        // Detect or use specified backend
        let backend = runtime.block_on(async {
            if let Some(name) = backend_name {
                match name.as_str() {
                    "rustyvault" => {
                        let addr = std::env::var("RUSTYVAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("RUSTYVAULT_TOKEN")
                            .map_err(|_| LabeledError::new("RUSTYVAULT_TOKEN not set"))?;
                        helpers::Backend::new_rustyvault(&addr, &token)
                            .await
                            .map_err(|e| LabeledError::new(e))
                    }
                    "age" => {
                        let identity = key
                            .clone()
                            .ok_or_else(|| LabeledError::new("Age requires --key identity_path"))?;
                        helpers::Backend::new_age("", Some(identity))
                            .map_err(|e| LabeledError::new(e))
                    }
                    "aws" => {
                        let key_id = key.clone();
                        Ok(helpers::Backend::new_aws_kms(&key_id.unwrap_or_default()))
                    }
                    "vault" => {
                        let addr = std::env::var("VAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("VAULT_TOKEN")
                            .map_err(|_| LabeledError::new("VAULT_TOKEN not set"))?;
                        Ok(helpers::Backend::new_vault(&addr, &token))
                    }
                    backend @ _ => {
                        let url = std::env::var("KMS_HTTP_URL")
                            .unwrap_or_else(|_| "http://localhost:8081".to_string());
                        Ok(helpers::Backend::new_http_fallback(backend, &url))
                    }
                }
            } else {
                Ok(helpers::detect_backend().await)
            }
        })?;

        // Decrypt based on backend
        let decrypted = match backend {
            helpers::Backend::RustyVault { ref client } => {
                let key_name = key.unwrap_or_else(|| "provisioning-main".to_string());
                helpers::decrypt_rustyvault(client, &key_name, &encrypted)
                    .map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::Age { ref identity, .. } => {
                let identity_path = identity.as_ref().ok_or_else(|| {
                    LabeledError::new("Age requires identity path for decryption")
                })?;
                helpers::decrypt_age(&encrypted, identity_path).map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::AwsKms { ref key_id } => runtime.block_on(async {
                helpers::decrypt_aws_kms(key_id, &encrypted)
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
            helpers::Backend::Vault {
                ref addr,
                ref token,
            } => {
                let key_name = key.unwrap_or_else(|| "provisioning-main".to_string());
                runtime.block_on(async {
                    helpers::decrypt_vault(addr, token, &key_name, &encrypted)
                        .await
                        .map_err(|e| LabeledError::new(e))
                })?
            }
            helpers::Backend::HttpFallback {
                ref backend_name,
                ref url,
            } => runtime.block_on(async {
                helpers::decrypt_http(url, backend_name, &encrypted)
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
        };

        // Convert bytes to string
        let plaintext = String::from_utf8(decrypted)
            .map_err(|e| LabeledError::new(format!("Failed to convert to UTF-8: {}", e)))?;

        Ok(Value::string(plaintext, call.head))
    }
}

// =============================================================================
// Generate Key Command
// =============================================================================

/// Generate data key command.
#[derive(Debug)]
pub struct KmsGenerateKey;

impl SimplePluginCommand for KmsGenerateKey {
    type Plugin = KmsPlugin;

    fn name(&self) -> &str {
        "kms generate-key"
    }

    fn signature(&self) -> Signature {
        Signature::build(PluginCommand::name(self))
            .input_output_type(Type::Nothing, Type::Record([].into()))
            .named(
                "spec",
                SyntaxShape::String,
                "Key spec: AES128, AES256",
                Some('s'),
            )
            .named("backend", SyntaxShape::String, "Backend", Some('b'))
            .category(Category::Custom("provisioning".into()))
    }

    fn description(&self) -> &str {
        "Generate data encryption key"
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![
            Example {
                example: "kms generate-key --spec AES256",
                description: "Generate AES-256 data key",
                result: None,
            },
            Example {
                example: "kms generate-key --backend age",
                description: "Generate Age key pair",
                result: None,
            },
        ]
    }

    fn run(
        &self,
        _plugin: &KmsPlugin,
        _engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        let key_spec: Option<String> = call.get_flag("spec")?;
        let backend_name: Option<String> = call.get_flag("backend")?;
        let key_spec = key_spec.unwrap_or_else(|| "AES256".to_string());

        // Create tokio runtime for async operations
        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| LabeledError::new(format!("Failed to create runtime: {}", e)))?;

        // Detect or use specified backend
        let backend = runtime.block_on(async {
            if let Some(name) = backend_name {
                match name.as_str() {
                    "rustyvault" => {
                        let addr = std::env::var("RUSTYVAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("RUSTYVAULT_TOKEN")
                            .map_err(|_| LabeledError::new("RUSTYVAULT_TOKEN not set"))?;
                        helpers::Backend::new_rustyvault(&addr, &token)
                            .await
                            .map_err(|e| LabeledError::new(e))
                    }
                    "age" => Ok(helpers::Backend::new_age("age1placeholder", None)
                        .map_err(|e| LabeledError::new(e))?),
                    "aws" => Ok(helpers::Backend::new_aws_kms("default")),
                    "vault" => {
                        let addr = std::env::var("VAULT_ADDR")
                            .unwrap_or_else(|_| "http://localhost:8200".to_string());
                        let token = std::env::var("VAULT_TOKEN")
                            .map_err(|_| LabeledError::new("VAULT_TOKEN not set"))?;
                        Ok(helpers::Backend::new_vault(&addr, &token))
                    }
                    backend @ _ => {
                        let url = std::env::var("KMS_HTTP_URL")
                            .unwrap_or_else(|_| "http://localhost:8081".to_string());
                        Ok(helpers::Backend::new_http_fallback(backend, &url))
                    }
                }
            } else {
                Ok(helpers::detect_backend().await)
            }
        })?;

        // Generate key based on backend
        let (plaintext, ciphertext) = match backend {
            helpers::Backend::RustyVault { ref client } => {
                let key_name = "provisioning-main";
                helpers::generate_data_key_rustyvault(client, key_name, &key_spec)
                    .map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::Age { .. } => {
                // Age generates key pairs, not data keys
                helpers::generate_age_key()
                    .map(|(secret, public)| (secret, public))
                    .map_err(|e| LabeledError::new(e))?
            }
            helpers::Backend::AwsKms { ref key_id } => runtime.block_on(async {
                helpers::generate_data_key_aws(key_id, &key_spec)
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
            helpers::Backend::Vault {
                ref addr,
                ref token,
            } => runtime.block_on(async {
                helpers::generate_data_key_vault(addr, token, "provisioning-main", &key_spec)
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
            helpers::Backend::HttpFallback {
                ref backend_name,
                ref url,
            } => runtime.block_on(async {
                helpers::generate_data_key_http(url, backend_name, &key_spec)
                    .await
                    .map_err(|e| LabeledError::new(e))
            })?,
        };

        Ok(Value::record(
            record! {
                "plaintext" => Value::string(plaintext, call.head),
                "ciphertext" => Value::string(ciphertext, call.head),
            },
            call.head,
        ))
    }
}

// =============================================================================
// Status Command
// =============================================================================

/// KMS status command.
#[derive(Debug)]
pub struct KmsStatus;

impl SimplePluginCommand for KmsStatus {
    type Plugin = KmsPlugin;

    fn name(&self) -> &str {
        "kms status"
    }

    fn signature(&self) -> Signature {
        Signature::build(PluginCommand::name(self))
            .input_output_type(Type::Nothing, Type::Record([].into()))
            .category(Category::Custom("provisioning".into()))
    }

    fn description(&self) -> &str {
        "Check KMS backend status and availability"
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![Example {
            example: "kms status",
            description: "Show current KMS backend status",
            result: None,
        }]
    }

    fn run(
        &self,
        _plugin: &KmsPlugin,
        _engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        // Create tokio runtime for async operations
        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| LabeledError::new(format!("Failed to create runtime: {}", e)))?;

        let backend = runtime.block_on(helpers::detect_backend());

        let (backend_type, available, config) = match backend {
            helpers::Backend::RustyVault { .. } => {
                let addr =
                    std::env::var("RUSTYVAULT_ADDR").unwrap_or_else(|_| "not set".to_string());
                ("rustyvault", true, format!("addr: {}", addr))
            }
            helpers::Backend::Age {
                ref recipient,
                ref identity,
            } => {
                let identity_status = identity.as_ref().map(|_| "set").unwrap_or("not set");
                (
                    "age",
                    true,
                    format!("recipient: {}, identity: {}", recipient, identity_status),
                )
            }
            helpers::Backend::AwsKms { ref key_id } => ("aws", true, format!("key_id: {}", key_id)),
            helpers::Backend::Vault { ref addr, .. } => ("vault", true, format!("addr: {}", addr)),
            helpers::Backend::HttpFallback {
                ref backend_name,
                ref url,
            } => (backend_name.as_str(), true, format!("url: {}", url)),
        };

        Ok(Value::record(
            record! {
                "backend" => Value::string(backend_type, call.head),
                "available" => Value::bool(available, call.head),
                "config" => Value::string(config, call.head),
            },
            call.head,
        ))
    }
}

// =============================================================================
// List Backends Command
// =============================================================================

/// List supported KMS backends command.
#[derive(Debug)]
pub struct KmsListBackends;

impl SimplePluginCommand for KmsListBackends {
    type Plugin = KmsPlugin;

    fn name(&self) -> &str {
        "kms list-backends"
    }

    fn signature(&self) -> Signature {
        Signature::build(PluginCommand::name(self))
            .input_output_type(Type::Nothing, Type::List(Box::new(Type::Record([].into()))))
            .category(Category::Custom("provisioning".into()))
    }

    fn description(&self) -> &str {
        "List all supported KMS backends and their availability"
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![Example {
            example: "kms list-backends",
            description: "Show all supported backends",
            result: None,
        }]
    }

    fn run(
        &self,
        _plugin: &KmsPlugin,
        _engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: &Value,
    ) -> Result<Value, LabeledError> {
        let backends = vec![
            (
                "rustyvault",
                "RustyVault Transit backend",
                "RUSTYVAULT_ADDR, RUSTYVAULT_TOKEN",
            ),
            (
                "age",
                "Age file-based encryption",
                "AGE_RECIPIENT, AGE_IDENTITY",
            ),
            (
                "aws",
                "AWS Key Management Service",
                "AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION",
            ),
            (
                "vault",
                "HashiCorp Vault Transit",
                "VAULT_ADDR, VAULT_TOKEN",
            ),
            (
                "cosmian",
                "Cosmian privacy-preserving encryption",
                "KMS_HTTP_URL",
            ),
        ];

        let backend_values: Vec<Value> = backends
            .iter()
            .map(|(name, description, env_vars)| {
                let available = helpers::check_backend_available(name);
                Value::record(
                    record! {
                        "name" => Value::string(*name, call.head),
                        "description" => Value::string(*description, call.head),
                        "available" => Value::bool(available, call.head),
                        "env_vars" => Value::string(*env_vars, call.head),
                    },
                    call.head,
                )
            })
            .collect();

        Ok(Value::list(backend_values, call.head))
    }
}

/// Entry point for the plugin binary.
fn main() {
    serve_plugin(&KmsPlugin, MsgPackSerializer);
}