klave_connector/
key.rs

1use std::fmt::Display;
2use aes_gcm::{aead::Aead, Aes256Gcm, Key as AesGcmKey, KeyInit, Nonce};
3use base64::{engine::general_purpose, Engine};
4use elliptic_curve::sec1::ToEncodedPoint;
5use p256::SecretKey;
6use serde_json::Value;
7use crate::utils::{get_sha256_bytes, get_random_bytes};
8
9
10#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)]
11pub struct AesKeyGenParams {
12    name: String,
13    length: u32,
14}
15
16#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, PartialEq)]
17pub struct CryptoKey {
18    crv: String,
19    ext: bool,
20    key_ops: Vec<String>,
21    kty: String,
22    x: String,
23    y: String,
24    d: Option<String>
25}
26
27impl Display for CryptoKey {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        write!(f, "CryptoKey {{ crv: {}, ext: {}, key_ops: {:?}, kty: {}, x: {}, y: {}, d: {:?} }}", 
30            self.crv, self.ext, self.key_ops, self.kty, self.x, self.y, self.d)
31    }
32}
33
34#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, PartialEq)]
35pub struct CryptoKeyPair {
36    public_key: CryptoKey,
37    private_key: CryptoKey,
38}
39
40impl Display for CryptoKeyPair {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        write!(f, "CryptoKeyPair {{ public_key: {}, private_key: {} }}", 
43            self.public_key, self.private_key)  
44    }
45}
46
47#[derive(serde::Deserialize, serde::Serialize)]
48pub struct EncryptedKeyPairV0 {
49    version: String,
50    name: String,
51    iv: String,
52    salt: String,
53    encrypted_keys: String,
54    encrypted_key_pair: String,
55}
56
57#[derive(serde::Deserialize, serde::Serialize)]
58pub struct EncryptedKeyPairV2 {
59    version: u64,
60    name: String,
61    iv: String,
62    salt: String,
63    data: String,
64}
65
66#[derive(serde::Deserialize, serde::Serialize)]
67pub enum EncryptedKeyPair {
68    V0(EncryptedKeyPairV0),
69    V2(EncryptedKeyPairV2),
70}
71
72pub enum CryptoKeyType {
73    Public,
74    Private,
75}
76
77#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, PartialEq)]
78pub struct Key {
79    crypto_key_pair: CryptoKeyPair,
80}
81
82impl Display for Key {
83    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84        write!(f, "Key {{ crypto_key_pair: {} }}", self.crypto_key_pair)
85    }
86}
87
88impl Key {
89    pub fn new(secret_key_input: Option<SecretKey>) -> Key {        
90        let secret_key = match secret_key_input {
91            Some(secret_key) => secret_key,
92            None => {
93                // Generate a new secret key if none is provided
94                SecretKey::random(&mut rand::thread_rng())
95            }
96        };
97        Key {
98            crypto_key_pair: CryptoKeyPair {
99                public_key: Key::key_pair_to_crypto_key(&secret_key, CryptoKeyType::Public).unwrap(),
100                private_key: Key::key_pair_to_crypto_key(&secret_key, CryptoKeyType::Private).unwrap(),
101            }
102        }
103    }
104
105    pub fn get_raw_public_key(&self) -> Vec<u8> {
106        let public_key = &self.crypto_key_pair.public_key;
107        let x = general_purpose::URL_SAFE_NO_PAD.decode(public_key.x.clone()).expect("Failed to decode x");
108        let y = general_purpose::URL_SAFE_NO_PAD.decode(public_key.y.clone()).expect("Failed to decode y");
109        let mut raw_public_key = vec![0; 64];
110        raw_public_key[0..32].copy_from_slice(&x);
111        raw_public_key[32..64].copy_from_slice(&y);
112        raw_public_key
113    }
114
115    pub fn get_raw_private_key(&self) -> Vec<u8> {
116        let private_key = &self.crypto_key_pair.private_key;
117        match &private_key.d {
118            Some(d) => {
119                let d = general_purpose::URL_SAFE_NO_PAD.decode(d).expect("Failed to decode d");
120                let mut raw_private_key = vec![0; 32];
121                raw_private_key[0..32].copy_from_slice(&d);
122                return raw_private_key;
123            }
124            None => {
125                return vec![];
126            }
127        }
128    }
129
130    pub fn key_pair_to_crypto_key(secret_key: &SecretKey, crypto_key_type: CryptoKeyType) -> Result<CryptoKey, Box<dyn std::error::Error>> {
131        // Derive the public key from the secret key
132        let public_key = secret_key.public_key();
133    
134        // Get the encoded point (x, y coordinates)
135        let encoded_point = public_key.to_encoded_point(false); // Uncompressed point
136    
137        // Extract x and y as byte slices
138        let x = match encoded_point.x() {
139            Some(x) => x,
140            None => return Err("Failed to extract x coordinate".into()),
141        };
142        let y = match encoded_point.y() {
143            Some(y) => y,
144            None => return Err("Failed to extract y coordinate".into()),
145        };
146    
147        // Convert x and y to Base64 strings
148        let x_b64 = general_purpose::URL_SAFE_NO_PAD.encode(x);
149        let y_b64 = general_purpose::URL_SAFE_NO_PAD.encode(y);
150    
151        // Create the CryptoKey
152        Ok(CryptoKey {
153            crv: "P-256".to_string(),
154            ext: true,
155            key_ops: vec!["sign".to_string(), "verify".to_string()],
156            kty: "EC".to_string(),
157            x: x_b64,
158            y: y_b64,
159            d: match crypto_key_type {
160                CryptoKeyType::Public => None,
161                CryptoKeyType::Private => Some(general_purpose::URL_SAFE_NO_PAD.encode(secret_key.to_bytes())),
162            },
163        })
164    }
165
166    pub fn import_jwk_from_file(        
167        file_path: &str,
168        password: &str
169    ) -> Result<Key, Box<dyn std::error::Error>> {
170        // Read the file content
171        let file_content = match std::fs::read_to_string(file_path) {
172            Ok(content) => content,
173            Err(_) => return Err("Failed to read file".into()),
174        };
175        // Parse the JSON content
176        let Ok(v) = serde_json::from_str::<Value>(&file_content.clone()) else {
177            return Err("Failed to parse JSON".into());
178        };            
179
180        //Check if the file did contain iv, salt and data
181        let iv = v.get("iv").and_then(|v| v.as_str()).unwrap_or("");
182        let salt = v.get("salt").and_then(|v| v.as_str()).unwrap_or("");
183        match v.get("data") {
184            Some(data) => {
185                println!("File contains a v2 key pair");
186                let data = data.as_str().unwrap_or("");
187                let encrypted_key_pair = EncryptedKeyPair::V2(EncryptedKeyPairV2 {
188                    version: 2,
189                    name: v.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(),
190                    iv: iv.to_string(),
191                    salt: salt.to_string(),
192                    data: data.to_string(),
193                });
194                return Ok(match Key::import_key_pair(encrypted_key_pair, password) {
195                    Ok(key) => key,
196                    Err(_) => {
197                        return Err("Failed to import key".into());
198                    }
199                });
200            }
201            None => {
202                println!("File contains a v0 key pair");
203                let encrypted_key_pair = EncryptedKeyPair::V0(EncryptedKeyPairV0 {
204                    version: 0.to_string(),
205                    name: v.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(),
206                    iv: iv.to_string(),
207                    salt: salt.to_string(),
208                    encrypted_keys: v.get("encrypted_keys").and_then(|v| v.as_str()).unwrap_or("").to_string(),
209                    encrypted_key_pair: v.get("encrypted_key_pair").and_then(|v| v.as_str()).unwrap_or("").to_string(),
210                });
211                return Ok(match Key::import_key_pair(encrypted_key_pair, password) {
212                    Ok(key) => key,
213                    Err(_) => {
214                        return Err("Failed to import key".into());
215                    }
216                });                
217            }
218        };
219    }
220
221    pub fn import_jwk(        
222        encrypted_key_pair_v2: &str,
223        password: &str
224    ) -> Result<Key, Box<dyn std::error::Error>> {
225
226        // Parse the JSON content
227        let Ok(encrypted_key_pair) = serde_json::from_str::<EncryptedKeyPairV2>(&encrypted_key_pair_v2) else {
228            return Err("Failed to parse JSON".into());
229        };            
230
231        return Ok(match Key::import_key_pair(EncryptedKeyPair::V2(encrypted_key_pair), password) {
232            Ok(key) => key,
233            Err(_) => {
234                return Err("Failed to import key".into());
235            }
236        });
237    }
238
239    pub fn export_jwk(self, password: &str) -> Result<String, Box<dyn std::error::Error>> {
240        // Convert the key to JSON
241        let json = serde_json::to_string(&self.crypto_key_pair)?;
242        // Encrypt the JSON data using AES-GCM
243        let iv = get_random_bytes(12);
244        let salt = get_random_bytes(16);
245        let weak_pwd = password.to_string().into_bytes();
246        let mut combined = salt.clone();
247        combined.extend_from_slice(&weak_pwd);
248        let strong_pwd = get_sha256_bytes(&combined);
249        // Create an AES key using the strong password imported as raw
250        let aes_key = AesGcmKey::<Aes256Gcm>::from_slice(&strong_pwd);
251        // Encrypt the data using the AES key
252        let cipher = Aes256Gcm::new(aes_key);            
253        let nonce = Nonce::from_slice(&iv);
254        let encrypted_data = match cipher.encrypt(nonce, json.as_bytes()) {
255            Ok(encrypted_data) => encrypted_data,
256            Err(_) => return Err("Failed to encrypt data".into())
257        };
258        
259        // Encode the encrypted data and IV as Base64
260        let encrypted_key_pair_v2 = EncryptedKeyPairV2{
261            iv: general_purpose::URL_SAFE.encode(iv),
262            salt: general_purpose::URL_SAFE.encode(salt),
263            data: general_purpose::URL_SAFE.encode(encrypted_data),
264            name: self.crypto_key_pair.public_key.kty,
265            version: 2,
266        };
267
268        // Create a JSON object with the encrypted data and IV
269        let json_output = serde_json::to_string::<EncryptedKeyPairV2>(&encrypted_key_pair_v2)?;        
270        Ok(json_output)
271    }
272
273    pub fn import_key_pair(
274        input_encrypted_key_pair: EncryptedKeyPair,
275        password: &str
276    ) -> Result<Key, Box<dyn std::error::Error>> {
277
278        let decrypt_v2 = |iv: &str, salt: &str, data: &str| {
279            let iv = match general_purpose::URL_SAFE.decode(iv) {
280                Ok(iv) => iv,
281                Err(_) => return Err(Box::<dyn std::error::Error>::from("Failed to decode IV"))
282            };                
283            let salt = match general_purpose::URL_SAFE.decode(salt) {
284                Ok(salt) => salt,
285                Err(_) => return Err("Failed to decode salt".into())
286            };
287            let encrypted_data = match general_purpose::URL_SAFE.decode(data) {
288                Ok(encrypted_data) => encrypted_data,
289                Err(_) => return Err("Failed to decode encrypted data".into())
290            };
291            let weak_pwd = password.to_string().into_bytes();
292            let mut combined = salt.clone();
293            combined.extend_from_slice(&weak_pwd);
294            let strong_pwd = get_sha256_bytes(&combined);
295            // Create an AES key using the strong password imported as raw
296            let aes_key = AesGcmKey::<Aes256Gcm>::from_slice(&strong_pwd);
297            // Decrypt the data using the AES key
298            let cipher = Aes256Gcm::new(aes_key);            
299            let nonce = Nonce::from_slice(&iv);
300            let decrypted = match cipher.decrypt(nonce, encrypted_data.as_slice()) {
301                Ok(decrypted) => decrypted,
302                Err(_) => return Err("Failed to decrypt data".into())
303            };
304            //Convert this decrypted data to a UTF8 string
305            let decrypted_str = match String::from_utf8(decrypted){
306                Ok(decrypted_str) => decrypted_str,
307                Err(_) => return Err("Failed to convert decrypted data to UTF-8".into())
308            };
309            // Parse the decrypted string as JSON
310            let decrypted_json: Value = match serde_json::from_str(&decrypted_str) {
311                Ok(decrypted_json) => decrypted_json,
312                Err(_) => return Err("Failed to parse decrypted data as JSON".into())
313            };
314            // Extract the public and private keys from the decrypted JSON
315            let public_key = match serde_json::from_value::<CryptoKey>(match decrypted_json.get("public_key") {
316                Some(public_key) => public_key.clone(),
317                None => return Err("Failed to get public key".into())
318            }) {
319                Ok(crypto_key) => crypto_key,
320                Err(_) => return Err("Failed to parse crypto key".into())
321            };
322            let private_key = match serde_json::from_value::<CryptoKey>(match decrypted_json.get("private_key") {
323                Some(private_key) => private_key.clone(),
324                None => return Err("Failed to get private key".into())
325            }) {
326                Ok(crypto_key) => crypto_key,
327                Err(_) => return Err("Failed to parse crypto key".into())
328            };
329            Ok((public_key, private_key))
330        };
331
332        match input_encrypted_key_pair {
333            EncryptedKeyPair::V0(v0) => {                
334                // Handle V0 key pair       
335                let data: String;
336                if v0.encrypted_keys.is_empty() {
337                    data = v0.encrypted_key_pair;
338                } else {
339                    data = v0.encrypted_keys;
340                }    
341                let (public_key, private_key) = match decrypt_v2(
342                    &v0.iv, 
343                    &v0.salt, 
344                    &data
345                ) {
346                    Ok((public_key, private_key)) => (public_key, private_key),
347                    Err(_) => return Err("Failed to decrypt V0 key pair".into())
348                };
349                // Convert the public and private keys to JWK format
350                Ok(Key {
351                    crypto_key_pair: CryptoKeyPair {
352                        public_key: public_key,
353                        private_key: private_key,
354                    }
355                })
356            }
357            EncryptedKeyPair::V2(v2) => {
358                // Handle V2 key pair
359                let (public_key, private_key) = match decrypt_v2(&v2.iv, &v2.salt, &v2.data) {
360                    Ok((public_key, private_key)) => (public_key, private_key),
361                    Err(_) => return Err("Failed to decrypt V2 key pair".into())
362                };
363                // Convert the public and private keys to JWK format          
364                Ok(Key {
365                    crypto_key_pair: CryptoKeyPair {
366                        public_key: public_key,
367                        private_key: private_key,
368                    }
369                })
370            }
371        }
372    }
373}
klave_connector/key.rs

klave_connector/
key.rs
