modality_utils/
keypair.rs

1use anyhow::{anyhow, Result};
2use base58::ToBase58;
3use base64::prelude::*;
4use libp2p::identity::{ed25519, Keypair as Libp2pKeypair, PublicKey as Libp2pPublicKey};
5use libp2p_identity::PeerId;
6use regex::Regex;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::fs;
10
11use crate::encrypted_text::EncryptedText;
12use crate::json_stringify_deterministic::stringify_deterministic;
13
14#[derive(Clone)]
15pub enum KeypairOrPublicKey {
16    Keypair(Libp2pKeypair),
17    PublicKey(Libp2pPublicKey),
18}
19
20#[derive(Clone)]
21pub struct Keypair {
22    pub inner: KeypairOrPublicKey,
23}
24
25#[derive(Serialize, Deserialize)]
26pub struct KeypairJSON {
27    pub id: String,
28    pub public_key: String,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub private_key: Option<String>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub encrypted_private_key: Option<String>,
33}
34
35impl KeypairJSON {
36    pub fn id(&self) -> &str {
37        &self.id
38    }
39
40    pub fn public_key(&self) -> &str {
41        &self.public_key
42    }
43
44    pub fn private_key(&self) -> Option<&str> {
45        self.private_key.as_deref()
46    }
47
48    pub fn encrypted_private_key(&self) -> Option<&str> {
49        self.encrypted_private_key.as_deref()
50    }
51}
52
53impl Keypair {
54    pub fn new(key: KeypairOrPublicKey) -> Self {
55        Self { inner: key }
56    }
57
58    pub fn generate() -> Result<Self> {
59        let key = ed25519::Keypair::generate();
60        let libp2p_keypair = Libp2pKeypair::from(key);
61        Ok(Self::new(KeypairOrPublicKey::Keypair(libp2p_keypair)))
62    }
63
64    pub async fn as_ssh_private_pem(&self, _comment: &str) -> Result<String> {
65        // TODO: Implement SSH PEM conversion
66        unimplemented!("SSH PEM conversion not implemented yet")
67    }
68
69    pub fn as_ssh_dot_pub(&self, _comment: &str) -> Result<String> {
70        // TODO: Implement SSH public key conversion
71        unimplemented!("SSH public key conversion not implemented yet")
72    }
73
74    pub fn from_libp2p_keypair(keypair: libp2p_identity::Keypair) -> Result<Self> {
75        Ok(Self::new(KeypairOrPublicKey::Keypair(keypair)))
76    }
77
78    pub fn from_ssh_dot_pub(_public_key_str: &str, _key_type: &str) -> Result<Self> {
79        // TODO: Implement SSH public key parsing
80        unimplemented!("SSH public key parsing not implemented yet")
81    }
82
83    fn uint8_array_as_base58_identity(bytes: &[u8]) -> String {
84        let mut identity_hash = vec![0x00, bytes.len() as u8];
85        identity_hash.extend_from_slice(bytes);
86        identity_hash.to_base58()
87    }
88
89    pub fn private_key(&self) -> String {
90        self.private_key_as_base64_pad().ok().unwrap()
91    }
92
93    fn public_key_bytes(&self) -> Vec<u8> {
94        match &self.inner {
95            KeypairOrPublicKey::Keypair(k) => k.public().encode_protobuf(),
96            KeypairOrPublicKey::PublicKey(pk) => pk.encode_protobuf(),
97        }
98    }
99
100    pub fn public_key_as_base58_identity(&self) -> String {
101        Self::uint8_array_as_base58_identity(&self.public_key_bytes())
102    }
103
104    pub fn as_public_key_id(&self) -> String {
105        self.public_key_as_base58_identity()
106    }
107
108    pub fn as_public_address(&self) -> String {
109        self.public_key_as_base58_identity()
110    }
111
112    fn uint8_array_as_base64_pad(bytes: &[u8]) -> String {
113        BASE64_STANDARD.encode(bytes)
114    }
115
116    pub fn public_key_as_base64_pad(&self) -> String {
117        Self::uint8_array_as_base64_pad(&self.public_key_bytes())
118    }
119
120    pub fn private_key_as_base64_pad(&self) -> Result<String> {
121        match &self.inner {
122            KeypairOrPublicKey::Keypair(k) => {
123                Ok(Self::uint8_array_as_base64_pad(&k.to_protobuf_encoding()?))
124            }
125            KeypairOrPublicKey::PublicKey(_) => Err(anyhow!("No private key available")),
126        }
127    }
128
129    pub fn public_key_to_multiaddr_string(&self) -> String {
130        let public_key_id = self.public_key_as_base58_identity();
131        format!("/ed25519-pub/{}", public_key_id)
132    }
133
134    pub fn as_public_multiaddress(&self) -> String {
135        self.public_key_to_multiaddr_string()
136    }
137
138    pub fn from_public_key(public_key_id: &str, key_type: &str) -> Result<Self> {
139        Self::from_public_multiaddress(&format!("/{}-pub/{}", key_type, public_key_id))
140    }
141
142    pub fn from_public_multiaddress(multiaddress: &str) -> Result<Self> {
143        let re = Regex::new(r"^(.+)-pub/(.+)$").unwrap();
144        let captures = re
145            .captures(multiaddress)
146            .ok_or_else(|| anyhow!("Invalid multiaddress format"))?;
147
148        // let _key_type = captures.get(1)
149        //     .ok_or_else(|| anyhow!("Failed to extract key type"))?
150        //     .as_str();
151        let peer_id_str = captures
152            .get(2)
153            .ok_or_else(|| anyhow!("Failed to extract public key ID"))?
154            .as_str();
155
156        // Parse the peer ID
157        let peer_id = peer_id_str
158            .parse::<PeerId>()
159            .map_err(|e| anyhow!("Failed to parse peer ID: {:?}", e))?;
160
161        let public_key = libp2p::identity::PublicKey::try_decode_protobuf(&peer_id.to_bytes())
162            .map_err(|e| anyhow!("Failed to decode public key from peer ID: {}", e))?;
163
164        Ok(Self::new(KeypairOrPublicKey::PublicKey(public_key)))
165    }
166
167    pub fn from_json(json: &KeypairJSON) -> Result<Self> {
168        if let Some(private_key) = &json.private_key {
169            let key_bytes = BASE64_STANDARD.decode(private_key)?;
170            let key = Libp2pKeypair::from_protobuf_encoding(&key_bytes)?;
171            Ok(Self::new(KeypairOrPublicKey::Keypair(key)))
172        } else {
173            let key_bytes = BASE64_STANDARD.decode(&json.public_key)?;
174            let public_key = Libp2pPublicKey::try_decode_protobuf(&key_bytes)?;
175            Ok(Self::new(KeypairOrPublicKey::PublicKey(public_key)))
176        }
177    }
178
179    pub fn from_json_string(json_str: &str) -> Result<Self> {
180        let json: KeypairJSON = serde_json::from_str(json_str)?;
181        Self::from_json(&json)
182    }
183
184    pub fn from_encrypted_json_file(filepath: &str, password: &str) -> Result<Self> {
185        let json_str = fs::read_to_string(filepath)?;
186        let json: KeypairJSON = serde_json::from_str(&json_str)?;
187
188        if let Some(encrypted_key) = json.encrypted_private_key() {
189            let decrypted_key = EncryptedText::decrypt(encrypted_key, password).ok();
190            let decrypted_json = KeypairJSON {
191                id: json.id().to_string(),
192                public_key: json.public_key().to_string(),
193                private_key: decrypted_key,
194                encrypted_private_key: None,
195            };
196            Self::from_json(&decrypted_json)
197        } else {
198            Self::from_json(&json)
199        }
200    }
201
202    pub fn from_json_file(filepath: &str) -> Result<Self> {
203        let json_str = fs::read_to_string(filepath)?;
204        Self::from_json_string(&json_str)
205    }
206
207    pub fn as_public_json(&self) -> Result<KeypairJSON> {
208        Ok(KeypairJSON {
209            id: self.public_key_as_base58_identity(),
210            public_key: self.public_key_as_base64_pad(),
211            private_key: None,
212            encrypted_private_key: None,
213        })
214    }
215
216    pub fn as_public_json_string(&self) -> Result<String> {
217        let json = self.as_public_json()?;
218        Ok(serde_json::to_string(&json)?)
219    }
220
221    pub fn as_public_json_file(&self, path: &str) -> Result<()> {
222        let json_string = self.as_public_json_string()?;
223        fs::write(path, json_string)?;
224        Ok(())
225    }
226
227    pub fn as_json(&self) -> Result<KeypairJSON> {
228        let private_key = self.private_key_as_base64_pad().ok();
229        let encrypted_private_key = if private_key.is_some() {
230            None
231        } else {
232            Some("".to_string()) // Use actual encrypted_private_key if available
233        };
234
235        Ok(KeypairJSON {
236            id: self.public_key_as_base58_identity(),
237            public_key: self.public_key_as_base64_pad(),
238            private_key,
239            encrypted_private_key,
240        })
241    }
242
243    pub fn as_json_string(&self) -> Result<String> {
244        let json = self.as_json()?;
245        Ok(serde_json::to_string(&json)?)
246    }
247
248    pub fn as_json_file(&self, path: &str) -> Result<()> {
249        let json_string = self.as_json_string()?;
250        fs::write(path, json_string)?;
251        Ok(())
252    }
253
254    pub fn as_encrypted_json(&self, password: &str) -> Result<KeypairJSON> {
255        let enc_pk = EncryptedText::encrypt(&self.private_key_as_base64_pad()?, password).ok();
256        Ok(KeypairJSON {
257            id: self.public_key_as_base58_identity(),
258            public_key: self.public_key_as_base64_pad(),
259            private_key: None,
260            encrypted_private_key: enc_pk,
261        })
262    }
263
264    pub fn as_encrypted_json_string(&self, password: &str) -> Result<String> {
265        let json = self.as_encrypted_json(password)?;
266        Ok(serde_json::to_string(&json)?)
267    }
268
269    pub fn as_encrypted_json_file(&self, path: &str, password: &str) -> Result<()> {
270        let json_string = self.as_encrypted_json_string(password)?;
271        fs::write(path, json_string)?;
272        Ok(())
273    }
274
275    pub fn sign_bytes(&self, bytes: &[u8]) -> Result<Vec<u8>> {
276        match &self.inner {
277            KeypairOrPublicKey::Keypair(k) => Ok(k.sign(bytes)?),
278            KeypairOrPublicKey::PublicKey(_) => Err(anyhow!("Cannot sign with public key only")),
279        }
280    }
281
282    pub fn sign_string(&self, s: &str) -> Result<Vec<u8>> {
283        self.sign_bytes(s.as_bytes())
284    }
285
286    pub fn sign_string_as_base64_pad(&self, s: &str) -> Result<String> {
287        let signature = self.sign_string(s)?;
288        Ok(BASE64_STANDARD.encode(signature))
289    }
290
291    pub fn sign_json(&self, json: &Value) -> Result<String> {
292        let str = stringify_deterministic(json, None);
293        self.sign_string_as_base64_pad(&str)
294    }
295
296    pub fn sign_json_element(&self, json: &mut Value, name: &str, suffix: &str) -> Result<()> {
297        let signature = self.sign_json(&json[name])?;
298        json[format!("{}{}", name, suffix)] = Value::String(signature);
299        Ok(())
300    }
301
302    pub fn sign_json_as_key(&self, json: &mut Value, key: &str) -> Result<()> {
303        let signature = self.sign_json(json)?;
304        json[key] = Value::String(signature);
305        Ok(())
306    }
307
308    pub fn verify_signature_for_bytes(&self, signature: &str, bytes: &[u8]) -> Result<bool> {
309        let signature_bytes = BASE64_STANDARD.decode(signature)?;
310        match &self.inner {
311            KeypairOrPublicKey::Keypair(k) => Ok(k.public().verify(bytes, &signature_bytes)),
312            KeypairOrPublicKey::PublicKey(pk) => Ok(pk.verify(bytes, &signature_bytes)),
313        }
314    }
315
316    pub fn verify_signature_for_string(&self, signature: &str, s: &str) -> Result<bool> {
317        self.verify_signature_for_bytes(signature, s.as_bytes())
318    }
319
320    pub fn verify_json(&self, signature: &str, json: &Value) -> Result<bool> {
321        let str = stringify_deterministic(json, None);
322        self.verify_signature_for_string(signature, &str)
323    }
324
325    pub fn verify_json_with_signature_key(
326        &self,
327        json: &Value,
328        signature_key: &str,
329    ) -> Result<bool> {
330        if let Value::Object(map) = json {
331            let signature = map
332                .get(signature_key)
333                .ok_or_else(|| anyhow!("Signature key not found"))?
334                .as_str()
335                .ok_or_else(|| anyhow!("Signature must be a string"))?;
336
337            let mut json_without_signature = json.clone();
338            if let Value::Object(map_without_signature) = &mut json_without_signature {
339                map_without_signature.remove(signature_key);
340            }
341
342            let stringified = stringify_deterministic(&json_without_signature, None);
343            self.verify_signature_for_string(signature, &stringified)
344        } else {
345            Err(anyhow!("Input must be a JSON object"))
346        }
347    }
348
349    pub fn verify_signatures_in_json(&self, json: &Value, suffix: Option<&str>) -> Result<bool> {
350        let suffix = suffix.unwrap_or(".signature");
351        let suffix_regex = Regex::new(&format!(r"(.+){}$", regex::escape(suffix)))?;
352
353        if let Value::Object(map) = json {
354            for (key, value) in map {
355                if let Some(captures) = suffix_regex.captures(key) {
356                    let original_key = captures.get(1).unwrap().as_str();
357                    if let Some(original_value) = map.get(original_key) {
358                        let signature = value
359                            .as_str()
360                            .ok_or_else(|| anyhow!("Signature must be a string"))?;
361                        let stringified = stringify_deterministic(original_value, None);
362
363                        if !self.verify_signature_for_string(signature, &stringified)? {
364                            return Ok(false);
365                        }
366                    } else {
367                        return Ok(false); // Original value not found
368                    }
369                }
370            }
371            Ok(true) // All signatures verified successfully
372        } else {
373            Err(anyhow!("Input must be a JSON object"))
374        }
375    }
376}