lego-obs-mcp 0.1.0

use aws_sdk_s3::config::{BehaviorVersion, Region};
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::{Client, Config};
use rmcp::model::ServerInfo;
use rmcp::{ServerHandler, ServiceExt, tool};
use schemars::JsonSchema;
use serde::Deserialize;
use std::env;
use std::sync::Arc;
use tracing::info;

/// Reads an environment variable by key.
///
/// Returns `Ok(value)` if the variable exists, or an `Err` describing
/// which variable was not found.
pub fn get_env(key: &str) -> Result<String, String> {
    env::var(key).map_err(|_| format!("Missing required environment variable: {key}"))
}

/// Creates an OBS (Object Storage Service) S3-compatible client.
///
/// # Arguments
///
/// * `access_key` - The OBS access key ID.
/// * `secret_key` - The OBS secret access key.
/// * `bucket` - The target bucket name.
/// * `endpoint` - The OBS endpoint URL (e.g. `https://obs.cn-east-3.myhuaweicloud.com`).
/// * `region` - The OBS region identifier (e.g. `cn-east-3`).
///
/// # Returns
///
/// An AWS S3 SDK `Client` configured for the specified OBS endpoint.
pub fn create_client(
    access_key: &str,
    secret_key: &str,
    bucket: &str,
    endpoint: &str,
    region: &str,
) -> Client {
    info!(
        "Connecting to OBS: bucket={}, endpoint={}, region={}",
        bucket, endpoint, region
    );

    let region = Region::new(region.to_string());
    let shared_config = Config::builder()
        .behavior_version(BehaviorVersion::latest())
        .region(region)
        .endpoint_url(endpoint)
        .credentials_provider(aws_sdk_s3::config::Credentials::new(
            access_key, secret_key, None, None, "obs",
        ))
        .build();
    Client::from_conf(shared_config)
}

/// Checks whether a given object exists in the specified bucket.
///
/// Uses a `HeadObject` request, which is efficient because it only
/// retrieves metadata rather than the full object body.
pub async fn object_exists(client: &Client, bucket: &str, key: &str) -> bool {
    client
        .head_object()
        .bucket(bucket)
        .key(key)
        .send()
        .await
        .is_ok()
}

/// The core MCP server struct that holds the OBS client and bucket name.
#[derive(Debug, Clone)]
pub struct ObsServer {
    /// The S3-compatible OBS client.
    pub client: Arc<Client>,
    /// The OBS bucket to operate on.
    pub bucket: String,
}

/// Input parameters for the `list_objects` tool.
#[derive(Debug, Deserialize, JsonSchema)]
pub struct ListObjectsInput {
    /// Object key prefix for filtering results.
    /// Example: `lego/configurations/app/`
    pub prefix: Option<String>,
}

/// Input parameters for the `read_object` tool.
#[derive(Debug, Deserialize, JsonSchema)]
pub struct ReadObjectInput {
    /// The full object key path.
    /// Example: `lego/configurations/app/com.example.app/lego_configuration.json`
    pub key: String,
}

/// Input parameters for the `put_object` tool.
#[derive(Debug, Deserialize, JsonSchema)]
pub struct PutObjectInput {
    /// The full object key path to write to.
    pub key: String,
    /// The complete content to store in the object.
    pub content: String,
}

/// Input parameters for the `delete_object` tool.
#[derive(Debug, Deserialize, JsonSchema)]
pub struct DeleteObjectInput {
    /// The full object key path to delete.
    pub key: String,
}

#[tool(tool_box)]
impl ObsServer {
    /// Lists objects in the OBS bucket, optionally filtered by a prefix.
    #[tool(
        description = "List objects in the OBS bucket. Use prefix to filter, e.g. 'lego/configurations/' to list all configurations."
    )]
    pub async fn list_objects(&self, #[tool(aggr)] input: ListObjectsInput) -> String {
        let prefix = input.prefix.as_deref().unwrap_or("");
        match self
            .client
            .list_objects_v2()
            .bucket(&self.bucket)
            .prefix(prefix)
            .send()
            .await
        {
            Ok(resp) => {
                let objects = resp
                    .contents()
                    .iter()
                    .map(|o| {
                        let key = o.key().unwrap_or("");
                        let size = o.size().unwrap_or(0);
                        format!("{key} ({size} bytes)")
                    })
                    .collect::<Vec<_>>();
                if objects.is_empty() {
                    format!("No objects found with prefix '{prefix}'.")
                } else {
                    format!("Found {} object(s):\n{}", objects.len(), objects.join("\n"))
                }
            }
            Err(e) => format!("Failed to list objects: {e}"),
        }
    }

    /// Reads and returns the full content of a specified object.
    #[tool(description = "Read the full content of a specified object from OBS.")]
    pub async fn read_object(&self, #[tool(aggr)] input: ReadObjectInput) -> String {
        if !object_exists(&self.client, &self.bucket, &input.key).await {
            return format!("Object not found: {}", input.key);
        }
        match self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(&input.key)
            .send()
            .await
        {
            Ok(resp) => {
                let data = resp.body.collect().await;
                match data {
                    Ok(bytes) => {
                        let bytes_vec = bytes.into_bytes().to_vec();
                        let content = String::from_utf8_lossy(&bytes_vec);
                        format!("Content of {}:\n{}", input.key, content)
                    }
                    Err(e) => format!("Failed to read object content: {e}"),
                }
            }
            Err(e) => format!("Failed to read object: {e}"),
        }
    }

    /// Uploads a new object or overwrites an existing one.
    #[tool(
        description = "Upload or update an object in OBS. Overwrites if exists, creates if not."
    )]
    pub async fn put_object(&self, #[tool(aggr)] input: PutObjectInput) -> String {
        match self
            .client
            .put_object()
            .bucket(&self.bucket)
            .key(&input.key)
            .body(ByteStream::from(input.content.as_bytes().to_vec()))
            .content_type("application/json")
            .send()
            .await
        {
            Ok(_) => format!(
                "Successfully wrote: {}\nSize: {} bytes",
                input.key,
                input.content.len()
            ),
            Err(e) => format!("Failed to write object: {e}"),
        }
    }

    /// Deletes a specified object from the OBS bucket.
    #[tool(description = "Delete a specified object from OBS.")]
    pub async fn delete_object(&self, #[tool(aggr)] input: DeleteObjectInput) -> String {
        if !object_exists(&self.client, &self.bucket, &input.key).await {
            return format!("Object does not exist, nothing to delete: {}", input.key);
        }
        match self
            .client
            .delete_object()
            .bucket(&self.bucket)
            .key(&input.key)
            .send()
            .await
        {
            Ok(_) => format!("Successfully deleted: {}", input.key),
            Err(e) => format!("Failed to delete object: {e}"),
        }
    }
}

#[tool(tool_box)]
impl ServerHandler for ObsServer {
    /// Returns server metadata advertised to the MCP client during initialization.
    fn get_info(&self) -> ServerInfo {
        ServerInfo {
            instructions: Some("Lego OBS MCP Server - Manage configuration files on OBS. Supports listing, reading, uploading/updating, and deleting objects.".to_string()),
            ..Default::default()
        }
    }
}

#[tokio::main]
async fn main() {
    // Configure structured logging to stderr so stdout remains clean for the MCP protocol.
    tracing_subscriber::fmt()
        .with_env_filter("lego_obs_mcp=info")
        .with_writer(std::io::stderr)
        .init();

    info!("Starting lego-obs-mcp server");

    // Read all required environment variables and construct the OBS client.
    let access_key = get_env("LEGO_OBS_ACCESS_KEY").expect("LEGO_OBS_ACCESS_KEY is required");
    let secret_key = get_env("LEGO_OBS_SECRET_KEY").expect("LEGO_OBS_SECRET_KEY is required");
    let bucket = get_env("LEGO_OBS_BUCKET").expect("LEGO_OBS_BUCKET is required");
    let endpoint = get_env("LEGO_OBS_ENDPOINT").expect("LEGO_OBS_ENDPOINT is required");
    let region = get_env("LEGO_OBS_REGION").expect("LEGO_OBS_REGION is required");

    let client = create_client(&access_key, &secret_key, &bucket, &endpoint, &region);
    let server = ObsServer {
        client: Arc::new(client),
        bucket,
    };

    // Start the MCP server over stdio transport.
    let transport = rmcp::transport::io::stdio();
    let server_handle = server.serve(transport).await.unwrap();
    server_handle.waiting().await.unwrap();
}

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

    /// Verifies that `get_env` returns the correct value when the variable is set.
    #[test]
    fn test_get_env_returns_value() {
        unsafe { std::env::set_var("TEST_LEGO_ENV", "test_value") };
        let result = get_env("TEST_LEGO_ENV");
        assert_eq!(result, Ok("test_value".to_string()));
        unsafe { std::env::remove_var("TEST_LEGO_ENV") };
    }

    /// Verifies that `get_env` returns an error when the variable is not set.
    #[test]
    fn test_get_env_returns_error_when_missing() {
        let result = get_env("NONEXISTENT_LEGO_ENV_VAR_12345");
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(err.contains("NONEXISTENT_LEGO_ENV_VAR_12345"));
    }

    /// Verifies deserialization of `ListObjectsInput` with a prefix provided.
    #[test]
    fn test_list_objects_input_deserialize_with_prefix() {
        let json = r#"{"prefix": "lego/configurations/"}"#;
        let input: ListObjectsInput = serde_json::from_str(json).unwrap();
        assert_eq!(input.prefix, Some("lego/configurations/".to_string()));
    }

    /// Verifies deserialization of `ListObjectsInput` without a prefix.
    #[test]
    fn test_list_objects_input_deserialize_without_prefix() {
        let json = r#"{}"#;
        let input: ListObjectsInput = serde_json::from_str(json).unwrap();
        assert_eq!(input.prefix, None);
    }

    /// Verifies deserialization of `ReadObjectInput`.
    #[test]
    fn test_read_object_input_deserialize() {
        let json = r#"{"key": "lego/configurations/app/config.json"}"#;
        let input: ReadObjectInput = serde_json::from_str(json).unwrap();
        assert_eq!(input.key, "lego/configurations/app/config.json");
    }

    /// Verifies deserialization of `PutObjectInput`.
    #[test]
    fn test_put_object_input_deserialize() {
        let json = r#"{"key": "test.json", "content": "{}"}"#;
        let input: PutObjectInput = serde_json::from_str(json).unwrap();
        assert_eq!(input.key, "test.json");
        assert_eq!(input.content, "{}");
    }

    /// Verifies deserialization of `DeleteObjectInput`.
    #[test]
    fn test_delete_object_input_deserialize() {
        let json = r#"{"key": "lego/configurations/app/config.json"}"#;
        let input: DeleteObjectInput = serde_json::from_str(json).unwrap();
        assert_eq!(input.key, "lego/configurations/app/config.json");
    }

    /// Verifies that `ServerHandler::get_info` returns a valid instructions string.
    #[test]
    fn test_server_info_instructions() {
        let client = create_client("ak", "sk", "bucket", "https://example.com", "us-east-1");
        let server = ObsServer {
            client: Arc::new(client),
            bucket: "test-bucket".to_string(),
        };
        let info = server.get_info();
        assert!(info.instructions.is_some());
        let instructions = info.instructions.unwrap();
        assert!(instructions.contains("Lego OBS MCP Server"));
        assert!(instructions.contains("listing"));
        assert!(instructions.contains("reading"));
        assert!(instructions.contains("deleting"));
    }

    /// Verifies that `create_client` does not panic when given valid parameters.
    #[test]
    fn test_create_client_does_not_panic_with_valid_input() {
        let client = create_client("ak", "sk", "bucket", "https://example.com", "us-east-1");
        drop(client);
    }

    /// Verifies the default prefix behavior when `prefix` is `None`.
    #[test]
    fn test_list_objects_input_default_prefix() {
        let input = ListObjectsInput { prefix: None };
        let prefix = input.prefix.as_deref().unwrap_or("");
        assert_eq!(prefix, "");
    }

    /// Verifies that a custom prefix is correctly extracted.
    #[test]
    fn test_list_objects_input_custom_prefix() {
        let input = ListObjectsInput {
            prefix: Some("custom/prefix/".to_string()),
        };
        let prefix = input.prefix.as_deref().unwrap_or("");
        assert_eq!(prefix, "custom/prefix/");
    }

    /// Verifies the content length calculation for `PutObjectInput`.
    #[test]
    fn test_put_object_content_length() {
        let input = PutObjectInput {
            key: "test.json".to_string(),
            content: "{\"test\": true}".to_string(),
        };
        assert_eq!(input.content.len(), 14);
    }

    /// Verifies that `ObsServer` can be cloned correctly.
    #[test]
    fn test_obs_server_clone() {
        let client = create_client("ak", "sk", "bucket", "https://example.com", "us-east-1");
        let server = ObsServer {
            client: Arc::new(client),
            bucket: "test-bucket".to_string(),
        };
        let cloned = server.clone();
        assert_eq!(cloned.bucket, "test-bucket");
    }

    /// Verifies that input structs produce correct Debug output.
    #[test]
    fn test_input_structs_debug_format() {
        let list_input = ListObjectsInput {
            prefix: Some("test/".to_string()),
        };
        let debug_str = format!("{:?}", list_input);
        assert!(debug_str.contains("ListObjectsInput"));
        assert!(debug_str.contains("test/"));

        let read_input = ReadObjectInput {
            key: "key.json".to_string(),
        };
        let debug_str = format!("{:?}", read_input);
        assert!(debug_str.contains("ReadObjectInput"));
        assert!(debug_str.contains("key.json"));
    }

    /// Verifies that `ServerHandler::get_info` returns instructions with expected keywords.
    #[test]
    fn test_server_info_has_instructions() {
        let client = create_client("ak", "sk", "bucket", "https://example.com", "us-east-1");
        let server = ObsServer {
            client: Arc::new(client),
            bucket: "test-bucket".to_string(),
        };
        let info = server.get_info();
        assert!(info.instructions.is_some());
        let instructions = info.instructions.unwrap();
        assert!(instructions.contains("Lego OBS MCP Server"));
        assert!(instructions.contains("listing"));
        assert!(instructions.contains("reading"));
        assert!(instructions.contains("deleting"));
    }
}