ether_link/
wallet.rs

1use k256::{
2    ecdsa::SigningKey,
3    elliptic_curve::rand_core::OsRng,
4    PublicKey,
5    SecretKey,
6};
7use sha3::{Digest, Keccak256};
8use hkdf::Hkdf;
9use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
10use aes_gcm::aead::Aead;
11use rand::RngCore;
12
13pub struct Wallet {
14    pub address: String,
15    pub privkey: String,
16    pub pubkey_compressed: String,
17    pub pubkey_uncompressed: String,
18}
19
20impl Wallet {
21    // Generate a new Ethereum key pair
22    pub fn new() -> Self {
23        Self::from_private_key(None).expect("Failed to generate wallet with random key")
24    }
25
26    // Create a wallet from an existing private key or generate a new one if None
27    pub fn from_private_key(private_key_hex: Option<&str>) -> Result<Self, String> {
28        let private_key = match private_key_hex {
29            Some(hex_key) => {
30                // Convert hex string to bytes
31                let key_bytes = hex::decode(hex_key.trim_start_matches("0x"))
32                    .map_err(|e| format!("Invalid private key format: {}", e))?;
33                
34                // Create SigningKey from bytes
35                let secret_key = SecretKey::from_slice(&key_bytes)
36                    .map_err(|e| format!("Invalid private key: {}", e))?;
37                SigningKey::from(secret_key)
38            },
39            None => {
40                // Generate random key
41                SigningKey::random(&mut OsRng)
42            }
43        };
44
45        let privkey_hex = format!("0x{}", hex::encode(private_key.to_bytes()));
46
47        let public_key = private_key.verifying_key();
48        // Compressed public key (starts with 0x02 or 0x03)
49        let pubkey_compressed = public_key.to_encoded_point(true).as_bytes().to_vec();
50        let pubkey_compressed_hex = format!("0x{}", hex::encode(&pubkey_compressed));
51        // Uncompressed public key (starts with 0x04)
52        let pubkey_uncompressed = public_key.to_encoded_point(false).as_bytes().to_vec();
53        let pubkey_uncompressed_hex = format!("0x{}", hex::encode(&pubkey_uncompressed));
54        
55        // Ethereum address (last 20 bytes of keccak256 of uncompressed key without prefix)
56        let hash = Keccak256::digest(&pubkey_uncompressed[1..]); // remove 0x04 prefix
57        let address = format!("0x{}", hex::encode(&hash[12..]));
58
59        Ok(Self {
60            address,
61            privkey: privkey_hex,
62            pubkey_compressed: pubkey_compressed_hex,
63            pubkey_uncompressed: pubkey_uncompressed_hex,
64        })
65    }
66    // Sign a message with the private key
67    pub async fn sign_message(&self, message: &str) -> Result<String, String> {
68        use ethers::core::types::Signature;
69        use ethers::signers::{LocalWallet, Signer};
70        use ethers::core::k256::SecretKey;
71        use std::time::{SystemTime, UNIX_EPOCH};
72        
73        // Convert hex private key to SecretKey
74        let secret_bytes = hex::decode(self.privkey.trim_start_matches("0x")).map_err(|e| e.to_string())?;
75        let secret_key = SecretKey::from_slice(&secret_bytes).map_err(|e| e.to_string())?;
76        let wallet = LocalWallet::from(secret_key);
77        
78        let address = &self.address;
79        let timestamps = SystemTime::now()
80            .duration_since(UNIX_EPOCH)
81            .map_err(|e| format!("Failed to get system time: {}", e))?
82            .as_secs();
83        let message_with_timestamp = format!("{}:{}", message, timestamps);
84        
85        // Verify the wallet address matches our derived address
86        let wallet_addr = format!("0x{}", hex::encode(wallet.address().as_bytes()));
87        debug_assert_eq!(wallet_addr.to_lowercase(), address.to_lowercase());
88
89        // Sign the hash with the private key
90        let signature: Signature = wallet
91            .sign_message(message_with_timestamp.clone())
92            .await
93            .map_err(|e| e.to_string())?;
94
95        // Return a JSON string containing both the message and signature for verification
96        let result = serde_json::json!({
97            "address": address,
98            "signed_message": message_with_timestamp,
99            "signature": format!("0x{}", hex::encode(signature.to_vec())),
100        });
101        Ok(result.to_string())
102    }
103    
104    // Encrypt a message using EIP-5630 compliant ECIES encryption
105    pub fn encrypt_message(&self, message: &str, recipient_compressed_pubkey: &str) -> Result<String, String> {
106        // Step 1: Parse the recipient's public key
107        let recipient_pubkey_bytes = hex::decode(recipient_compressed_pubkey.trim_start_matches("0x"))
108            .map_err(|e| format!("Invalid recipient public key: {}", e))?;
109        
110        let recipient_public_key = PublicKey::from_sec1_bytes(&recipient_pubkey_bytes)
111            .map_err(|e| format!("Failed to parse recipient public key: {}", e))?;
112            
113        // Step 3: Generate ephemeral key pair for ECIES
114        let ephemeral_private_key = SigningKey::random(&mut OsRng);
115        let ephemeral_public_key = ephemeral_private_key.verifying_key();
116        
117        // Step 4: Perform ECDH to derive shared secret
118        let shared_secret = {
119            use ethers::core::k256::ecdh::diffie_hellman;
120            let secret = diffie_hellman(
121                ephemeral_private_key.as_nonzero_scalar(), 
122                recipient_public_key.as_affine()
123            );
124            secret.raw_secret_bytes().to_vec()
125        };
126        
127        // Step 5: Derive encryption key using HKDF-SHA256
128        let kdf = Hkdf::<Keccak256>::new(None, &shared_secret);
129        let mut derived_key = [0u8; 32]; // 256 bits for AES-256-GCM
130        kdf.expand(b"EIP-5630-ECIES-AES-256-GCM", &mut derived_key)
131            .map_err(|_| "HKDF expansion failed".to_string())?;
132        
133        // Step 6: Generate random nonce for AES-GCM
134        let mut nonce_bytes = [0u8; 12]; // 96 bits
135        rand::rng().fill_bytes(&mut nonce_bytes);
136        
137        // Step 7: Encrypt the message using AES-256-GCM
138        let cipher = Aes256Gcm::new_from_slice(&derived_key)
139            .map_err(|_| "Failed to create cipher".to_string())?;
140        let nonce = Nonce::from_slice(&nonce_bytes);
141        
142        let ciphertext = cipher.encrypt(nonce, message.as_bytes().as_ref())
143            .map_err(|e| format!("Encryption failed: {}", e))?;
144        
145        // Step 8: Prepare the ephemeral public key for transmission
146        let ephemeral_pubkey_bytes = ephemeral_public_key.to_encoded_point(true).as_bytes().to_vec();
147        
148        // Step 9: Format according to EIP-5630 (ephemeral_pubkey || nonce || ciphertext)
149        let encrypted_data = [
150            ephemeral_pubkey_bytes.as_slice(),
151            &nonce_bytes,
152            ciphertext.as_slice(),
153        ].concat();
154        
155        // Return as hex string with 0x prefix
156        Ok(format!("0x{}", hex::encode(encrypted_data)))
157    }
158    pub fn decrypt_message(&self, encrypted_message: &str) -> Result<String, String> {
159        // Step 1: Decode the encrypted message
160        let encrypted_data = hex::decode(encrypted_message.trim_start_matches("0x"))
161            .map_err(|e| format!("Invalid encrypted message: {}", e))?;
162
163        // Ensure message has minimum required length (33 bytes pubkey + 12 bytes nonce + at least 16 bytes ciphertext)
164        if encrypted_data.len() < 33 + 12 + 16 {
165            return Err("Encrypted message is too short".to_string());
166        }
167
168        // Step 2: Extract components from encrypted data (ephemeral_pubkey || nonce || ciphertext)
169        let ephemeral_pubkey_bytes = &encrypted_data[0..33]; // Compressed public key is 33 bytes
170        let nonce_bytes = &encrypted_data[33..45]; // 12 bytes nonce
171        let ciphertext = &encrypted_data[45..]; // Rest is ciphertext
172
173        // Step 3: Parse the ephemeral public key
174        let ephemeral_public_key = PublicKey::from_sec1_bytes(ephemeral_pubkey_bytes)
175            .map_err(|e| format!("Failed to parse ephemeral public key: {}", e))?;
176
177        // Step 4: Parse our private key for ECDH
178        let private_key_bytes = hex::decode(self.privkey.trim_start_matches("0x"))
179            .map_err(|e| format!("Failed to decode private key: {}", e))?;
180        
181        // Use SecretKey for conversion
182        let secret_key = SecretKey::from_slice(&private_key_bytes)
183            .map_err(|e| format!("Failed to create secret key: {}", e))?;
184        let private_key = SigningKey::from(secret_key);
185
186        // Step 5: Perform ECDH to derive shared secret (same as in encrypt but with roles reversed)
187        let shared_secret = {
188            use ethers::core::k256::ecdh::diffie_hellman;
189            let secret = diffie_hellman(
190                private_key.as_nonzero_scalar(),
191                ephemeral_public_key.as_affine()
192            );
193            secret.raw_secret_bytes().to_vec()
194        };
195
196        // Step 6: Derive encryption key using HKDF-SHA256 (same as in encrypt)
197        let kdf = Hkdf::<Keccak256>::new(None, &shared_secret);
198        let mut derived_key = [0u8; 32]; // 256 bits for AES-256-GCM
199        kdf.expand(b"EIP-5630-ECIES-AES-256-GCM", &mut derived_key)
200            .map_err(|_| "HKDF expansion failed".to_string())?;
201        // Step 7: Decrypt the message using AES-256-GCM
202        let cipher = Aes256Gcm::new_from_slice(&derived_key)
203            .map_err(|_| "Failed to create cipher".to_string())?;
204        let nonce = Nonce::from_slice(nonce_bytes);
205
206        let plaintext = cipher.decrypt(nonce, ciphertext.as_ref())
207            .map_err(|_| "Decryption failed".to_string())?;
208
209        // Step 8: Convert plaintext bytes to string
210        let message = String::from_utf8(plaintext)
211            .map_err(|e| format!("Failed to decode message: {}", e))?;
212
213        Ok(message)
214    }
215    
216}
217
218