Skip to main content

newton_cli/commands/
policy.rs

1use alloy::primitives::Address;
2use clap::{Parser, Subcommand};
3use eyre::{Context, Result};
4use newton_prover_chainio::policy::PolicyController;
5use newton_prover_core::{
6    common::parse_intent,
7    config::NewtonAvsConfig,
8    rego::{evaluate, Value as RegoValue},
9};
10use serde_json::Value;
11use std::path::PathBuf;
12use tracing::{self, info};
13
14use crate::{commands::utils, config::NewtonCliConfig};
15
16/// Policy commands
17#[derive(Debug, Parser)]
18#[command(name = "policy")]
19pub struct PolicyCommand {
20    #[command(subcommand)]
21    pub subcommand: PolicySubcommand,
22}
23
24#[derive(Debug, Subcommand)]
25pub enum PolicySubcommand {
26    /// Deploy policy
27    Deploy(DeployCommand),
28    /// Simulate Rego policy execution
29    Simulate(SimulateCommand),
30}
31
32fn validate_non_empty_path(s: &str) -> Result<PathBuf, String> {
33    if s.trim().is_empty() {
34        Err(String::from("Path cannot be empty"))
35    } else {
36        Ok(PathBuf::from(s))
37    }
38}
39
40/// Normalize intent JSON by converting value and chainId to number strings
41/// Supports: number strings, numbers, and hex strings (with or without 0x prefix)
42fn normalize_intent(mut intent: serde_json::Value) -> eyre::Result<serde_json::Value> {
43    // Normalize value field
44    if let Some(value_field) = intent.get_mut("value") {
45        let normalized = normalize_number_field(value_field).with_context(|| "Failed to normalize value field")?;
46        *value_field = serde_json::Value::String(normalized);
47    }
48
49    // Normalize chainId field (optional)
50    if let Some(chain_id_field) = intent.get_mut("chainId") {
51        let normalized = normalize_number_field(chain_id_field).with_context(|| "Failed to normalize chainId field")?;
52        *chain_id_field = serde_json::Value::String(normalized);
53    }
54
55    Ok(intent)
56}
57
58/// Normalize a number field that can be a string, number, or hex string
59/// Returns a decimal number string
60fn normalize_number_field(value: &serde_json::Value) -> eyre::Result<String> {
61    match value {
62        serde_json::Value::String(s) => {
63            // Check if it's a hex string
64            if s.starts_with("0x") || s.starts_with("0X") {
65                // Parse hex string - try as u64 first, then U256 for very large values
66                let hex_str = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")).unwrap();
67
68                // Try parsing as u64 first (for chainId)
69                if let Ok(num) = u64::from_str_radix(hex_str, 16) {
70                    Ok(num.to_string())
71                } else {
72                    // Try parsing as U256 (for large value fields)
73                    let num = alloy::primitives::U256::from_str_radix(hex_str, 16)
74                        .map_err(|e| eyre::eyre!("Invalid hex string '{}': {}", s, e))?;
75                    Ok(num.to_string())
76                }
77            } else {
78                // Already a decimal string, validate it's a valid number
79                // Try parsing as U256 first (supports very large numbers), then u64
80                if s.parse::<alloy::primitives::U256>().is_ok() || s.parse::<u64>().is_ok() {
81                    Ok(s.clone())
82                } else {
83                    Err(eyre::eyre!("Invalid number string '{}'", s))
84                }
85            }
86        }
87        serde_json::Value::Number(n) => {
88            // Convert JSON number to string
89            if let Some(u) = n.as_u64() {
90                Ok(u.to_string())
91            } else if let Some(i) = n.as_i64() {
92                if i < 0 {
93                    return Err(eyre::eyre!("Negative numbers are not supported: {}", i));
94                }
95                Ok(i.to_string())
96            } else {
97                // For very large numbers (f64), convert to string
98                // Note: JSON numbers are limited to f64 precision, so this may lose precision
99                Ok(n.to_string())
100            }
101        }
102        _ => Err(eyre::eyre!("Expected string or number, got: {}", value)),
103    }
104}
105
106/// Simulate Rego policy execution command
107#[derive(Debug, Parser)]
108pub struct SimulateCommand {
109    /// Path to the WASM component file
110    #[arg(long)]
111    wasm_file: PathBuf,
112
113    /// Path to the Rego policy file
114    #[arg(long)]
115    rego_file: PathBuf,
116
117    /// Path to the intent JSON file
118    #[arg(long)]
119    intent_json: PathBuf,
120
121    /// Entrypoint rule to evaluate (e.g., "max_gas_price.allow" or "data.max_gas_price.allow")
122    /// The "data." prefix will be automatically added if not present
123    #[arg(long, default_value = "allow")]
124    entrypoint: String,
125
126    /// Path to JSON file containing arguments for WASM execution
127    /// If not provided, an empty JSON object "{}" will be used
128    #[arg(long)]
129    wasm_args: Option<PathBuf>,
130
131    /// Path to JSON file containing policy parameters data
132    /// If not provided, an empty JSON object "{}" will be used for the params field
133    #[arg(long)]
134    policy_params_data: Option<PathBuf>,
135}
136
137/// Deploy policy command
138#[derive(Debug, Parser)]
139pub struct DeployCommand {
140    #[arg(long, env = "PRIVATE_KEY")]
141    private_key: Option<String>,
142
143    #[arg(long, env = "RPC_URL")]
144    rpc_url: Option<String>,
145
146    /// Address of the deployed policy data contract
147    #[arg(long)]
148    policy_data_address: Address,
149
150    #[arg(long, value_parser = validate_non_empty_path)]
151    policy_cids: PathBuf,
152}
153
154impl DeployCommand {
155    /// Deploy policy contract
156    ///
157    /// This function deploys a policy contract using the Rust chainio functions.
158    /// It returns the deployed policy address.
159    async fn deploy_policy(
160        private_key: &str,
161        rpc_url: &str,
162        policy_cid: &str,
163        schema_cid: &str,
164        entrypoint: &str,
165        policy_data_address: Address,
166        policy_metadata_cid: &str,
167    ) -> Result<Address> {
168        let controller = PolicyController::new(private_key.to_string(), rpc_url.to_string());
169
170        // Create policy_data array with the deployed policy data address
171        let policy_data = vec![policy_data_address];
172
173        tracing::info!(
174            "Deploying policy:\n  entrypoint: {}\n  policyCid: {}\n  schemaCid: {}\n  policyData: {:?}\n  metadataCid: {} \n",
175            entrypoint,
176            policy_cid,
177            schema_cid,
178            policy_data,
179            policy_metadata_cid
180        );
181
182        let (_receipt, policy_address) = controller
183            .deploy_policy(
184                policy_cid.to_string(),
185                schema_cid.to_string(),
186                entrypoint.to_string(),
187                policy_data,
188                policy_metadata_cid.to_string(),
189            )
190            .await
191            .map_err(|e| eyre::eyre!("Failed to deploy policy: {}", e))?;
192
193        tracing::info!("Policy deployed successfully at address: {}", policy_address);
194
195        Ok(policy_address)
196    }
197
198    /// Execute the deploy command
199    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
200        // Get values from args or env, with error if still missing
201        let private_key = self
202            .private_key
203            .ok_or_else(|| eyre::eyre!("private_key is required (use --private-key or PRIVATE_KEY env var)"))?;
204
205        let rpc_url = self
206            .rpc_url
207            .ok_or_else(|| eyre::eyre!("rpc_url is required (use --rpc-url or RPC_URL env var)"))?;
208
209        // Read from policy_cids.json file
210        let json_content = std::fs::read_to_string(&self.policy_cids)
211            .with_context(|| format!("Failed to read policy_cids.json: {:?}", self.policy_cids))?;
212        let json: Value = serde_json::from_str(&json_content)
213            .with_context(|| format!("Failed to parse policy_cids.json: {:?}", self.policy_cids))?;
214
215        // Extract values from policy_cids.json
216        let policy_cid = json.get("policyCid").and_then(|v| v.as_str()).unwrap_or("");
217        let schema_cid = json.get("schemaCid").and_then(|v| v.as_str()).unwrap_or("");
218        let entrypoint = json.get("entrypoint").and_then(|v| v.as_str()).unwrap_or("");
219        let policy_metadata_cid = json.get("policyMetadataCid").and_then(|v| v.as_str()).unwrap_or("");
220
221        // Deploy policy using the provided policy data address
222        Self::deploy_policy(
223            &private_key,
224            &rpc_url,
225            policy_cid,
226            schema_cid,
227            entrypoint,
228            self.policy_data_address,
229            policy_metadata_cid,
230        )
231        .await?;
232        Ok(())
233    }
234}
235
236impl SimulateCommand {
237    /// Execute the simulate command
238    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
239        info!("Starting Rego simulation...");
240
241        // Read intent JSON file
242        info!("Reading intent JSON from: {:?}", self.intent_json);
243        let intent_json_str = std::fs::read_to_string(&self.intent_json)
244            .with_context(|| format!("Failed to read intent JSON file: {:?}", self.intent_json))?;
245        let mut intent_value: serde_json::Value = serde_json::from_str(&intent_json_str)
246            .with_context(|| format!("Failed to parse intent JSON: {:?}", self.intent_json))?;
247
248        // Normalize intent (convert value and chainId to number strings)
249        info!("Normalizing intent fields...");
250        intent_value = normalize_intent(intent_value).with_context(|| "Failed to normalize intent fields")?;
251
252        // Read WASM args if provided, otherwise use empty JSON object
253        let wasm_input = if let Some(wasm_args_path) = &self.wasm_args {
254            info!("Reading WASM args from: {:?}", wasm_args_path);
255            let wasm_args_str = std::fs::read_to_string(wasm_args_path)
256                .with_context(|| format!("Failed to read WASM args file: {:?}", wasm_args_path))?;
257            // Validate it's valid JSON
258            serde_json::from_str::<serde_json::Value>(&wasm_args_str)
259                .with_context(|| format!("Failed to parse WASM args as JSON: {:?}", wasm_args_path))?;
260            wasm_args_str
261        } else {
262            "{}".to_string()
263        };
264
265        // Execute WASM to get policy data
266        info!("Executing WASM file: {:?}", self.wasm_file);
267        let wasm_output = utils::execute_wasm(self.wasm_file, wasm_input, config)
268            .await
269            .with_context(|| "Failed to execute WASM file")?;
270
271        // Parse WASM output as JSON
272        let policy_data: serde_json::Value = serde_json::from_str(&wasm_output)
273            .with_context(|| format!("Failed to parse WASM output as JSON: {}", wasm_output))?;
274
275        // Parse intent
276        info!("Parsing intent...");
277        let parsed_intent = parse_intent(intent_value).with_context(|| "Failed to parse intent")?;
278
279        // Convert parsed intent to string
280        let parsed_intent_str: String = parsed_intent.into();
281
282        // Read Rego policy file
283        info!("Reading Rego policy from: {:?}", self.rego_file);
284        let policy = std::fs::read_to_string(&self.rego_file)
285            .with_context(|| format!("Failed to read Rego policy file: {:?}", self.rego_file))?;
286
287        // Read policy params data if provided, otherwise use empty JSON object
288        let policy_params: serde_json::Value = if let Some(policy_params_path) = &self.policy_params_data {
289            info!("Reading policy params data from: {:?}", policy_params_path);
290            let policy_params_str = std::fs::read_to_string(policy_params_path)
291                .with_context(|| format!("Failed to read policy params data file: {:?}", policy_params_path))?;
292            serde_json::from_str(&policy_params_str)
293                .with_context(|| format!("Failed to parse policy params data as JSON: {:?}", policy_params_path))?
294        } else {
295            serde_json::json!({})
296        };
297
298        // Construct policy params and data
299        let policy_params_and_data = serde_json::json!({
300            "params": policy_params,
301            "data": policy_data,
302        });
303        let policy_params_and_data_str = policy_params_and_data.to_string();
304
305        // Print the data object that will be used in the rego evaluation
306        info!("\n=== Data object for Rego evaluation ===");
307        info!("{}", serde_json::to_string_pretty(&policy_params_and_data)?);
308        info!("==========================================\n");
309
310        // Normalize entrypoint: add "data." prefix if not present
311        let entrypoint = if self.entrypoint.starts_with("data.") {
312            self.entrypoint.clone()
313        } else {
314            format!("data.{}", self.entrypoint)
315        };
316
317        // Evaluate the policy
318        info!("Evaluating policy with entrypoint: {}", entrypoint);
319        let result = evaluate(policy, &policy_params_and_data_str, &parsed_intent_str, &entrypoint)
320            .with_context(|| "Failed to evaluate Rego policy")?;
321
322        // Print the result
323        match result {
324            RegoValue::Bool(b) => {
325                info!("Evaluation result: {}", b);
326                if b {
327                    info!("  Policy evaluation: ALLOWED");
328                } else {
329                    info!("  Policy evaluation: DENIED");
330                }
331            }
332            _ => {
333                info!("Evaluation result: {:?}", result);
334                if result == RegoValue::Undefined {
335                    info!("  Policy evaluation: UNDEFINED (denied)");
336                } else {
337                    info!("? Policy evaluation: {:?}", result);
338                }
339            }
340        }
341
342        Ok(())
343    }
344}
345
346impl PolicyCommand {
347    /// Execute the policy command
348    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
349        match self.subcommand {
350            PolicySubcommand::Deploy(command) => {
351                Box::new(command).execute(config).await?;
352            }
353            PolicySubcommand::Simulate(command) => {
354                Box::new(command).execute(config).await?;
355            }
356        }
357        Ok(())
358    }
359}