newton_cli/commands/

policy.rs

use alloy::primitives::Address;
use clap::{Parser, Subcommand};
use eyre::{Context, Result};
use newton_prover_chainio::policy::PolicyController;
use newton_prover_core::{
    common::parse_intent,
    config::NewtonAvsConfig,
    rego::{evaluate, KycIdentityData, Value as RegoValue},
};
use serde_json::Value;
use std::path::PathBuf;
use tracing::{self, info};

use crate::{commands::utils, config::NewtonCliConfig};

/// Policy commands
#[derive(Debug, Parser)]
#[command(name = "policy")]
pub struct PolicyCommand {
    #[command(subcommand)]
    pub subcommand: PolicySubcommand,
}

#[derive(Debug, Subcommand)]
pub enum PolicySubcommand {
    /// Deploy policy
    Deploy(DeployCommand),
    /// Simulate Rego policy execution
    Simulate(SimulateCommand),
}

fn validate_non_empty_path(s: &str) -> Result<PathBuf, String> {
    if s.trim().is_empty() {
        Err(String::from("Path cannot be empty"))
    } else {
        Ok(PathBuf::from(s))
    }
}

/// Normalize intent JSON by converting value and chainId to number strings
/// Supports: number strings, numbers, and hex strings (with or without 0x prefix)
fn normalize_intent(mut intent: serde_json::Value) -> eyre::Result<serde_json::Value> {
    // Normalize value field
    if let Some(value_field) = intent.get_mut("value") {
        let normalized = normalize_number_field(value_field).with_context(|| "Failed to normalize value field")?;
        *value_field = serde_json::Value::String(normalized);
    }

    // Normalize chainId field (optional)
    if let Some(chain_id_field) = intent.get_mut("chainId") {
        let normalized = normalize_number_field(chain_id_field).with_context(|| "Failed to normalize chainId field")?;
        *chain_id_field = serde_json::Value::String(normalized);
    }

    Ok(intent)
}

/// Normalize a number field that can be a string, number, or hex string
/// Returns a decimal number string
fn normalize_number_field(value: &serde_json::Value) -> eyre::Result<String> {
    match value {
        serde_json::Value::String(s) => {
            // Check if it's a hex string
            if s.starts_with("0x") || s.starts_with("0X") {
                // Parse hex string - try as u64 first, then U256 for very large values
                let hex_str = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")).unwrap();

                // Try parsing as u64 first (for chainId)
                if let Ok(num) = u64::from_str_radix(hex_str, 16) {
                    Ok(num.to_string())
                } else {
                    // Try parsing as U256 (for large value fields)
                    let num = alloy::primitives::U256::from_str_radix(hex_str, 16)
                        .map_err(|e| eyre::eyre!("Invalid hex string '{}': {}", s, e))?;
                    Ok(num.to_string())
                }
            } else {
                // Already a decimal string, validate it's a valid number
                // Try parsing as U256 first (supports very large numbers), then u64
                if s.parse::<alloy::primitives::U256>().is_ok() || s.parse::<u64>().is_ok() {
                    Ok(s.clone())
                } else {
                    Err(eyre::eyre!("Invalid number string '{}'", s))
                }
            }
        }
        serde_json::Value::Number(n) => {
            // Convert JSON number to string
            if let Some(u) = n.as_u64() {
                Ok(u.to_string())
            } else if let Some(i) = n.as_i64() {
                if i < 0 {
                    return Err(eyre::eyre!("Negative numbers are not supported: {}", i));
                }
                Ok(i.to_string())
            } else {
                // For very large numbers (f64), convert to string
                // Note: JSON numbers are limited to f64 precision, so this may lose precision
                Ok(n.to_string())
            }
        }
        _ => Err(eyre::eyre!("Expected string or number, got: {}", value)),
    }
}

/// Simulate Rego policy execution command
#[derive(Debug, Parser)]
pub struct SimulateCommand {
    /// Path to the WASM component file
    #[arg(long)]
    wasm_file: PathBuf,

    /// Path to the Rego policy file
    #[arg(long)]
    rego_file: PathBuf,

    /// Path to the intent JSON file
    #[arg(long)]
    intent_json: PathBuf,

    /// Entrypoint rule to evaluate (e.g., "max_gas_price.allow" or "data.max_gas_price.allow")
    /// The "data." prefix will be automatically added if not present
    #[arg(long, default_value = "allow")]
    entrypoint: String,

    /// Path to JSON file containing arguments for WASM execution
    /// If not provided, an empty JSON object "{}" will be used
    #[arg(long)]
    wasm_args: Option<PathBuf>,

    /// Path to JSON file containing policy parameters data
    /// If not provided, an empty JSON object "{}" will be used for the params field
    #[arg(long)]
    policy_params_data: Option<PathBuf>,

    /// Path to JSON file containing KYC identity data for simulation.
    /// Expected fields: status, selected_country_code, address_subdivision,
    /// address_country_code, birthdate, expiration_date, issue_date, issuing_authority.
    /// If not provided, a default approved US/CA identity is used.
    #[arg(long)]
    identity_data: Option<PathBuf>,
}

/// Deploy policy command
#[derive(Debug, Parser)]
pub struct DeployCommand {
    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,

    /// Address of the deployed policy data contract
    #[arg(long)]
    policy_data_address: Address,

    #[arg(long, value_parser = validate_non_empty_path)]
    policy_cids: PathBuf,
}

impl DeployCommand {
    /// Deploy policy contract
    ///
    /// This function deploys a policy contract using the Rust chainio functions.
    /// It returns the deployed policy address.
    async fn deploy_policy(
        private_key: &str,
        rpc_url: &str,
        policy_cid: &str,
        schema_cid: &str,
        entrypoint: &str,
        policy_data_address: Address,
        policy_metadata_cid: &str,
    ) -> Result<Address> {
        let controller = PolicyController::new(private_key.to_string(), rpc_url.to_string());

        // Create policy_data array with the deployed policy data address
        let policy_data = vec![policy_data_address];

        tracing::info!(
            "Deploying policy:\n  entrypoint: {}\n  policyCid: {}\n  schemaCid: {}\n  policyData: {:?}\n  metadataCid: {} \n",
            entrypoint,
            policy_cid,
            schema_cid,
            policy_data,
            policy_metadata_cid
        );

        let (_receipt, policy_address) = controller
            .deploy_policy(
                policy_cid.to_string(),
                schema_cid.to_string(),
                entrypoint.to_string(),
                policy_data,
                policy_metadata_cid.to_string(),
            )
            .await
            .map_err(|e| eyre::eyre!("Failed to deploy policy: {}", e))?;

        tracing::info!("Policy deployed successfully at address: {}", policy_address);

        Ok(policy_address)
    }

    /// Execute the deploy command
    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        // Get values from args or env, with error if still missing
        let private_key = self
            .private_key
            .ok_or_else(|| eyre::eyre!("private_key is required (use --private-key or PRIVATE_KEY env var)"))?;

        let rpc_url = self
            .rpc_url
            .ok_or_else(|| eyre::eyre!("rpc_url is required (use --rpc-url or RPC_URL env var)"))?;

        // Read from policy_cids.json file
        let json_content = std::fs::read_to_string(&self.policy_cids)
            .with_context(|| format!("Failed to read policy_cids.json: {:?}", self.policy_cids))?;
        let json: Value = serde_json::from_str(&json_content)
            .with_context(|| format!("Failed to parse policy_cids.json: {:?}", self.policy_cids))?;

        // Extract values from policy_cids.json
        let policy_cid = json.get("policyCid").and_then(|v| v.as_str()).unwrap_or("");
        let schema_cid = json.get("schemaCid").and_then(|v| v.as_str()).unwrap_or("");
        let entrypoint = json.get("entrypoint").and_then(|v| v.as_str()).unwrap_or("");
        let policy_metadata_cid = json.get("policyMetadataCid").and_then(|v| v.as_str()).unwrap_or("");

        // Deploy policy using the provided policy data address
        Self::deploy_policy(
            &private_key,
            &rpc_url,
            policy_cid,
            schema_cid,
            entrypoint,
            self.policy_data_address,
            policy_metadata_cid,
        )
        .await?;
        Ok(())
    }
}

impl SimulateCommand {
    /// Execute the simulate command
    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        info!("Starting Rego simulation...");

        // Read intent JSON file
        info!("Reading intent JSON from: {:?}", self.intent_json);
        let intent_json_str = std::fs::read_to_string(&self.intent_json)
            .with_context(|| format!("Failed to read intent JSON file: {:?}", self.intent_json))?;
        let mut intent_value: serde_json::Value = serde_json::from_str(&intent_json_str)
            .with_context(|| format!("Failed to parse intent JSON: {:?}", self.intent_json))?;

        // Normalize intent (convert value and chainId to number strings)
        info!("Normalizing intent fields...");
        intent_value = normalize_intent(intent_value).with_context(|| "Failed to normalize intent fields")?;

        // Read WASM args if provided, otherwise use empty JSON object
        let wasm_input = if let Some(wasm_args_path) = &self.wasm_args {
            info!("Reading WASM args from: {:?}", wasm_args_path);
            let wasm_args_str = std::fs::read_to_string(wasm_args_path)
                .with_context(|| format!("Failed to read WASM args file: {:?}", wasm_args_path))?;
            // Validate it's valid JSON
            serde_json::from_str::<serde_json::Value>(&wasm_args_str)
                .with_context(|| format!("Failed to parse WASM args as JSON: {:?}", wasm_args_path))?;
            wasm_args_str
        } else {
            "{}".to_string()
        };

        // Execute WASM to get policy data
        info!("Executing WASM file: {:?}", self.wasm_file);
        let wasm_output = utils::execute_wasm(self.wasm_file, wasm_input, config)
            .await
            .with_context(|| "Failed to execute WASM file")?;

        // Parse WASM output as JSON
        let policy_data: serde_json::Value = serde_json::from_str(&wasm_output)
            .with_context(|| format!("Failed to parse WASM output as JSON: {}", wasm_output))?;

        // Parse intent
        info!("Parsing intent...");
        let parsed_intent = parse_intent(intent_value).with_context(|| "Failed to parse intent")?;

        // Convert parsed intent to string
        let parsed_intent_str: String = parsed_intent.into();

        // Read Rego policy file
        info!("Reading Rego policy from: {:?}", self.rego_file);
        let policy = std::fs::read_to_string(&self.rego_file)
            .with_context(|| format!("Failed to read Rego policy file: {:?}", self.rego_file))?;

        // Read policy params data if provided, otherwise use empty JSON object
        let policy_params: serde_json::Value = if let Some(policy_params_path) = &self.policy_params_data {
            info!("Reading policy params data from: {:?}", policy_params_path);
            let policy_params_str = std::fs::read_to_string(policy_params_path)
                .with_context(|| format!("Failed to read policy params data file: {:?}", policy_params_path))?;
            serde_json::from_str(&policy_params_str)
                .with_context(|| format!("Failed to parse policy params data as JSON: {:?}", policy_params_path))?
        } else {
            serde_json::json!({})
        };

        // Construct policy params and data
        let policy_params_and_data = serde_json::json!({
            "params": policy_params,
            "data": policy_data,
        });
        let policy_params_and_data_str = policy_params_and_data.to_string();

        // Print the data object that will be used in the rego evaluation
        info!("\n=== Data object for Rego evaluation ===");
        info!("{}", serde_json::to_string_pretty(&policy_params_and_data)?);
        info!("==========================================\n");

        // Normalize entrypoint: add "data." prefix if not present
        let entrypoint = if self.entrypoint.starts_with("data.") {
            self.entrypoint.clone()
        } else {
            format!("data.{}", self.entrypoint)
        };

        let identity_data = if let Some(identity_path) = &self.identity_data {
            info!("Reading identity data from: {:?}", identity_path);
            let identity_str = std::fs::read_to_string(identity_path)
                .with_context(|| format!("Failed to read identity data file: {:?}", identity_path))?;
            let v: serde_json::Value = serde_json::from_str(&identity_str)
                .with_context(|| format!("Failed to parse identity data JSON: {:?}", identity_path))?;
            KycIdentityData {
                reference_date: v
                    .get("reference_date")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
                status: v.get("status").and_then(|x| x.as_str()).unwrap_or("").to_string(),
                selected_country_code: v
                    .get("selected_country_code")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
                address_subdivision: v
                    .get("address_subdivision")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
                address_country_code: v
                    .get("address_country_code")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
                birthdate: v.get("birthdate").and_then(|x| x.as_str()).unwrap_or("").to_string(),
                expiration_date: v
                    .get("expiration_date")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
                issue_date: v.get("issue_date").and_then(|x| x.as_str()).unwrap_or("").to_string(),
                issuing_authority: v
                    .get("issuing_authority")
                    .and_then(|x| x.as_str())
                    .unwrap_or("")
                    .to_string(),
            }
        } else {
            KycIdentityData {
                status: "approved".to_string(),
                address_subdivision: "CA".to_string(),
                address_country_code: "US".to_string(),
                ..Default::default()
            }
        };

        // Evaluate the policy
        info!("Evaluating policy with entrypoint: {}", entrypoint);
        let result = evaluate(
            policy,
            &policy_params_and_data_str,
            &parsed_intent_str,
            Some(Box::new(identity_data)),
            &entrypoint,
            None,
        )
        .with_context(|| "Failed to evaluate Rego policy")?;

        // Print the result
        match result {
            RegoValue::Bool(b) => {
                info!("Evaluation result: {}", b);
                if b {
                    info!("  Policy evaluation: ALLOWED");
                } else {
                    info!("  Policy evaluation: DENIED");
                }
            }
            _ => {
                info!("Evaluation result: {:?}", result);
                if result == RegoValue::Undefined {
                    info!("  Policy evaluation: UNDEFINED (denied)");
                } else {
                    info!("? Policy evaluation: {:?}", result);
                }
            }
        }

        Ok(())
    }
}

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