newton_cli/commands/

policy_client.rs

use alloy::{
    network::EthereumWallet,
    primitives::{Address, Bytes},
    providers::ProviderBuilder,
    signers::local::PrivateKeySigner,
    transports::http::reqwest::Url,
};
use clap::{Parser, Subcommand};
use eyre::Context;
use newton_prover_chainio::{
    policy_client::PolicyClientController, policy_client_registry::PolicyClientRegistryController,
};
use newton_prover_core::{config::NewtonAvsConfig, newton_policy_client::NewtonPolicyClient};
use std::{path::PathBuf, str::FromStr};
use tracing::info;

use crate::config::NewtonCliConfig;

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

#[derive(Debug, Subcommand)]
pub enum PolicyClientSubcommand {
    /// Register a policy client with the PolicyClientRegistry
    #[command(name = "register")]
    Register(RegisterCommand),

    /// Deactivate a registered policy client
    #[command(name = "deactivate")]
    Deactivate(DeactivateCommand),

    /// Reactivate a previously deactivated policy client
    #[command(name = "activate")]
    Activate(ActivateCommand),

    /// Transfer registry ownership record of a policy client
    #[command(name = "transfer-ownership")]
    TransferOwnership(TransferOwnershipCommand),

    /// Query registration status and record for a policy client
    #[command(name = "status")]
    Status(StatusCommand),

    /// List all policy clients owned by an address
    #[command(name = "list")]
    List(ListCommand),

    /// Set the policy address on a policy client contract (owner-only)
    #[command(name = "set-policy")]
    SetPolicy(SetPolicyCommand),

    /// Set policy parameters for a policy client
    #[command(name = "set-policy-params")]
    SetPolicyParams(SetPolicyParamsCommand),
}

/// Common arguments for registry write operations
#[derive(Debug, Parser)]
struct RegistryWriteArgs {
    /// PolicyClientRegistry contract address
    #[arg(long)]
    registry: Address,

    /// Policy client contract address
    #[arg(long)]
    client: Address,

    /// Private key of the transaction signer
    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    /// RPC endpoint URL
    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,
}

impl RegistryWriteArgs {
    fn validate(&self) -> eyre::Result<(String, String)> {
        let private_key = self
            .private_key
            .clone()
            .ok_or_else(|| eyre::eyre!("PRIVATE_KEY is required (set via --private-key or env var)"))?;
        let rpc_url = self
            .rpc_url
            .clone()
            .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;
        Ok((private_key, rpc_url))
    }
}

/// Register a policy client with the PolicyClientRegistry
#[derive(Debug, Parser)]
pub struct RegisterCommand {
    #[command(flatten)]
    args: RegistryWriteArgs,
}

/// Deactivate a registered policy client
#[derive(Debug, Parser)]
pub struct DeactivateCommand {
    #[command(flatten)]
    args: RegistryWriteArgs,
}

/// Reactivate a deactivated policy client
#[derive(Debug, Parser)]
pub struct ActivateCommand {
    #[command(flatten)]
    args: RegistryWriteArgs,
}

/// Transfer registry ownership of a policy client
#[derive(Debug, Parser)]
pub struct TransferOwnershipCommand {
    /// PolicyClientRegistry contract address
    #[arg(long)]
    registry: Address,

    /// Policy client contract address
    #[arg(long)]
    client: Address,

    /// Address of the new owner
    #[arg(long)]
    new_owner: Address,

    /// Private key of the current owner
    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    /// RPC endpoint URL
    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,
}

/// Query status of a policy client in the registry
#[derive(Debug, Parser)]
pub struct StatusCommand {
    /// PolicyClientRegistry contract address
    #[arg(long)]
    registry: Address,

    /// Policy client contract address
    #[arg(long)]
    client: Address,

    /// RPC endpoint URL
    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,
}

/// List all policy clients for an owner
#[derive(Debug, Parser)]
pub struct ListCommand {
    /// PolicyClientRegistry contract address
    #[arg(long)]
    registry: Address,

    /// Owner address to query
    #[arg(long)]
    owner: Address,

    /// RPC endpoint URL
    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,
}

/// Set the policy address on a policy client contract
#[derive(Debug, Parser)]
pub struct SetPolicyCommand {
    /// Policy client contract address
    #[arg(long)]
    client: Address,

    /// New policy contract address
    #[arg(long)]
    policy: Address,

    /// Private key of the policy client owner
    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    /// RPC endpoint URL
    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,
}

/// Set policy parameters command
#[derive(Debug, Parser)]
pub struct SetPolicyParamsCommand {
    /// Policy client address
    #[arg(long)]
    policy_client: Address,

    /// Path to JSON file containing policy parameters
    #[arg(long)]
    policy_params: PathBuf,

    /// Expire after (in blocks)
    #[arg(long)]
    expire_after: u32,

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

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

impl PolicyClientCommand {
    /// Execute the policy-client command
    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        match self.subcommand {
            PolicyClientSubcommand::Register(cmd) => {
                let (private_key, rpc_url) = cmd.args.validate()?;
                let controller = PolicyClientRegistryController::new(private_key, rpc_url, cmd.args.registry);
                let receipt = controller
                    .register_client(cmd.args.client)
                    .await
                    .map_err(|e| eyre::eyre!("failed to register client: {}", e))?;
                info!("Transaction hash: {}", receipt.transaction_hash);
                Ok(())
            }

            PolicyClientSubcommand::Deactivate(cmd) => {
                let (private_key, rpc_url) = cmd.args.validate()?;
                let controller = PolicyClientRegistryController::new(private_key, rpc_url, cmd.args.registry);
                let receipt = controller
                    .deactivate_client(cmd.args.client)
                    .await
                    .map_err(|e| eyre::eyre!("failed to deactivate client: {}", e))?;
                info!("Transaction hash: {}", receipt.transaction_hash);
                Ok(())
            }

            PolicyClientSubcommand::Activate(cmd) => {
                let (private_key, rpc_url) = cmd.args.validate()?;
                let controller = PolicyClientRegistryController::new(private_key, rpc_url, cmd.args.registry);
                let receipt = controller
                    .activate_client(cmd.args.client)
                    .await
                    .map_err(|e| eyre::eyre!("failed to activate client: {}", e))?;
                info!("Transaction hash: {}", receipt.transaction_hash);
                Ok(())
            }

            PolicyClientSubcommand::TransferOwnership(cmd) => {
                let private_key = cmd
                    .private_key
                    .ok_or_else(|| eyre::eyre!("PRIVATE_KEY is required (set via --private-key or env var)"))?;
                let rpc_url = cmd
                    .rpc_url
                    .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;

                let controller = PolicyClientRegistryController::new(private_key, rpc_url, cmd.registry);
                let receipt = controller
                    .set_client_owner(cmd.client, cmd.new_owner)
                    .await
                    .map_err(|e| eyre::eyre!("failed to transfer ownership: {}", e))?;
                info!("Transaction hash: {}", receipt.transaction_hash);
                Ok(())
            }

            PolicyClientSubcommand::Status(cmd) => {
                let rpc_url = cmd
                    .rpc_url
                    .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;

                let controller = PolicyClientRegistryController::new_read_only(rpc_url, cmd.registry);

                let is_active = controller
                    .is_registered_client(cmd.client)
                    .await
                    .map_err(|e| eyre::eyre!("failed to query registration status: {}", e))?;

                // Try to fetch the full record (reverts if not registered at all)
                match controller.get_client_record(cmd.client).await {
                    Ok(record) => {
                        println!("Client: {}", cmd.client);
                        println!("  Owner: {}", record.owner);
                        println!("  Active: {}", record.active);
                        println!("  Registered At: {} (unix timestamp)", record.registeredAt);
                        println!("  isRegisteredClient: {}", is_active);
                    }
                    Err(_) => {
                        println!("Client: {}", cmd.client);
                        println!("  Not registered in registry {}", cmd.registry);
                    }
                }
                Ok(())
            }

            PolicyClientSubcommand::List(cmd) => {
                let rpc_url = cmd
                    .rpc_url
                    .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;

                let controller = PolicyClientRegistryController::new_read_only(rpc_url, cmd.registry);

                let clients = controller
                    .get_clients_by_owner(cmd.owner)
                    .await
                    .map_err(|e| eyre::eyre!("failed to query clients: {}", e))?;

                println!("Owner: {}", cmd.owner);
                println!("Registered clients: {}", clients.len());
                for client in &clients {
                    println!("  {}", client);
                }
                Ok(())
            }

            PolicyClientSubcommand::SetPolicy(cmd) => {
                let private_key = cmd
                    .private_key
                    .ok_or_else(|| eyre::eyre!("PRIVATE_KEY is required (set via --private-key or env var)"))?;
                let rpc_url = cmd
                    .rpc_url
                    .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;

                let private_key_str = private_key.strip_prefix("0x").unwrap_or(&private_key);
                let signer = PrivateKeySigner::from_str(private_key_str).context("failed to parse private key")?;
                let wallet = EthereumWallet::from(signer);
                let url = Url::parse(&rpc_url).context("invalid RPC URL")?;
                let provider = ProviderBuilder::new().wallet(wallet).connect_http(url);

                let policy_client = NewtonPolicyClient::new(cmd.client, provider);

                info!(
                    "Calling setPolicyAddress({}) on policy client {}",
                    cmd.policy, cmd.client
                );

                let tx = policy_client
                    .setPolicyAddress(cmd.policy)
                    .send()
                    .await
                    .context("failed to send setPolicyAddress transaction")?;

                info!("Transaction sent: {:?}", tx.tx_hash());

                let receipt = tx
                    .get_receipt()
                    .await
                    .context("failed to get setPolicyAddress receipt")?;

                if !receipt.status() {
                    return Err(eyre::eyre!(
                        "setPolicyAddress transaction reverted (tx: {:?})",
                        receipt.transaction_hash
                    ));
                }

                info!("Policy client {} now points to policy {}", cmd.client, cmd.policy);
                info!("Transaction hash: {}", receipt.transaction_hash);
                Ok(())
            }

            PolicyClientSubcommand::SetPolicyParams(cmd) => {
                let private_key = cmd
                    .private_key
                    .ok_or_else(|| eyre::eyre!("PRIVATE_KEY is required (set via --private-key or env var)"))?;
                let rpc_url = cmd
                    .rpc_url
                    .ok_or_else(|| eyre::eyre!("RPC_URL is required (set via --rpc-url or env var)"))?;

                info!("Setting policy params for policy client: {}", cmd.policy_client);

                // Read policy params JSON file
                let policy_params_str = std::fs::read_to_string(&cmd.policy_params)
                    .with_context(|| format!("failed to read policy params file: {:?}", cmd.policy_params))?;

                // Validate it's valid JSON
                let _policy_params_json: serde_json::Value = serde_json::from_str(&policy_params_str)
                    .with_context(|| format!("failed to parse policy params as JSON: {:?}", cmd.policy_params))?;

                // Convert to bytes
                let policy_params_bytes = Bytes::copy_from_slice(policy_params_str.as_bytes());

                // Create controller and set policy
                let controller = PolicyClientController::new(private_key, rpc_url, cmd.policy_client);
                let (receipt, policy_id) = controller
                    .set_policy(policy_params_bytes, cmd.expire_after)
                    .await
                    .map_err(|e| eyre::eyre!("failed to set policy: {}", e))?;

                info!("Policy set successfully");
                info!("Transaction hash: {:?}", receipt.transaction_hash);
                info!("Policy ID: {:?}", policy_id);

                Ok(())
            }
        }
    }
}