northroot-proof-engine 0.3.2

//! Northroot Proof Engine CLI
//!
//! Command-line tool for demonstrating end-to-end workflows,
//! validating proofs, and benchmarking operations.

use clap::{Parser, Subcommand};
use colored::*;
use ed25519_dalek::SigningKey;
use northroot_proof_engine::{
    canonicalize_span_shape, compute_data_shape_hash, compute_identity_root,
    compute_method_shape_hash, compute_sha256, compute_span_commitment_hash,
    compute_span_shape_hash, compute_trace_hash, compute_trace_roots_hash, compute_trace_seq_root,
    compute_trace_set_root, normalize_timestamp, sign_span_commitment, validate_signature_format,
    verify_span_commitment_signature, ArtifactRef, MethodRef, MethodShapeManifest, RuntimePins,
    Signature, SpanCommitment, SpanShape, SpanStatus, Trace, TraceRoots,
};
use rand::rngs::OsRng;
use std::fs;
use std::time::Instant;

#[derive(Parser)]
#[command(name = "northroot-cli")]
#[command(about = "Northroot Proof Engine CLI - Demo and validation tool")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Demonstrate end-to-end PoSH → PoX workflow
    Demo {
        #[command(subcommand)]
        workflow: DemoWorkflow,
    },
    /// Benchmark hash and canonicalization operations
    Bench {
        #[command(subcommand)]
        bench_type: BenchType,
    },
    /// Validate proof JSON files
    Validate {
        /// Path to JSON file to validate
        file: String,
    },
}

#[derive(Subcommand)]
enum DemoWorkflow {
    /// Full PoSH → PoX workflow demonstration
    Workflow,
}

#[derive(Subcommand)]
enum BenchType {
    /// Benchmark hash operations
    Hash,
}

fn main() {
    let cli = Cli::parse();

    match cli.command {
        Commands::Demo { workflow } => match workflow {
            DemoWorkflow::Workflow => demo_workflow(),
        },
        Commands::Bench { bench_type } => match bench_type {
            BenchType::Hash => bench_hash(),
        },
        Commands::Validate { file } => validate_file(&file),
    }
}

fn demo_workflow() {
    println!("{}", "=== Northroot Proof Engine Demo ===\n".bold().cyan());

    // Step 1: Create Method Shape (PoSH)
    println!("{}", "Step 1: Creating Method Shape (PoSH)".bold());
    let span_shape = SpanShape {
        tool_id: "io.read_csv".to_string(),
        tool_version: "1".to_string(),
        param_schema: serde_json::json!({
            "delimiter": {"type": "string", "default": ","},
            "header": {"type": "boolean", "default": true}
        }),
        checker_ref: Some("chk:csv-valid@1".to_string()),
    };

    let span_shape_hash =
        compute_span_shape_hash(&span_shape).expect("Failed to compute span shape hash");
    println!("  Span Shape Hash: {}", span_shape_hash.green());

    // Create method shape manifest
    let mut csprng = OsRng;
    let _publisher_key = SigningKey::generate(&mut csprng);
    let publisher_kid = "did:key:zPublisher".to_string();

    let method_shape_manifest = MethodShapeManifest {
        schema_version: "method.v1".to_string(),
        method_id: "com.example/data-processor".to_string(),
        version: "1.0.0".to_string(),
        method_shape_root: span_shape_hash.clone(),
        method_dag_hash: None,
        runtime: RuntimePins {
            image_digest: "sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef"
                .to_string(),
            method_fingerprint: "git:abc123".to_string(),
        },
        checker_refs: vec!["chk:csv-valid@1".to_string()],
        default_policy_id: Some("pol:standard-v1".to_string()),
        publisher_sig: Signature {
            alg: "ed25519".to_string(),
            kid: publisher_kid.clone(),
            sig: "dGVzdHNpZ25hdHVyZQ".to_string(), // Placeholder
        },
    };

    let method_shape_hash = compute_method_shape_hash(&method_shape_manifest)
        .expect("Failed to compute method shape hash");
    println!("  Method Shape Hash: {}", method_shape_hash.green());
    println!();

    // Step 2: Create Span Commitment (PoX)
    println!("{}", "Step 2: Creating Span Commitment (PoX)".bold());
    let commitment = SpanCommitment {
        schema_version: "proof.v2".to_string(),
        trace_id: "tr_01H1234567890".to_string(),
        span_id: "sp-0001".to_string(),
        method_ref: MethodRef {
            method_id: method_shape_manifest.method_id.clone(),
            version: method_shape_manifest.version.clone(),
            method_shape_root: method_shape_hash.clone(),
            image_digest: method_shape_manifest.runtime.image_digest.clone(),
            method_fingerprint: method_shape_manifest.runtime.method_fingerprint.clone(),
        },
        data_shape_hash: compute_data_shape_hash(&serde_json::json!({"type": "object"}), None)
            .expect("Failed to compute data shape hash"),
        span_shape_hash: span_shape_hash.clone(),
        inputs: vec![ArtifactRef {
            name: "input.csv".to_string(),
            content_hash: "sha256:1111111111111111111111111111111111111111111111111111111111111111"
                .to_string(),
            uri: None,
        }],
        outputs: vec![ArtifactRef {
            name: "output.json".to_string(),
            content_hash: "sha256:2222222222222222222222222222222222222222222222222222222222222222"
                .to_string(),
            uri: None,
        }],
        start_time: normalize_timestamp("2025-11-06T12:00:00.000Z")
            .expect("Failed to normalize timestamp"),
        end_time: normalize_timestamp("2025-11-06T12:00:00.042Z")
            .expect("Failed to normalize timestamp"),
        status: SpanStatus::OK,
        meta: None,
    };

    let commitment_hash =
        compute_span_commitment_hash(&commitment).expect("Failed to compute span commitment hash");
    println!("  Span Commitment Hash: {}", commitment_hash.green());
    println!();

    // Step 3: Sign Span Commitment
    println!("{}", "Step 3: Signing Span Commitment".bold());
    let executor_key = SigningKey::generate(&mut csprng);
    let executor_kid = "did:key:zExecutor".to_string();

    let signature = sign_span_commitment(&commitment, &executor_key, executor_kid.clone())
        .expect("Failed to sign span commitment");
    println!("  Signature Algorithm: {}", signature.alg.green());
    println!("  Signature KID: {}", signature.kid.green());
    println!("  Signature: {}", signature.sig[..20].to_string() + "...");
    println!();

    // Step 4: Verify Signature
    println!("{}", "Step 4: Verifying Signature".bold());
    let verifying_key = executor_key.verifying_key();
    let verified =
        verify_span_commitment_signature(&commitment, &signature, verifying_key.as_bytes())
            .expect("Failed to verify signature");

    if verified {
        println!("  {} Signature verification: PASSED", "✓".green());
    } else {
        println!("  {} Signature verification: FAILED", "✗".red());
    }
    println!();

    // Step 5: Create Trace with Multiple Spans
    println!("{}", "Step 5: Creating Trace with Multiple Spans".bold());
    let trace = Trace {
        schema_version: "trace.v2".to_string(),
        trace_id: commitment.trace_id.clone(),
        method_ref: commitment.method_ref.clone(),
        data_shape_hash: commitment.data_shape_hash.clone(),
        policy_id: "pol:standard-v1".to_string(),
        identity_route: "did:key:zTenant".to_string(),
        content_route: "s3://bucket/trace".to_string(),
        execution_route: None,
    };

    let trace_hash = compute_trace_hash(&trace).expect("Failed to compute trace hash");
    println!("  Trace Hash: {}", trace_hash.green());
    println!();

    // Step 6: Compute Trace Roots
    println!("{}", "Step 6: Computing Trace Roots".bold());
    let span_hashes = vec![
        commitment_hash.clone(),
        "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string(),
    ];

    let trace_set_root =
        compute_trace_set_root(&span_hashes).expect("Failed to compute trace set root");
    println!("  Trace Set Root: {}", trace_set_root.green());

    let trace_seq_root =
        compute_trace_seq_root(&span_hashes).expect("Failed to compute trace seq root");
    println!("  Trace Seq Root: {}", trace_seq_root.green());

    let identity_root = compute_identity_root(&trace.trace_id, &trace.identity_route, &span_hashes)
        .expect("Failed to compute identity root");
    println!("  Identity Root: {}", identity_root.green());
    println!();

    // Summary
    println!("{}", "=== Summary ===".bold().cyan());
    println!("✓ Method Shape (PoSH) created and hashed");
    println!("✓ Span Commitment (PoX) created and hashed");
    println!("✓ Signature generated and verified");
    println!("✓ Trace created with multiple spans");
    println!("✓ Trace roots computed");
    println!();
    println!("{}", "Demo completed successfully!".green().bold());
}

fn bench_hash() {
    println!("{}", "=== Hash Operations Benchmark ===\n".bold().cyan());

    const ITERATIONS: usize = 10000;
    let test_data = "x".repeat(1000); // 1KB test data

    // Benchmark SHA-256 computation
    println!("Benchmarking SHA-256 computation...");
    let start = Instant::now();
    for _ in 0..ITERATIONS {
        let _ = compute_sha256(&test_data);
    }
    let duration = start.elapsed();
    let ops_per_sec = ITERATIONS as f64 / duration.as_secs_f64();
    println!("  {} iterations in {:?}", ITERATIONS, duration);
    println!(
        "  Throughput: {}\n",
        format!("{:.2} ops/sec", ops_per_sec).green()
    );

    // Benchmark canonicalization
    println!("Benchmarking canonicalization...");
    let span_shape = SpanShape {
        tool_id: "io.read_csv".to_string(),
        tool_version: "1".to_string(),
        param_schema: serde_json::json!({"key": "value"}),
        checker_ref: None,
    };

    let start = Instant::now();
    for _ in 0..ITERATIONS {
        let _ = canonicalize_span_shape(&span_shape);
    }
    let duration = start.elapsed();
    let ops_per_sec = ITERATIONS as f64 / duration.as_secs_f64();
    println!("  {} iterations in {:?}", ITERATIONS, duration);
    println!(
        "  Throughput: {}\n",
        format!("{:.2} ops/sec", ops_per_sec).green()
    );

    // Benchmark span commitment hash
    println!("Benchmarking span commitment hash...");
    let commitment = SpanCommitment {
        schema_version: "proof.v2".to_string(),
        trace_id: "tr_01H1234567890".to_string(),
        span_id: "sp-0001".to_string(),
        method_ref: MethodRef {
            method_id: "com.test/method".to_string(),
            version: "1.0.0".to_string(),
            method_shape_root:
                "sha256:0000000000000000000000000000000000000000000000000000000000000000"
                    .to_string(),
            image_digest: "sha256:0000000000000000000000000000000000000000000000000000000000000000"
                .to_string(),
            method_fingerprint: "git:deadbeef".to_string(),
        },
        data_shape_hash: "sha256:1111111111111111111111111111111111111111111111111111111111111111"
            .to_string(),
        span_shape_hash: "sha256:2222222222222222222222222222222222222222222222222222222222222222"
            .to_string(),
        inputs: vec![],
        outputs: vec![],
        start_time: "2025-11-06T12:00:00.000Z".to_string(),
        end_time: "2025-11-06T12:00:00.042Z".to_string(),
        status: SpanStatus::OK,
        meta: None,
    };

    let start = Instant::now();
    for _ in 0..ITERATIONS {
        let _ = compute_span_commitment_hash(&commitment);
    }
    let duration = start.elapsed();
    let ops_per_sec = ITERATIONS as f64 / duration.as_secs_f64();
    println!("  {} iterations in {:?}", ITERATIONS, duration);
    println!(
        "  Throughput: {}\n",
        format!("{:.2} ops/sec", ops_per_sec).green()
    );

    println!("{}", "Benchmark completed!".green().bold());
}

fn validate_file(file_path: &str) {
    println!("{}", "=== Validating Proof File ===\n".bold().cyan());
    println!("File: {}\n", file_path);

    let content = match fs::read_to_string(file_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("{} Failed to read file: {}", "✗".red(), e);
            std::process::exit(1);
        }
    };

    // Try to parse as various proof types
    let mut validated = false;

    // Try as MethodShapeManifest (PoSH)
    if let Ok(manifest) = serde_json::from_str::<MethodShapeManifest>(&content) {
        println!("{} Detected: Method Shape Manifest (PoSH)", "✓".green());
        println!("  Method ID: {}", manifest.method_id);
        println!("  Version: {}", manifest.version);

        match validate_signature_format(&manifest.publisher_sig) {
            Ok(_) => println!("  {} Signature format: Valid", "✓".green()),
            Err(e) => println!("  {} Signature format: Invalid - {}", "✗".red(), e),
        }

        match compute_method_shape_hash(&manifest) {
            Ok(hash) => println!("  Method Shape Hash: {}", hash),
            Err(e) => println!("  {} Hash computation failed: {}", "✗".red(), e),
        }

        validated = true;
    }

    // Try as SpanCommitment (PoX)
    if let Ok(commitment) = serde_json::from_str::<SpanCommitment>(&content) {
        println!("{} Detected: Span Commitment (PoX)", "✓".green());
        println!("  Trace ID: {}", commitment.trace_id);
        println!("  Span ID: {}", commitment.span_id);

        match compute_span_commitment_hash(&commitment) {
            Ok(hash) => println!("  Span Commitment Hash: {}", hash),
            Err(e) => println!("  {} Hash computation failed: {}", "✗".red(), e),
        }

        validated = true;
    }

    // Try as Trace (PoX)
    if let Ok(trace) = serde_json::from_str::<Trace>(&content) {
        println!("{} Detected: Trace (PoX)", "✓".green());
        println!("  Trace ID: {}", trace.trace_id);

        match compute_trace_hash(&trace) {
            Ok(hash) => println!("  Trace Hash: {}", hash),
            Err(e) => println!("  {} Hash computation failed: {}", "✗".red(), e),
        }

        validated = true;
    }

    // Try as TraceRoots (PoX)
    if let Ok(roots) = serde_json::from_str::<TraceRoots>(&content) {
        println!("{} Detected: Trace Roots (PoX)", "✓".green());
        println!("  Trace ID: {}", roots.trace_id);

        if let Some(ref sig) = roots.server_sig {
            match validate_signature_format(sig) {
                Ok(_) => println!("  {} Signature format: Valid", "✓".green()),
                Err(e) => println!("  {} Signature format: Invalid - {}", "✗".red(), e),
            }
        } else {
            println!("  {} No server signature present", "⚠".yellow());
        }

        match compute_trace_roots_hash(&roots) {
            Ok(hash) => println!("  Trace Roots Hash: {}", hash),
            Err(e) => println!("  {} Hash computation failed: {}", "✗".red(), e),
        }

        validated = true;
    }

    if !validated {
        eprintln!(
            "{} Could not identify proof type. File may be invalid JSON or unsupported type.",
            "✗".red()
        );
        std::process::exit(1);
    }

    println!();
    println!("{}", "Validation completed!".green().bold());
}