vault-hier/src/main.rs

use anyhow::Result;
use clap::{Parser, Subcommand};
use std::path::PathBuf;
use tracing::{info};
use tracing_subscriber::{fmt, EnvFilter};

// Import our library
use vault_hier::{start_api, initialize_vault};

#[derive(Parser)]
#[command(author, version, about = "Hierarchical Document Signing with HashiCorp Vault")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Start the vault-hier server
    Server {
        /// Address of the Vault server
        #[arg(long, env = "VAULT_ADDR", default_value = "http://127.0.0.1:8200")]
        vault_addr: String,

        /// Port for the API server
        #[arg(long, env = "API_PORT", default_value_t = 3000)]
        api_port: u16,
    },

    /// Login to get an authentication token
    Login {
        /// Username for authentication
        #[arg(short, long)]
        username: String,

        /// Password for authentication
        #[arg(short, long)]
        password: String,

        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },

    /// Upload a document
    Upload {
        /// Name of the document
        #[arg(short, long)]
        name: String,

        /// Path to the document file
        #[arg(short, long)]
        file: PathBuf,

        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },

    /// Sign a document
    Sign {
        /// Document ID to sign
        #[arg(short, long)]
        document_id: String,

        /// Username for signing
        #[arg(short, long)]
        username: String,

        /// Authentication token
        #[arg(short, long)]
        token: String,

        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },

    /// Verify a document's signatures
    Verify {
        /// Document ID to verify
        #[arg(short, long)]
        document_id: String,

        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },

    /// List all documents
    List {
        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },

    /// Get document details
    Get {
        /// Document ID to retrieve
        #[arg(short, long)]
        document_id: String,

        /// API server address
        #[arg(long, default_value = "http://localhost:3000")]
        api_url: String,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    fmt()
        .with_env_filter(EnvFilter::from_default_env().add_directive("vault_hier=info".parse()?))
        .with_target(false)
        .init();

    let cli = Cli::parse();

    match cli.command {
        Commands::Server { vault_addr, api_port } => {
            run_server(&vault_addr, api_port).await?;
        },
        Commands::Login { username, password, api_url } => {
            login(&username, &password, &api_url).await?;
        },
        Commands::Upload { name, file, api_url } => {
            upload_document(&name, file, &api_url).await?;
        },
        Commands::Sign { document_id, username, token, api_url } => {
            sign_document(&document_id, &username, &token, &api_url).await?;
        },
        Commands::Verify { document_id, api_url } => {
            verify_document(&document_id, &api_url).await?;
        },
        Commands::List { api_url } => {
            list_documents(&api_url).await?;
        },
        Commands::Get { document_id, api_url } => {
            get_document(&document_id, &api_url).await?;
        },
    }

    Ok(())
}

async fn run_server(vault_addr: &str, api_port: u16) -> Result<()> {
    info!("Vault address: {}", vault_addr);
    info!("Connecting to Vault at: {}", vault_addr);

    // Initialize and unseal Vault, get the root token
    let root_token = initialize_vault(vault_addr).await?;

    info!("Starting hierarchical document signing API...");

    // Start the hierarchical signing API
    start_api(vault_addr, &root_token, api_port).await?;

    info!("API server shutdown. Exiting.");

    Ok(())
}

async fn login(username: &str, password: &str, api_url: &str) -> Result<()> {
    info!("Logging in as user: {}", username);

    let client = reqwest::Client::new();
    let response = client
        .post(&format!("{}/api/login", api_url))
        .json(&serde_json::json!({
            "username": username,
            "password": password,
        }))
        .send()
        .await?;

    if response.status().is_success() {
        let data: serde_json::Value = response.json().await?;
        println!("Login successful!");
        println!("Token: {}", data["token"]);
        println!("Department: {}", data["department"]);
    } else {
        let error_text = response.text().await?;
        println!("Login failed: {}", error_text);
    }

    Ok(())
}

async fn upload_document(name: &str, file_path: PathBuf, api_url: &str) -> Result<()> {
    info!("Uploading document: {}", name);

    let file_content = tokio::fs::read(&file_path).await?;
    let file_name = file_path.file_name().unwrap().to_string_lossy();

    let form = reqwest::multipart::Form::new()
        .text("name", name.to_string())
        .part("file", reqwest::multipart::Part::bytes(file_content)
            .file_name(file_name.to_string()));

    let client = reqwest::Client::new();
    let response = client
        .post(&format!("{}/api/documents", api_url))
        .multipart(form)
        .send()
        .await?;

    if response.status().is_success() {
        let data: serde_json::Value = response.json().await?;
        println!("Document uploaded successfully!");
        println!("Document ID: {}", data["id"]);
    } else {
        let error_text = response.text().await?;
        println!("Upload failed: {}", error_text);
    }

    Ok(())
}

async fn sign_document(document_id: &str, username: &str, token: &str, api_url: &str) -> Result<()> {
    info!("Signing document: {}", document_id);

    let client = reqwest::Client::new();
    let response = client
        .post(&format!("{}/api/documents/{}/sign", api_url, document_id))
        .json(&serde_json::json!({
            "username": username,
            "token": token,
        }))
        .send()
        .await?;

    if response.status().is_success() {
        let data: serde_json::Value = response.json().await?;
        println!("Document signed successfully!");
        println!("Signature ID: {}", data["signature_id"]);
    } else {
        let error_text = response.text().await?;
        println!("Signing failed: {}", error_text);
    }

    Ok(())
}

async fn verify_document(document_id: &str, api_url: &str) -> Result<()> {
    info!("Verifying document: {}", document_id);

    let client = reqwest::Client::new();
    let response = client
        .get(&format!("{}/api/documents/{}/verify", api_url, document_id))
        .send()
        .await?;

    if response.status().is_success() {
        let data: serde_json::Value = response.json().await?;
        println!("Verification result:");
        println!("  Valid: {}", data["valid"]);
        println!("  Total signatures: {}", data["signature_count"]);
        println!("  Departments represented: {}", data["departments_represented"]);

        if let Some(signatures) = data["signatures"].as_array() {
            println!("\nSignatures:");
            for (i, sig) in signatures.iter().enumerate() {
                println!("  {}. User: {}, Department: {}, Time: {}",
                    i+1,
                    sig["username"],
                    sig["department"],
                    sig["timestamp"]);
            }
        }
    } else {
        let error_text = response.text().await?;
        println!("Verification failed: {}", error_text);
    }

    Ok(())
}

async fn list_documents(api_url: &str) -> Result<()> {
    info!("Listing all documents");

    let client = reqwest::Client::new();
    let response = client
        .get(&format!("{}/api/documents", api_url))
        .send()
        .await?;

    if response.status().is_success() {
        let data: Vec<serde_json::Value> = response.json().await?;
        println!("Documents:");

        for (i, doc) in data.iter().enumerate() {
            println!("  {}. ID: {}, Name: {}", i+1, doc["id"], doc["name"]);
        }

        if data.is_empty() {
            println!("  No documents found.");
        }
    } else {
        let error_text = response.text().await?;
        println!("Failed to list documents: {}", error_text);
    }

    Ok(())
}

async fn get_document(document_id: &str, api_url: &str) -> Result<()> {
    info!("Getting document: {}", document_id);

    let client = reqwest::Client::new();
    let response = client
        .get(&format!("{}/api/documents/{}", api_url, document_id))
        .send()
        .await?;

    if response.status().is_success() {
        let doc: serde_json::Value = response.json().await?;
        println!("Document details:");
        println!("  ID: {}", doc["id"]);
        println!("  Name: {}", doc["name"]);
        println!("  Hash: {}", doc["hash"]);
        println!("  Created: {}", doc["created_at"]);

        if let Some(signatures) = doc["signatures"].as_array() {
            println!("\nSignatures:");
            for (i, sig) in signatures.iter().enumerate() {
                println!("  {}. User: {}, Department: {}, Time: {}",
                    i+1,
                    sig["username"],
                    sig["department"],
                    sig["timestamp"]);
            }

            if signatures.is_empty() {
                println!("  No signatures yet.");
            }
        }
    } else {
        let error_text = response.text().await?;
        println!("Failed to get document: {}", error_text);
    }

    Ok(())
}