vault-hier/src/api.rs

use anyhow::Result;
use axum::{
    extract::{Multipart, Path, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::{get, post},
    Json, Router,
    Server,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tracing::{info, error, debug, instrument};

use crate::document_service::{Document, DocumentService, SignatureVerification};
use crate::vault_setup::VaultClient;

// API state containing services
#[derive(Clone)]
pub struct ApiState {
    document_service: Arc<DocumentService>,
    vault_client: Arc<VaultClient>,
}

// API error
#[derive(Debug)]
pub struct ApiError(anyhow::Error);

// Login request
#[derive(Debug, Deserialize)]
pub struct LoginRequest {
    username: String,
    password: String,
}

// Login response
#[derive(Debug, Serialize)]
pub struct LoginResponse {
    token: String,
}

// Sign document request
#[derive(Debug, Deserialize)]
pub struct SignDocumentRequest {
    username: String,
    token: String,
}

// API response implementations
impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        error!("API error: {}", self.0);
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Error: {}", self.0),
        )
            .into_response()
    }
}

impl<E> From<E> for ApiError
where
    E: Into<anyhow::Error>,
{
    fn from(err: E) -> Self {
        Self(err.into())
    }
}

// Start the API server
#[instrument(skip(vault_addr, root_token))]
pub async fn start_api(
    vault_addr: &str,
    root_token: &str,
    api_port: u16,
) -> Result<()> {
    info!("Starting API server on port {}...", api_port);

    // Initialize Vault client
    let vault_client = VaultClient::new(vault_addr, root_token);

    // Setup required secrets engines and auth methods
    vault_client.setup_secrets_engines().await?;

    // Setup 10 users in hierarchical structure
    vault_client.setup_hierarchical_users().await?;

    // Initialize document service
    let document_service = DocumentService::new(vault_client.clone());

    // Create API state
    let state = ApiState {
        document_service: Arc::new(document_service),
        vault_client: Arc::new(vault_client),
    };

    // Setup router
    let app = Router::new()
        .route("/health", get(health_check))
        .route("/api/login", post(login))
        .route("/api/documents", post(upload_document))
        .route("/api/documents/:id", get(get_document))
        .route("/api/documents/:id/sign", post(sign_document))
        .route("/api/documents/:id/verify", get(verify_document))
        .with_state(state);

    info!("API routes configured");

    // Get the socket address
    let addr = std::net::SocketAddr::from(([0, 0, 0, 0], api_port));

    // Bind and serve
    info!("Serving API at {}", addr);
    Server::bind(&addr)
        .serve(app.into_make_service())
        .await?;

    Ok(())
}

// Health check endpoint
#[instrument]
async fn health_check() -> &'static str {
    debug!("Health check endpoint called");
    "OK"
}

// Login endpoint
#[instrument(skip(state, request), fields(username = %request.username))]
async fn login(
    State(state): State<ApiState>,
    Json(request): Json<LoginRequest>,
) -> Result<Json<LoginResponse>, ApiError> {
    info!("Login attempt for user: {}", request.username);

    let token = state.vault_client
        .login_user(&request.username, &request.password)
        .await?;

    info!("User {} successfully authenticated", request.username);
    Ok(Json(LoginResponse { token }))
}

// Upload document endpoint
#[instrument(skip(state, multipart))]
async fn upload_document(
    State(state): State<ApiState>,
    mut multipart: Multipart,
) -> Result<Json<Document>, ApiError> {
    info!("Document upload request received");

    let mut document_name = String::new();
    let mut document_content = Vec::new();

    // Process multipart form
    while let Some(field) = multipart.next_field().await? {
        let name = field.name().unwrap_or("").to_string();

        if name == "name" {
            document_name = field.text().await?;
            debug!("Received document name: {}", document_name);
        } else if name == "file" {
            document_content = field.bytes().await?.to_vec();
            debug!("Received document content: {} bytes", document_content.len());
        }
    }

    if document_name.is_empty() || document_content.is_empty() {
        error!("Missing document name or content");
        return Err(anyhow::anyhow!("Missing document name or content").into());
    }

    // Upload document
    let document_id = state.document_service
        .upload_document(&document_name, &document_content)
        .await?;

    // Return document metadata
    let document = state.document_service
        .get_document(&document_id)
        .await?;

    info!("Document uploaded successfully with ID: {}", document_id);
    Ok(Json(document))
}

// Get document endpoint
#[instrument(skip(state))]
async fn get_document(
    State(state): State<ApiState>,
    Path(document_id): Path<String>,
) -> Result<Json<Document>, ApiError> {
    info!("Fetching document: {}", document_id);

    let document = state.document_service
        .get_document(&document_id)
        .await?;

    debug!("Retrieved document {} with {} signatures",
           document.id, document.signatures.len());

    Ok(Json(document))
}

// Sign document endpoint
#[instrument(skip(state, request), fields(document_id = %document_id, username = %request.username))]
async fn sign_document(
    State(state): State<ApiState>,
    Path(document_id): Path<String>,
    Json(request): Json<SignDocumentRequest>,
) -> Result<Json<Document>, ApiError> {
    info!("Signing request for document {} by user {}", document_id, request.username);

    state.document_service
        .sign_document(&document_id, &request.username, &request.token)
        .await?;

    let document = state.document_service
        .get_document(&document_id)
        .await?;

    info!("Document {} successfully signed by {}", document_id, request.username);
    Ok(Json(document))
}

// Verify document endpoint
#[instrument(skip(state))]
async fn verify_document(
    State(state): State<ApiState>,
    Path(document_id): Path<String>,
) -> Result<Json<SignatureVerification>, ApiError> {
    info!("Verifying document signatures: {}", document_id);

    let verification = state.document_service
        .verify_document_signatures(&document_id)
        .await?;

    info!("Document {} verification result: {}",
           document_id, if verification.is_verified { "VERIFIED" } else { "PENDING" });

    Ok(Json(verification))
}