1  /* This file is part of DarkFi (https://dark.fi)
  2   *
  3   * Copyright (C) 2020-2025 Dyne.org foundation
  4   *
  5   * This program is free software: you can redistribute it and/or modify
  6   * it under the terms of the GNU Affero General Public License as
  7   * published by the Free Software Foundation, either version 3 of the
  8   * License, or (at your option) any later version.
  9   *
 10   * This program is distributed in the hope that it will be useful,
 11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
 12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 13   * GNU Affero General Public License for more details.
 14   *
 15   * You should have received a copy of the GNU Affero General Public License
 16   * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 17   */
 18  
 19  use std::{
 20      process::{exit, ExitCode},
 21      sync::{mpsc::channel, Arc},
 22      thread::available_parallelism,
 23  };
 24  
 25  use arg::Args;
 26  use darkfi::{util::cli::ProgressInc, ANSI_LOGO};
 27  use darkfi_money_contract::{model::TokenId, MoneyFunction};
 28  use darkfi_sdk::crypto::{
 29      contract_id::MONEY_CONTRACT_ID, poseidon_hash, BaseBlind, ContractId, FuncRef, PublicKey,
 30      SecretKey,
 31  };
 32  use rand::rngs::OsRng;
 33  use rayon::iter::ParallelIterator;
 34  
 35  const ABOUT: &str =
 36      concat!("vanityaddr ", env!("CARGO_PKG_VERSION"), '\n', env!("CARGO_PKG_DESCRIPTION"));
 37  
 38  const USAGE: &str = r#"
 39  Usage: vanityaddr [OPTIONS] <PREFIX> <PREFIX> ...
 40  
 41  Arguments:
 42    <PREFIX>    Prefixes to search
 43  
 44  Options:
 45    -c    Make the search case-sensitive
 46    -t    Number of threads to use (defaults to number of available CPUs)
 47    -A    Search for an address
 48    -C    Search for a Contract ID
 49    -T    Search for a Token ID
 50  "#;
 51  
 52  fn usage() {
 53      print!("{ANSI_LOGO}{ABOUT}\n{USAGE}");
 54  }
 55  
 56  struct DrkAddr {
 57      pub public: PublicKey,
 58      pub secret: SecretKey,
 59  }
 60  
 61  struct DrkToken {
 62      pub token_id: TokenId,
 63      pub secret: SecretKey,
 64      pub blind: BaseBlind,
 65  }
 66  
 67  struct DrkContract {
 68      pub contract_id: ContractId,
 69      pub secret: SecretKey,
 70  }
 71  
 72  trait Prefixable {
 73      fn new() -> Self;
 74      fn to_string(&self) -> String;
 75      fn _get_secret(&self) -> SecretKey;
 76  
 77      fn starts_with(&self, prefix: &str, case_sensitive: bool) -> bool {
 78          if case_sensitive {
 79              self.to_string().starts_with(prefix)
 80          } else {
 81              self.to_string().to_lowercase().starts_with(prefix.to_lowercase().as_str())
 82          }
 83      }
 84  
 85      fn starts_with_any(&self, prefixes: &[String], case_sensitive: bool) -> bool {
 86          prefixes.iter().any(|prefix| self.starts_with(prefix, case_sensitive))
 87      }
 88  }
 89  
 90  impl Prefixable for DrkAddr {
 91      fn new() -> Self {
 92          let secret = SecretKey::random(&mut OsRng);
 93          let public = PublicKey::from_secret(secret);
 94          Self { public, secret }
 95      }
 96  
 97      fn to_string(&self) -> String {
 98          self.public.to_string()
 99      }
100  
101      fn _get_secret(&self) -> SecretKey {
102          self.secret
103      }
104  }
105  
106  impl Prefixable for DrkToken {
107      fn new() -> Self {
108          // Generate the mint authority secret key and blind
109          let secret = SecretKey::random(&mut OsRng);
110          let blind = BaseBlind::random(&mut OsRng);
111  
112          // Create the Auth FuncID
113          let func_id = FuncRef {
114              contract_id: *MONEY_CONTRACT_ID,
115              func_code: MoneyFunction::AuthTokenMintV1 as u8,
116          }
117          .to_func_id();
118  
119          // Grab the mint authority user data
120          let (auth_x, auth_y) = PublicKey::from_secret(secret).xy();
121          let user_data = poseidon_hash([auth_x, auth_y]);
122  
123          // Derive the Token ID
124          let token_id = TokenId::derive_from(func_id.inner(), user_data, blind.inner());
125  
126          Self { token_id, secret, blind }
127      }
128  
129      fn to_string(&self) -> String {
130          self.token_id.to_string()
131      }
132  
133      fn _get_secret(&self) -> SecretKey {
134          self.secret
135      }
136  }
137  
138  impl Prefixable for DrkContract {
139      fn new() -> Self {
140          let secret = SecretKey::random(&mut OsRng);
141          let contract_id = ContractId::derive(secret);
142          Self { contract_id, secret }
143      }
144  
145      fn to_string(&self) -> String {
146          self.contract_id.to_string()
147      }
148  
149      fn _get_secret(&self) -> SecretKey {
150          self.secret
151      }
152  }
153  
154  fn main() -> ExitCode {
155      let argv;
156      let mut hflag = false;
157      let mut cflag = false;
158      let mut addrflag = false;
159      let mut toknflag = false;
160      let mut ctrcflag = false;
161  
162      let mut n_threads = available_parallelism().unwrap().get();
163  
164      {
165          let mut args = Args::new().with_cb(|args, flag| match flag {
166              'c' => cflag = true,
167              'A' => addrflag = true,
168              'T' => toknflag = true,
169              'C' => ctrcflag = true,
170              't' => n_threads = args.eargf().parse::<usize>().unwrap(),
171              _ => hflag = true,
172          });
173  
174          argv = args.parse();
175      }
176  
177      if hflag || argv.is_empty() {
178          usage();
179          return ExitCode::FAILURE
180      }
181  
182      if (addrflag as u8 + toknflag as u8 + ctrcflag as u8) != 1 {
183          eprintln!("The search flags are mutually exclusive. Use only one of -A/-C/-T.");
184          return ExitCode::FAILURE
185      }
186  
187      // Validate search prefixes
188      for (idx, prefix) in argv.iter().enumerate() {
189          match bs58::decode(prefix).into_vec() {
190              Ok(_) => {}
191              Err(e) => {
192                  eprintln!("Error: Invalid base58 for prefix #{idx}: {e}");
193                  return ExitCode::FAILURE
194              }
195          }
196      }
197  
198      // Handle SIGINT
199      let (tx, rx) = channel();
200      ctrlc::set_handler(move || tx.send(()).expect("Could not send signal on channel"))
201          .expect("Error setting SIGINT handler");
202  
203      // Something fancy
204      let progress = Arc::new(ProgressInc::new());
205  
206      // Threadpool
207      let progress_ = progress.clone();
208      let rayon_pool = rayon::ThreadPoolBuilder::new().num_threads(n_threads).build().unwrap();
209      rayon_pool.spawn(move || {
210          if addrflag {
211              let addr = rayon::iter::repeat(DrkAddr::new)
212                  .inspect(|_| progress_.inc(1))
213                  .map(|create| create())
214                  .find_any(|address| address.starts_with_any(&argv, cflag))
215                  .expect("Failed to find an address match");
216  
217              // The above will keep running until it finds a match or until
218              // the program terminates. Only if a match is found shall the
219              // following code be executed and the program exit successfully:
220              let attempts = progress_.position();
221              progress_.finish_and_clear();
222  
223              println!(
224                  "{{\"address\":\"{}\",\"attempts\":{attempts},\"secret\":\"{}\"}}",
225                  addr.public, addr.secret,
226              );
227          }
228  
229          if toknflag {
230              let tid = rayon::iter::repeat(DrkToken::new)
231                  .inspect(|_| progress_.inc(1))
232                  .map(|create| create())
233                  .find_any(|token_id| token_id.starts_with_any(&argv, cflag))
234                  .expect("Failed to find a token ID match");
235  
236              let attempts = progress_.position();
237              progress_.finish_and_clear();
238  
239              println!(
240                  "{{\"token_id\":\"{}\",\"attempts\":{attempts},\"secret\":\"{}\",\"blind\":\"{}\"}}",
241                  tid.token_id, tid.secret, tid.blind
242              );
243          }
244  
245          if ctrcflag {
246              let cid = rayon::iter::repeat(DrkContract::new)
247                  .inspect(|_| progress_.inc(1))
248                  .map(|create| create())
249                  .find_any(|contract_id| contract_id.starts_with_any(&argv, cflag))
250                  .expect("Failed to find a contract ID match");
251  
252              let attempts = progress_.position();
253              progress_.finish_and_clear();
254  
255              println!(
256                  "{{\"contract_id\":\"{attempts}\",\"attempts\":{},\"secret\":\"{}\"}}",
257                  cid.contract_id, cid.secret,
258              );
259          }
260  
261          exit(0);
262      });
263  
264      // This now blocks and lets our threadpool execute in the background.
265      rx.recv().expect("Could not receive from channel");
266      progress.finish_and_clear();
267      eprintln!("\r\x1b[2KCaught SIGINT, exiting...");
268      ExitCode::FAILURE
269  }