payload-loader 0.1.0

//! # Loader
//!
//! A secure payload loader that fetches encrypted blobs, decrypts them in-memory using AEAD,
//! and consumes the plaintext as JSON configuration. No code injection—just safe serialization.
//!
//! ## Features
//! - Flexible URL resolution (plain, reversed, or base64-encoded)
//! - AEAD encryption with authenticated data (AAD)
//! - In-memory decryption with optional disk output (with warnings)
//! - Integrity verification via SHA-256
//! - Dry-run mode for validation without decryption
//! - Local blob decryption support
//! - Comprehensive input validation
//! - Structured logging for debugging and operations

use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand};
use mal_dev_core::{
    container::ScfBlob,
    crypto::decrypt_blob,
    transform::{decode_base64_to_bytes, reverse_string},
    transport::{build_client, fetch_bytes},
};
use serde::Deserialize;
use sha2::{Digest, Sha256};
use std::{path::PathBuf, time::Instant};
use tokio::fs;
use tracing::{debug, info, warn};

mod config;
mod metrics;
mod security;

pub use config::{Config, Profile};
pub use metrics::{OperationMetrics, Timer};
pub use security::{RateLimiter, SensitiveData};

// Constants
const DEFAULT_AAD: &str = "scf-aad-v1";
const EXPECTED_KEY_LENGTH: usize = 32;
const SAMPLE_CONFIG_NAME: &str = "lab-config";
const SAMPLE_CONFIG_VERSION: &str = "1.0";

#[derive(Parser, Debug)]
#[command(name = "loader")]
#[command(about = "Safe loader.exe: fetch + decrypt in-memory + consume (no injection)", long_about = None)]
struct Cli {
    /// AEAD AAD must match what was used during encryption
    #[arg(long, default_value = DEFAULT_AAD)]
    aad: Option<String>,

    /// Load configuration from TOML/JSON file
    #[arg(long)]
    config: Option<PathBuf>,

    /// Configuration profile to use (if config file provided)
    #[arg(long)]
    profile: Option<String>,

    /// Increase verbosity via RUST_LOG, e.g. `RUST_LOG=debug`
    #[command(subcommand)]
    cmd: Commands,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Full pipeline: fetch -> transform -> decrypt (memory) -> consume (parse JSON)
    Run {
        #[arg(long)]
        url: Option<String>,

        /// URL stored reversed (runtime reconstructed)
        #[arg(long)]
        url_rev: Option<String>,

        /// URL stored base64 (runtime decoded)
        #[arg(long)]
        url_b64: Option<String>,

        /// 32-byte key (base64)
        #[arg(long, default_value = "")]
        key_b64: String,

        /// Optional: write decrypted bytes to disk (defaults to memory-only)
        #[arg(long)]
        out: Option<PathBuf>,
    },

    /// Only fetch + parse container + hash ciphertext. No decrypt.
    DryRun {
        #[arg(long)]
        url: Option<String>,
        #[arg(long)]
        url_rev: Option<String>,
        #[arg(long)]
        url_b64: Option<String>,
    },

    /// Decrypt a local .scf blob and consume it (parse JSON)
    Decrypt {
        #[arg(long)]
        input: PathBuf,
        #[arg(long, default_value = "")]
        key_b64: String,
        #[arg(long)]
        out: Option<PathBuf>,
    },

    /// Generate a sample JSON payload file (safe) for lab testing
    GenJson {
        #[arg(long)]
        out: PathBuf,
    },
}

#[derive(Debug, Deserialize)]
struct PayloadConfig {
    name: String,
    version: String,
    features: Vec<String>,
    // Add fields as needed for your lab
}

/// Main entry point: parse CLI args and dispatch to appropriate command handler
#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env().add_directive("info".parse()?))
        .init();

    let mut cli = Cli::parse();

    // Load configuration file if provided
    let config_profile = if let Some(config_path) = &cli.config {
        debug!("loading config from {}", config_path.display());
        let config = if config_path.to_str().map_or(false, |s| s.ends_with(".json")) {
            Config::from_json_file(config_path).await?
        } else {
            Config::from_toml_file(config_path).await?
        };

        let profile = config.get_profile(cli.profile.as_deref());
        match profile {
            Some(p) => p,
            None => {
                let profile_name = cli.profile.as_deref().unwrap_or("default");
                bail!("profile '{}' not found in config file", profile_name);
            }
        }
    } else {
        Profile::default()
    };

    // Resolve AAD from config or CLI
    let aad_str = cli.aad.clone().or(config_profile.aad).unwrap_or_else(|| DEFAULT_AAD.to_string());
    let aad = aad_str.as_bytes();

    match cli.cmd {
        Commands::Run { url, url_rev, url_b64, mut key_b64, out } => {
            // Merge config with CLI args (CLI takes precedence)
            if key_b64.is_empty() && config_profile.key_b64.is_some() {
                key_b64 = config_profile.key_b64.unwrap();
            }
            if url.is_none() && config_profile.url.is_some() {
                handle_run(None, config_profile.url, None, key_b64, out.or(config_profile.output.map(PathBuf::from)), aad).await
            } else if url_rev.is_none() && config_profile.url_rev.is_some() {
                handle_run(None, None, config_profile.url_rev, key_b64, out.or(config_profile.output.map(PathBuf::from)), aad).await
            } else if url_b64.is_none() && config_profile.url_b64.is_some() {
                handle_run(None, None, config_profile.url_b64, key_b64, out.or(config_profile.output.map(PathBuf::from)), aad).await
            } else {
                handle_run(url, None, None, key_b64, out.or(config_profile.output.map(PathBuf::from)), aad).await
            }
        }

        Commands::DryRun { url, url_rev, url_b64 } => {
            let resolved_url = if let Some(u) = url {
                Some(u)
            } else if let Some(ur) = url_rev {
                Some(ur)
            } else if let Some(ub) = url_b64 {
                Some(ub)
            } else {
                config_profile.url.or(config_profile.url_rev).or(config_profile.url_b64)
            };

            if let Some(resolved) = resolved_url {
                handle_dry_run(Some(resolved), None, None).await
            } else {
                bail!("no URL provided via CLI or config")
            }
        }

        Commands::Decrypt { input, mut key_b64, out } => {
            if key_b64.is_empty() && config_profile.key_b64.is_some() {
                key_b64 = config_profile.key_b64.unwrap();
            }
            handle_decrypt(input, key_b64, out.or(config_profile.output.map(PathBuf::from)), aad).await
        }

        Commands::GenJson { out } => {
            handle_gen_json(out).await
        }
    }
}

/// Validate that a base64-encoded key decodes to exactly 32 bytes
fn validate_key(key_b64: &str) -> Result<Vec<u8>> {
    let key = decode_base64_to_bytes(key_b64).context("decode key_b64")?;
    if key.len() != EXPECTED_KEY_LENGTH {
        bail!(
            "key must be exactly {} bytes when decoded, got {}",
            EXPECTED_KEY_LENGTH,
            key.len()
        );
    }
    debug!("key validation passed: {} bytes", key.len());
    Ok(key)
}

/// Consume and log a JSON payload config from plaintext
fn consume_payload(pt: &[u8]) -> Result<PayloadConfig> {
    let cfg: PayloadConfig = serde_json::from_slice(pt).context("parse JSON payload config")?;
    info!("consume ok: name={} version={}", cfg.name, cfg.version);
    debug!("features: {:?}", cfg.features);
    Ok(cfg)
}

/// Handle the `Run` command: fetch -> decrypt -> consume
async fn handle_run(
    url: Option<String>,
    url_rev: Option<String>,
    url_b64: Option<String>,
    key_b64: String,
    out: Option<PathBuf>,
    aad: &[u8],
) -> Result<()> {
    let start = Instant::now();

    let resolved = resolve_url(url, url_rev, url_b64)?;
    debug!("resolved url: {}", &resolved);
    info!("resolved url ok");

    let client = build_client().context("build http client")?;
    let bytes = fetch_bytes(&client, &resolved).await.context("fetch blob")?;
    info!("downloaded {} bytes", bytes.len());

    let blob = ScfBlob::decode(&bytes).context("parse SCF container")?;
    info!(
        "container ok: version={} alg={} ct_len={}",
        blob.version, blob.alg, blob.ciphertext.len()
    );

    let key = validate_key(&key_b64)?;
    let pt = decrypt_blob(&blob, &key, aad).context("decrypt (AEAD)")?;
    info!("decrypted {} bytes in {:?}", pt.len(), start.elapsed());

    // Safe "consume": parse JSON config
    consume_payload(&pt)?;

    // Integrity signal for ops/debug: SHA-256 of plaintext
    let sha = Sha256::digest(&pt);
    info!("plaintext sha256={:x}", sha);

    if let Some(out_path) = out {
        fs::write(&out_path, &pt).await.context("write plaintext")?;
        warn!(
            "wrote plaintext to disk at {} (memory-only is safer)",
            out_path.display()
        );
    }

    Ok(())
}

/// Handle the `DryRun` command: fetch + parse + hash (no decrypt)
async fn handle_dry_run(
    url: Option<String>,
    url_rev: Option<String>,
    url_b64: Option<String>,
) -> Result<()> {
    let resolved = resolve_url(url, url_rev, url_b64)?;
    debug!("resolved url: {}", &resolved);

    let client = build_client().context("build http client")?;
    let bytes = fetch_bytes(&client, &resolved).await.context("fetch blob")?;
    info!("downloaded {} bytes", bytes.len());

    let blob = ScfBlob::decode(&bytes).context("parse SCF container")?;
    info!(
        "container ok: version={} alg={} nonce_len=12 ct_len={}",
        blob.version, blob.alg, blob.ciphertext.len()
    );

    let ct_sha = Sha256::digest(&blob.ciphertext);
    info!("ciphertext sha256={:x}", ct_sha);

    Ok(())
}

/// Handle the `Decrypt` command: decrypt a local blob file
async fn handle_decrypt(
    input: PathBuf,
    key_b64: String,
    out: Option<PathBuf>,
    aad: &[u8],
) -> Result<()> {
    debug!("reading local blob from {}", input.display());
    let raw = fs::read(&input).await.context("read input")?;
    info!("read {} bytes from {}", raw.len(), input.display());

    let blob = ScfBlob::decode(&raw).context("parse SCF container")?;
    info!(
        "container ok: version={} alg={} ct_len={}",
        blob.version, blob.alg, blob.ciphertext.len()
    );

    let key = validate_key(&key_b64)?;
    let pt = decrypt_blob(&blob, &key, aad).context("decrypt (AEAD)")?;
    info!("decrypted {} bytes", pt.len());

    consume_payload(&pt)?;

    let sha = Sha256::digest(&pt);
    info!("plaintext sha256={:x}", sha);

    if let Some(out_path) = out {
        fs::write(&out_path, &pt).await.context("write plaintext")?;
        warn!(
            "wrote plaintext to disk at {} (memory-only is safer)",
            out_path.display()
        );
    } else {
        warn!("no --out provided; plaintext kept in memory only");
    }

    Ok(())
}

/// Handle the `GenJson` command: generate a sample JSON payload file
async fn handle_gen_json(out: PathBuf) -> Result<()> {
    let sample = serde_json::json!({
        "name": SAMPLE_CONFIG_NAME,
        "version": SAMPLE_CONFIG_VERSION,
        "features": ["telemetry", "memory-only", "aead"]
    });

    let bytes = serde_json::to_vec_pretty(&sample).context("serialize sample json")?;
    fs::write(&out, &bytes)
        .await
        .context("write json")?;
    info!(
        "wrote sample json to {} ({} bytes)",
        out.display(),
        bytes.len()
    );
    debug!("sample json: {:?}", sample);
    Ok(())
}

/// Resolve URL from one of three sources: plain, reversed, or base64-encoded
fn resolve_url(url: Option<String>, url_rev: Option<String>, url_b64: Option<String>) -> Result<String> {
    let provided = url.is_some() as u8 + url_rev.is_some() as u8 + url_b64.is_some() as u8;
    if provided != 1 {
        bail!("provide exactly one of --url, --url-rev, or --url-b64");
    }

    if let Some(u) = url {
        return Ok(u);
    }
    if let Some(r) = url_rev {
        return Ok(reverse_string(&r));
    }
    if let Some(b64) = url_b64 {
        let bytes = decode_base64_to_bytes(&b64).context("decode url_b64")?;
        let s = String::from_utf8(bytes).context("url_b64 is not valid utf-8")?;
        return Ok(s);
    }

    unreachable!()
}