newton_cli/commands/

policy_files.rs

use clap::{Parser, Subcommand};
use eyre::{Context, Result};
use newton_prover_core::config::NewtonAvsConfig;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use tracing;

use crate::config::NewtonCliConfig;

/// Policy files commands
#[derive(Debug, Parser)]
#[command(name = "policy-files")]
pub struct PolicyFilesCommand {
    #[command(subcommand)]
    pub subcommand: PolicyFilesSubcommand,
}

#[derive(Debug, Subcommand)]
pub enum PolicyFilesSubcommand {
    /// Generate CIDs for policy files
    #[command(name = "generate-cids")]
    GenerateCids(GenerateCidsCommand),
}

/// Generate CIDs command
///
/// Uploads policy files to IPFS via Pinata and generates a policy_cids.json file
/// containing all CIDs needed for on-chain policy and policy-data deployment.
///
/// NOTE: Pinata is the currently supported IPFS hosting backend, but the Newton
/// protocol itself is IPFS-backend agnostic. Any pinning service that produces
/// valid IPFS CIDs can be used.
#[derive(Debug, Parser)]
pub struct GenerateCidsCommand {
    /// Directory containing policy files (defaults to policy-files)
    #[arg(short, long, default_value = "policy-files")]
    directory: PathBuf,

    /// Entrypoint (e.g., "newton_trading_agent.allow")
    #[arg(long)]
    entrypoint: String,

    /// Chain ID for attester resolution from config
    #[arg(long)]
    chain_id: u64,

    /// Deployment environment for config resolution (reads newton_prover_config.{env}.json)
    #[arg(long, default_value = "prod")]
    env: String,

    /// Optional path to secrets schema JSON file. If provided, uploads to IPFS
    /// and includes secretsSchemaCid in the output.
    #[arg(long)]
    secrets_schema_file: Option<PathBuf>,

    /// Pinata JWT for IPFS uploads (or set PINATA_JWT env var).
    /// Pinata is one supported IPFS backend; the protocol is backend-agnostic.
    #[arg(long, env = "PINATA_JWT")]
    pinata_jwt: Option<String>,

    /// Pinata Gateway URL for constructing IPFS links (or set PINATA_GATEWAY env var)
    #[arg(long, env = "PINATA_GATEWAY")]
    pinata_gateway: Option<String>,

    /// Output file path (defaults to policy-files/policy_cids.json)
    #[arg(short, long, default_value = "policy-files/policy_cids.json")]
    output: PathBuf,
}

#[derive(Debug, Serialize, Deserialize)]
#[allow(non_snake_case)]
struct PolicyCids {
    wasmCid: String,
    policyCid: String,
    schemaCid: String,
    attester: String,
    entrypoint: String,
    policyMetadataCid: String,
    policyDataMetadataCid: String,
    secretsSchemaCid: String,
}

#[derive(Debug)]
struct FileUpload {
    name: String,
    path: PathBuf,
    mime_type: String,
    default_name_prefix: String,
    header: String,
    json_key: String, // Key in the PolicyCids struct
}

/// Resolve the attester address from newton_prover_config.{env}.json.
///
/// Resolution order:
/// 1. `taskGenerator[0]` in the chain's config section
/// 2. `task_generator_addr` (legacy key) in the chain's config section
/// 3. Hardcoded per-chain fallback addresses
fn resolve_attester(chain_id: u64, env: &str) -> Result<String> {
    let config_path = format!("contracts/newton_prover_config.{}.json", env);
    let chain_key = chain_id.to_string();

    if let Ok(content) = std::fs::read_to_string(&config_path) {
        let config: serde_json::Value =
            serde_json::from_str(&content).with_context(|| format!("failed to parse {}", config_path))?;

        if let Some(chain_config) = config.get(&chain_key) {
            // Try taskGenerator[0]
            if let Some(addr) = chain_config
                .get("taskGenerator")
                .and_then(|v| v.as_array())
                .and_then(|arr| arr.first())
                .and_then(|v| v.as_str())
            {
                tracing::info!("Resolved attester from {}: taskGenerator[0] = {}", config_path, addr);
                return Ok(addr.to_string());
            }

            // Fallback: task_generator_addr
            if let Some(addr) = chain_config.get("task_generator_addr").and_then(|v| v.as_str()) {
                tracing::info!("Resolved attester from {}: task_generator_addr = {}", config_path, addr);
                return Ok(addr.to_string());
            }

            tracing::warn!(
                "Config {} has chain {} but no taskGenerator or task_generator_addr, using fallback",
                config_path,
                chain_key
            );
        } else {
            tracing::warn!(
                "Config {} has no entry for chain {}, using fallback",
                config_path,
                chain_key
            );
        }
    } else {
        tracing::warn!("Config file {} not found, using fallback attester", config_path);
    }

    let fallback = match chain_id {
        31337 => "0x70997970C51812dc3A010C7d01b50e0d17dc79C8",
        11155111 | 17000 => "0xD45062003a4626a532F30A4596aB253c45AE0647",
        1 => "0x4883282094755C01cd0d15dFE74753c9E189d194",
        _ => {
            return Err(eyre::eyre!(
                "no attester found in config and no hardcoded fallback for chain {}",
                chain_id
            ));
        }
    };

    tracing::info!("Using hardcoded fallback attester for chain {}: {}", chain_id, fallback);
    Ok(fallback.to_string())
}

/// Check that a required policy file exists in the directory.
fn require_file(dir: &Path, filename: &str) -> Result<()> {
    let path = dir.join(filename);
    if !path.exists() {
        eyre::bail!("{} not found in {:?}", filename, dir);
    }
    Ok(())
}

impl GenerateCidsCommand {
    /// Upload a file to Pinata IPFS
    pub async fn upload_file_to_pinata(
        file: PathBuf,
        jwt: String,
        file_name: String,
        mime_type: &str,
    ) -> Result<String> {
        // Check if file exists
        if !file.exists() {
            eyre::bail!("Error: file not found: {:?}", file);
        }

        // Read file
        let file_data = std::fs::read(&file).with_context(|| format!("Failed to read file: {:?}", file))?;

        // Get file name from path
        let file_name_from_path = file.file_name().and_then(|n| n.to_str()).unwrap_or("file");

        // Create multipart form
        let form = reqwest::multipart::Form::new()
            .text("pinataOptions", r#"{"cidVersion":1}"#)
            .text("pinataMetadata", format!(r#"{{"name":"{}"}}"#, file_name))
            .part(
                "file",
                reqwest::multipart::Part::bytes(file_data)
                    .file_name(file_name_from_path.to_string())
                    .mime_str(mime_type)
                    .unwrap(),
            );

        // Upload to Pinata
        let client = reqwest::Client::new();
        let response = client
            .post("https://api.pinata.cloud/pinning/pinFileToIPFS")
            .header("Authorization", format!("Bearer {}", jwt))
            .multipart(form)
            .send()
            .await
            .context("Failed to send request to Pinata")?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response.text().await.unwrap_or_default();
            eyre::bail!("Pinata upload failed with status {}: {}", status, error_text);
        }

        // Parse response
        let response_json: serde_json::Value = response.json().await.context("Failed to parse Pinata response")?;

        // Extract IPFS hash
        let ipfs_hash = response_json
            .get("IpfsHash")
            .and_then(|v| v.as_str())
            .ok_or_else(|| eyre::eyre!("IPFS hash not found in response"))?;

        Ok(ipfs_hash.to_string())
    }

    /// Display upload results
    pub fn display_upload_results(ipfs_hash: &str, gateway: &str) {
        tracing::info!("\n=== IPFS Upload Results ===");
        tracing::info!("IPFS Hash: {}", ipfs_hash);
        // Construct gateway links
        let gateway_link = if gateway.ends_with('/') {
            format!("{}{}", gateway, ipfs_hash)
        } else {
            format!("{}/{}", gateway, ipfs_hash)
        };

        tracing::info!(gateway_link = %gateway_link, "Direct IPFS Link");
        tracing::info!("Public IPFS Link: https://ipfs.io/ipfs/{}", ipfs_hash);
    }

    /// Execute the generate-cids command
    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> Result<()> {
        // Get Pinata credentials
        let jwt = self
            .pinata_jwt
            .ok_or_else(|| eyre::eyre!("Pinata JWT is required. Set PINATA_JWT env var or use --pinata-jwt"))?;

        let gateway = self.pinata_gateway.ok_or_else(|| {
            eyre::eyre!("Pinata Gateway is required. Set PINATA_GATEWAY env var or use --pinata-gateway")
        })?;

        // Validate required files exist before starting uploads
        for filename in &[
            "policy.wasm",
            "policy.rego",
            "params_schema.json",
            "policy_metadata.json",
            "policy_data_metadata.json",
        ] {
            require_file(&self.directory, filename)?;
        }

        if let Some(ref secrets_path) = self.secrets_schema_file {
            if !secrets_path.exists() {
                eyre::bail!("secrets schema file not found: {:?}", secrets_path);
            }
        }

        // Resolve attester from config
        let attester = resolve_attester(self.chain_id, &self.env)?;

        // Define all files to upload in the correct order
        let mut files_to_upload = vec![
            FileUpload {
                name: "policy.wasm".to_string(),
                path: self.directory.join("policy.wasm"),
                mime_type: "application/wasm".to_string(),
                default_name_prefix: "newton-policy-wasm".to_string(),
                header: "============ Upload policy.wasm ================".to_string(),
                json_key: "wasmCid".to_string(),
            },
            FileUpload {
                name: "policy.rego".to_string(),
                path: self.directory.join("policy.rego"),
                mime_type: "text/plain".to_string(),
                default_name_prefix: "newton-policy".to_string(),
                header: "============== Upload policy.rego ==============".to_string(),
                json_key: "policyCid".to_string(),
            },
            FileUpload {
                name: "params_schema.json".to_string(),
                path: self.directory.join("params_schema.json"),
                mime_type: "application/json".to_string(),
                default_name_prefix: "newton-policy-schema".to_string(),
                header: "========== Upload params_schema.json ===========".to_string(),
                json_key: "schemaCid".to_string(),
            },
            FileUpload {
                name: "policy_metadata.json".to_string(),
                path: self.directory.join("policy_metadata.json"),
                mime_type: "application/json".to_string(),
                default_name_prefix: "newton-policy-metadata".to_string(),
                header: "========== Upload policy_metadata.json =========".to_string(),
                json_key: "policyMetadataCid".to_string(),
            },
            FileUpload {
                name: "policy_data_metadata.json".to_string(),
                path: self.directory.join("policy_data_metadata.json"),
                mime_type: "application/json".to_string(),
                default_name_prefix: "newton-policy-data-metadata".to_string(),
                header: "======== Upload policy_data_metadata.json ======".to_string(),
                json_key: "policyDataMetadataCid".to_string(),
            },
        ];

        // Add secrets schema if provided
        if let Some(ref secrets_path) = self.secrets_schema_file {
            files_to_upload.push(FileUpload {
                name: secrets_path
                    .file_name()
                    .and_then(|n| n.to_str())
                    .unwrap_or("secrets_schema.json")
                    .to_string(),
                path: secrets_path.clone(),
                mime_type: "application/json".to_string(),
                default_name_prefix: "newton-secrets-schema".to_string(),
                header: "========== Upload secrets_schema.json ==========".to_string(),
                json_key: "secretsSchemaCid".to_string(),
            });
        }

        tracing::info!("================================================");
        tracing::info!("========== Uploading All Policy Files ==========");
        tracing::info!("================================================");
        tracing::info!("Directory: {:?}", self.directory);
        tracing::info!("Chain ID: {}", self.chain_id);
        tracing::info!("Env: {}", self.env);
        tracing::info!("Attester: {}", attester);
        tracing::info!("");

        let mut ipfs_hashes: std::collections::HashMap<String, String> = std::collections::HashMap::new();

        // Upload each file and collect IPFS hashes
        for file_upload in files_to_upload {
            tracing::info!("================================================");
            tracing::info!("{}", file_upload.header);
            tracing::info!("================================================");

            // Generate file name with timestamp
            let file_name = format!(
                "{}-{}",
                file_upload.default_name_prefix,
                chrono::Utc::now().format("%Y%m%d-%H%M%S")
            );

            tracing::info!(path = %file_upload.path.display(), "Uploading to Pinata IPFS");

            let ipfs_hash =
                Self::upload_file_to_pinata(file_upload.path.clone(), jwt.clone(), file_name, &file_upload.mime_type)
                    .await
                    .with_context(|| format!("failed to upload {}", file_upload.name))?;

            Self::display_upload_results(&ipfs_hash, &gateway);
            ipfs_hashes.insert(file_upload.json_key.clone(), ipfs_hash);
            tracing::info!("");
        }

        // Build output JSON
        let get_cid = |key: &str| -> Result<String> {
            ipfs_hashes
                .get(key)
                .cloned()
                .ok_or_else(|| eyre::eyre!("missing {}", key))
        };

        let policy_cids = PolicyCids {
            wasmCid: get_cid("wasmCid")?,
            policyCid: get_cid("policyCid")?,
            schemaCid: get_cid("schemaCid")?,
            attester,
            entrypoint: self.entrypoint.clone(),
            policyMetadataCid: get_cid("policyMetadataCid")?,
            policyDataMetadataCid: get_cid("policyDataMetadataCid")?,
            secretsSchemaCid: ipfs_hashes.get("secretsSchemaCid").cloned().unwrap_or_default(),
        };

        // Write to JSON file
        let json_content =
            serde_json::to_string_pretty(&policy_cids).context("failed to serialize policy_cids to JSON")?;

        // Create parent directory if it doesn't exist
        if let Some(parent) = self.output.parent() {
            std::fs::create_dir_all(parent).with_context(|| format!("failed to create directory: {:?}", parent))?;
        }

        std::fs::write(&self.output, &json_content)
            .with_context(|| format!("failed to write to file: {:?}", self.output))?;

        tracing::info!("================================================");
        tracing::info!("========== Created policy_cids.json ============");
        tracing::info!("================================================");
        tracing::info!("Output file: {:?}", self.output);
        tracing::info!("");
        tracing::info!("{}", json_content);

        Ok(())
    }
}

impl PolicyFilesCommand {
    /// Execute the policy-files command
    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        match self.subcommand {
            PolicyFilesSubcommand::GenerateCids(command) => {
                Box::new(command).execute(config).await?;
            }
        }
        Ok(())
    }
}