eigen_cli/
generate.rs

1use crate::args::KeyType;
2use crate::EigenKeyCliError;
3use ark_ff::UniformRand;
4use ark_serialize::{CanonicalSerialize, SerializationError};
5use eth_keystore::encrypt_key;
6use rand::{distributions::Alphanumeric, Rng};
7use rand_core::OsRng;
8use std::io::Write;
9use std::{
10    fs::{self, File},
11    path::Path,
12};
13use uuid::Uuid;
14
15const PASSWORD_LENGTH: usize = 20;
16pub const DEFAULT_KEY_FOLDER: &str = "keys";
17pub const PASSWORD_FILE: &str = "password.txt";
18pub const PRIVATE_KEY_HEX_FILE: &str = "private_key_hex.txt";
19
20pub enum KeyGenerator {
21    ECDSAKeyGenerator,
22    BLSKeyGenerator,
23}
24
25impl KeyGenerator {
26    /// Generates a number of private keys in the given directory.
27    ///
28    /// # Arguments
29    ///
30    /// * `num_keys` - The number of keys to generate.
31    /// * `output_dir` - The directory where the key files are generated.
32    pub fn generate(
33        self,
34        num_keys: u32,
35        output_dir: Option<String>,
36    ) -> Result<(), EigenKeyCliError> {
37        let dir_name = output_dir.unwrap_or_else(|| {
38            let id = Uuid::new_v4();
39            format!("{}-{}", self.key_name(), id)
40        });
41
42        let dir_path = Path::new(&dir_name);
43        let key_path = dir_path.join(DEFAULT_KEY_FOLDER);
44        fs::create_dir_all(key_path).map_err(EigenKeyCliError::FileError)?;
45
46        self.generate_keys(num_keys, dir_path)
47    }
48
49    /// Generates a number of private keys and stores them both encrypted and in plaintext.
50    /// It creates the following files:
51    /// - `passwords.txt`: contains all passwords to decrypt keys
52    /// - `private_key_hex.txt`: plaintext private keys
53    /// - `keys/*`: all the encrypted json files in this folder
54    ///
55    /// # Arguments
56    ///
57    /// * `num_keys` - The number of keys to generate.
58    /// * `path` - The path to the directory where the generated files are stored.
59    fn generate_keys(self, num_keys: u32, path: &Path) -> Result<(), EigenKeyCliError> {
60        let key_path = path.join(DEFAULT_KEY_FOLDER);
61        let private_key_path = path.join(PRIVATE_KEY_HEX_FILE);
62        let password_path = path.join(PASSWORD_FILE);
63
64        for i in 0..num_keys {
65            let password = KeyGenerator::generate_random_password();
66            let private_key = self
67                .random_key()
68                .map_err(EigenKeyCliError::SerializationError)?;
69            let private_key_hex = hex::encode(private_key.clone());
70
71            // encrypt the private key into `path` directory
72            let name = format!("{}.{}.key.json", i + 1, self.key_name());
73            encrypt_key(
74                key_path.clone(),
75                &mut OsRng,
76                private_key,
77                password.clone(),
78                Some(&name),
79            )
80            .map_err(EigenKeyCliError::KeystoreError)?;
81
82            // write the private key into `private_key_file`
83            File::create(private_key_path.clone())
84                .and_then(|mut file| file.write_all(private_key_hex.as_bytes()))
85                .map_err(EigenKeyCliError::FileError)?;
86
87            // write the password into `password_file`
88            File::create(password_path.clone())
89                .and_then(|mut file| file.write_all(password.as_bytes()))
90                .map_err(EigenKeyCliError::FileError)?;
91
92            if (i + 1) % 50 == 0 {
93                println!("Generated {} keys\n", i + 1);
94            }
95        }
96        Ok(())
97    }
98
99    /// Generates a random key which can be type ecdsa or BLS.
100    ///
101    /// # Returns
102    ///
103    /// * A private key as a vector of bytes.
104    fn random_key(&self) -> Result<Vec<u8>, SerializationError> {
105        match self {
106            KeyGenerator::ECDSAKeyGenerator => Ok(Self::random_ecdsa_key()),
107            KeyGenerator::BLSKeyGenerator => Ok(Self::random_bls_key()?),
108        }
109    }
110
111    /// Generates a random ecdsa key.
112    ///
113    /// # Returns
114    ///
115    /// * An ecdsa private key as a vector of bytes.
116    pub fn random_ecdsa_key() -> Vec<u8> {
117        let private_key = k256::SecretKey::random(&mut OsRng);
118        private_key.to_bytes().as_slice().to_vec()
119    }
120
121    /// Generates a random BLS key.
122    ///
123    /// # Returns
124    ///
125    /// * A BLS private key as a vector of bytes.
126    fn random_bls_key() -> Result<Vec<u8>, SerializationError> {
127        let mut buffer = Vec::new();
128        let private_key = eigen_crypto_bls::PrivateKey::rand(&mut OsRng);
129        private_key.serialize_uncompressed(&mut buffer)?;
130        Ok(buffer)
131    }
132
133    /// Generates a 20-character random password.
134    ///
135    /// # Returns
136    ///
137    /// * A random password.
138    pub fn generate_random_password() -> String {
139        rand::thread_rng()
140            .sample_iter(&Alphanumeric)
141            .take(PASSWORD_LENGTH)
142            .map(char::from)
143            .collect()
144    }
145
146    /// Get the key type.
147    ///
148    /// # Returns
149    ///
150    /// * The key type as a string.
151    fn key_name(&self) -> String {
152        match self {
153            KeyGenerator::ECDSAKeyGenerator => "ecdsa",
154            KeyGenerator::BLSKeyGenerator => "bls",
155        }
156        .to_string()
157    }
158}
159
160impl From<KeyType> for KeyGenerator {
161    fn from(value: KeyType) -> Self {
162        match value {
163            KeyType::Ecdsa => KeyGenerator::ECDSAKeyGenerator,
164            KeyType::Bls => KeyGenerator::BLSKeyGenerator,
165        }
166    }
167}