yellowstone_shield_cli/
lib.rs

1mod command;
2
3use anyhow::{Context, Result};
4use bs58::decode;
5use clap_derive::{Parser as DeriveParser, Subcommand};
6use log::info;
7use serde_json::from_str as parse_json_str;
8use solana_cli_config::Config;
9use solana_client::nonblocking::rpc_client::RpcClient;
10use solana_commitment_config::{CommitmentConfig, ParseCommitmentLevelError};
11use solana_keypair::Keypair;
12use solana_pubkey::Pubkey;
13use spl_token_metadata_interface::state::TokenMetadata;
14use std::fmt;
15use std::fs::read_to_string as read_path;
16use std::path::PathBuf;
17use std::sync::Arc;
18use std::{str::FromStr, time::Duration};
19use yellowstone_shield_client::types::PermissionStrategy;
20
21pub use command::*;
22
23use crate::command::policy::PolicyVersion;
24
25#[derive(Debug, DeriveParser)]
26#[command(
27    author,
28    version,
29    name = "Yellowstone Shield CLI",
30    about = "CLI for managing Yellowstone shield policies"
31)]
32pub struct Args {
33    /// RPC endpoint url to override using the Solana config
34    #[arg(short, long, global = true)]
35    pub rpc: Option<String>,
36
37    /// Log level
38    #[arg(short, long, global = true, default_value = "off")]
39    pub log_level: String,
40
41    /// Path to the owner keypair file
42    #[arg(short, long, global = true)]
43    pub keypair: Option<String>,
44
45    #[command(subcommand)]
46    pub command: Command,
47}
48
49#[derive(Subcommand, Debug)]
50pub enum Command {
51    /// Manage policies
52    Policy {
53        #[command(subcommand)]
54        action: PolicyAction,
55    },
56    /// Manage identities
57    Identities {
58        #[command(subcommand)]
59        action: IdentitiesAction,
60    },
61}
62
63#[derive(Subcommand, Debug)]
64pub enum PolicyAction {
65    /// Create a new policy
66    Create {
67        /// The strategy to use for the policy
68        #[arg(long)]
69        strategy: PermissionStrategy,
70
71        /// The name of the policy
72        #[arg(long)]
73        name: String,
74
75        /// The symbol of the policy
76        #[arg(long)]
77        symbol: String,
78
79        /// The URI of the policy
80        #[arg(long)]
81        uri: String,
82    },
83    /// Delete a policy
84    Delete {
85        /// The mint address associated with the policy
86        #[arg(long)]
87        mint: Pubkey,
88    },
89    /// Show policy details
90    Show {
91        /// The mint address associated with the policy
92        #[arg(long)]
93        mint: Pubkey,
94    },
95}
96
97#[derive(Subcommand, Debug)]
98pub enum IdentitiesAction {
99    /// Add identities to a policy
100    Add {
101        /// The mint address associated with the policy
102        #[arg(long)]
103        mint: Pubkey,
104        /// The identities to add to the policy
105        #[arg(long)]
106        identities_path: PathBuf,
107    },
108    /// Update/Replace Identities for a Policy
109    Update {
110        /// The mint address associated with the policy
111        #[arg(long)]
112        mint: Pubkey,
113        /// The identities to update/replace
114        #[arg(long)]
115        identities_path: PathBuf,
116    },
117
118    /// Remove identities from a policy
119    Remove {
120        /// The mint address associated with the policy
121        #[arg(long)]
122        mint: Pubkey,
123        /// The identities to remove from the policy
124        #[arg(long)]
125        identities_path: PathBuf,
126    },
127}
128
129#[derive(thiserror::Error, Debug)]
130pub enum CliError {
131    #[error("unable to get config file path")]
132    ConfigFilePathError,
133    #[error(transparent)]
134    Io(#[from] std::io::Error),
135    #[error(transparent)]
136    Other(#[from] anyhow::Error),
137    #[error(transparent)]
138    ParseCommitmentLevelError(#[from] ParseCommitmentLevelError),
139    #[error("unable to parse keypair")]
140    Keypair,
141}
142
143pub async fn run(config: Arc<Config>, command: Command) -> RunResult {
144    let client = RpcClient::new_with_timeout_and_commitment(
145        config.json_rpc_url.clone(),
146        Duration::from_secs(90),
147        CommitmentConfig::from_str(&config.commitment).map_err::<CliError, _>(Into::into)?,
148    );
149    let keypair = parse_keypair(&config.keypair_path)?;
150    let context = command::CommandContext { keypair, client };
151
152    match &command {
153        Command::Policy { action } => match action {
154            PolicyAction::Create {
155                strategy,
156                name,
157                symbol,
158                uri,
159            } => {
160                policy::CreateCommandBuilder::new()
161                    .strategy(*strategy)
162                    .name(name.clone())
163                    .symbol(symbol.clone())
164                    .uri(uri.clone())
165                    .run(context)
166                    .await
167            }
168            PolicyAction::Delete { mint } => {
169                policy::DeleteCommandBuilder::new()
170                    .mint(mint)
171                    .run(context)
172                    .await
173            }
174            PolicyAction::Show { mint } => {
175                policy::ShowCommandBuilder::new()
176                    .mint(mint)
177                    .run(context)
178                    .await
179            }
180        },
181        Command::Identities { action } => match action {
182            IdentitiesAction::Add {
183                mint,
184                identities_path,
185            } => {
186                let identities: Vec<Pubkey> = read_path(identities_path)?
187                    .lines()
188                    .filter_map(|s| Pubkey::from_str(s.trim()).ok())
189                    .collect();
190
191                identity::AddBatchCommandBuilder::new()
192                    .mint(mint)
193                    .identities(identities)
194                    .run(context)
195                    .await
196            }
197            IdentitiesAction::Update {
198                mint,
199                identities_path,
200            } => {
201                let identities: Vec<Pubkey> = read_path(identities_path)?
202                    .lines()
203                    .filter_map(|s| Pubkey::from_str(s.trim()).ok())
204                    .collect();
205
206                identity::UpdateBatchCommandBuilder::new()
207                    .mint(mint)
208                    .identities(identities)
209                    .run(context)
210                    .await
211            }
212            IdentitiesAction::Remove {
213                mint,
214                identities_path,
215            } => {
216                let identities: Vec<Pubkey> = read_path(identities_path)?
217                    .lines()
218                    .filter_map(|s| Pubkey::from_str(s.trim()).ok())
219                    .collect();
220
221                identity::RemoveBatchCommandBuilder::new()
222                    .mint(mint)
223                    .identities(identities)
224                    .run(context)
225                    .await
226            }
227        },
228    }
229}
230
231fn parse_keypair(keypair_path: &str) -> Result<Keypair, CliError> {
232    let secret_string = read_path(keypair_path).context("Can't find key file")?;
233    let secret_bytes = parse_json_str(&secret_string)
234        .or_else(|_| decode(&secret_string.trim()).into_vec())
235        .map_err(|_| CliError::ConfigFilePathError)?;
236
237    Keypair::try_from(secret_bytes.as_slice()).map_err(|_| CliError::Keypair)
238}
239
240pub struct LogPolicy<'a> {
241    token_mint: &'a Pubkey,
242    token_metadata: &'a TokenMetadata,
243    policy_address: &'a Pubkey,
244    policy_info: &'a PolicyVersion,
245    identities: Option<&'a Vec<Pubkey>>,
246}
247
248impl<'a> LogPolicy<'a> {
249    pub fn new(
250        token_mint: &'a Pubkey,
251        token_metadata: &'a TokenMetadata,
252        policy_address: &'a Pubkey,
253        policy_info: &'a PolicyVersion,
254        identities: Option<&'a Vec<Pubkey>>,
255    ) -> Self {
256        LogPolicy {
257            token_mint,
258            token_metadata,
259            policy_address,
260            policy_info,
261            identities,
262        }
263    }
264
265    fn log(&self) {
266        info!("{}", self);
267    }
268}
269
270impl fmt::Display for LogPolicy<'_> {
271    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
272        writeln!(f)?;
273        writeln!(f)?;
274        writeln!(f, "📜 Policy")?;
275        writeln!(f, "--------------------------------")?;
276        writeln!(f, "🏠 Addresses")?;
277        writeln!(f, "  📜 Policy: {}", self.policy_address)?;
278        writeln!(f, "  🪙 Mint: {}", self.token_mint)?;
279        writeln!(f, "--------------------------------")?;
280        writeln!(f, "🔍 Details")?;
281        let strategy = match self.policy_info.strategy() {
282            0 => "❌ Strategy: Deny",
283            1 => "✅ Strategy: Allow",
284            _ => "❓ Strategy: Unknown",
285        };
286        writeln!(f, "  {}", strategy)?;
287        writeln!(f, "  🏷️  Name: {}", self.token_metadata.name)?;
288        writeln!(f, "  🔖 Symbol: {}", self.token_metadata.symbol)?;
289        writeln!(f, "  🌐 URI: {}", self.token_metadata.uri)?;
290        writeln!(f, "--------------------------------")?;
291        if let Some(identities) = self.identities {
292            writeln!(f, "  🔑 Identities in policy:")?;
293            if !identities.is_empty() {
294                for (i, identity) in identities.iter().enumerate() {
295                    writeln!(f, "    {}. {}", i, identity)?;
296                }
297            } else {
298                writeln!(f, "    []")?;
299            }
300            writeln!(f, "--------------------------------")?;
301        }
302        Ok(())
303    }
304}