zombienet_orchestrator/generators/
key.rs

1use sp_core::{crypto::SecretStringError, ecdsa, ed25519, keccak_256, sr25519, Pair, H160, H256};
2
3use super::errors::GeneratorError;
4use crate::shared::types::{Accounts, NodeAccount};
5const KEYS: [&str; 5] = ["sr", "sr_stash", "ed", "ec", "eth"];
6
7pub fn generate_pair<T: Pair>(seed: &str) -> Result<T::Pair, SecretStringError> {
8    let pair = T::Pair::from_string(seed, None)?;
9    Ok(pair)
10}
11
12pub fn generate(seed: &str) -> Result<Accounts, GeneratorError> {
13    let mut accounts: Accounts = Default::default();
14    for k in KEYS {
15        let (address, public_key) = match k {
16            "sr" => {
17                let pair = generate_pair::<sr25519::Pair>(seed)
18                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
19                (pair.public().to_string(), hex::encode(pair.public()))
20            },
21            "sr_stash" => {
22                let pair = generate_pair::<sr25519::Pair>(&format!("{seed}//stash"))
23                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
24                (pair.public().to_string(), hex::encode(pair.public()))
25            },
26            "ed" => {
27                let pair = generate_pair::<ed25519::Pair>(seed)
28                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
29                (pair.public().to_string(), hex::encode(pair.public()))
30            },
31            "ec" => {
32                let pair = generate_pair::<ecdsa::Pair>(seed)
33                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
34                (pair.public().to_string(), hex::encode(pair.public()))
35            },
36            "eth" => {
37                let pair = generate_pair::<ecdsa::Pair>(seed)
38                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?;
39
40                let decompressed = libsecp256k1::PublicKey::parse_compressed(&pair.public().0)
41                    .map_err(|_| GeneratorError::KeyGeneration(k.into(), seed.into()))?
42                    .serialize();
43                let mut m = [0u8; 64];
44                m.copy_from_slice(&decompressed[1..65]);
45                let account = H160::from(H256::from(keccak_256(&m)));
46
47                (hex::encode(account), hex::encode(account))
48            },
49            _ => unreachable!(),
50        };
51        accounts.insert(k.into(), NodeAccount::new(address, public_key));
52    }
53    Ok(accounts)
54}
55
56#[cfg(test)]
57mod tests {
58
59    use super::*;
60    #[test]
61    fn generate_for_alice() {
62        use sp_core::crypto::Ss58Codec;
63        let s = "Alice";
64        let seed = format!("//{s}");
65
66        let pair = generate_pair::<sr25519::Pair>(&seed).unwrap();
67        assert_eq!(
68            "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY",
69            pair.public().to_ss58check()
70        );
71
72        let pair = generate_pair::<ecdsa::Pair>(&seed).unwrap();
73        assert_eq!(
74            "0x020a1091341fe5664bfa1782d5e04779689068c916b04cb365ec3153755684d9a1",
75            format!("0x{}", hex::encode(pair.public()))
76        );
77
78        let pair = generate_pair::<ed25519::Pair>(&seed).unwrap();
79        assert_eq!(
80            "5FA9nQDVg267DEd8m1ZypXLBnvN7SFxYwV7ndqSYGiN9TTpu",
81            pair.public().to_ss58check()
82        );
83    }
84
85    #[test]
86    fn generate_for_zombie() {
87        use sp_core::crypto::Ss58Codec;
88        let s = "Zombie";
89        let seed = format!("//{s}");
90
91        let pair = generate_pair::<sr25519::Pair>(&seed).unwrap();
92        assert_eq!(
93            "5FTcLfwFc7ctvqp3RhbEig6UuHLHcHVRujuUm8r21wy4dAR8",
94            pair.public().to_ss58check()
95        );
96    }
97
98    #[test]
99    fn generate_pair_invalid_should_fail() {
100        let s = "Alice";
101        let seed = s.to_string();
102
103        let pair = generate_pair::<sr25519::Pair>(&seed);
104        assert!(pair.is_err());
105    }
106
107    #[test]
108    fn generate_invalid_should_fail() {
109        let s = "Alice";
110        let seed = s.to_string();
111
112        let pair = generate(&seed);
113        assert!(pair.is_err());
114        assert!(matches!(pair, Err(GeneratorError::KeyGeneration(_, _))));
115    }
116
117    #[test]
118    fn generate_work() {
119        let s = "Alice";
120        let seed = format!("//{s}");
121
122        let pair = generate(&seed).unwrap();
123        let sr = pair.get("sr").unwrap();
124        let sr_stash = pair.get("sr_stash").unwrap();
125        let ed = pair.get("ed").unwrap();
126        let ec = pair.get("ec").unwrap();
127        let eth = pair.get("eth").unwrap();
128
129        assert_eq!(
130            sr.address,
131            "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY"
132        );
133        assert_eq!(
134            sr_stash.address,
135            "5GNJqTPyNqANBkUVMN1LPPrxXnFouWXoe2wNSmmEoLctxiZY"
136        );
137        assert_eq!(
138            ed.address,
139            "5FA9nQDVg267DEd8m1ZypXLBnvN7SFxYwV7ndqSYGiN9TTpu"
140        );
141        assert_eq!(
142            format!("0x{}", ec.public_key),
143            "0x020a1091341fe5664bfa1782d5e04779689068c916b04cb365ec3153755684d9a1"
144        );
145
146        assert_eq!(
147            format!("0x{}", eth.public_key),
148            "0xe04cc55ebee1cbce552f250e85c57b70b2e2625b"
149        )
150    }
151}