1use alloc::vec::Vec;
2use core::fmt;
3
4use crate::ed25519::{Ed25519PrivateKey, Ed25519PublicKey};
5use crate::hkdf;
6use crate::sha256;
7use crate::token::{Token, TokenError};
8use crate::x25519::{X25519PrivateKey, X25519PublicKey};
9use crate::Rng;
10
11pub const KEYSIZE: usize = 512; pub const DERIVED_KEY_LENGTH: usize = 64; pub const TRUNCATED_HASHLENGTH: usize = 128; #[derive(Debug)]
16pub enum CryptoError {
17 NoPrivateKey,
18 NoPublicKey,
19 TokenError(TokenError),
20 HkdfError(hkdf::HkdfError),
21 InvalidCiphertext,
22}
23
24impl fmt::Display for CryptoError {
25 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26 match self {
27 CryptoError::NoPrivateKey => write!(f, "No private key"),
28 CryptoError::NoPublicKey => write!(f, "No public key"),
29 CryptoError::TokenError(e) => write!(f, "Token error: {}", e),
30 CryptoError::HkdfError(e) => write!(f, "HKDF error: {}", e),
31 CryptoError::InvalidCiphertext => write!(f, "Invalid ciphertext"),
32 }
33 }
34}
35
36pub struct Identity {
37 prv: Option<X25519PrivateKey>,
38 sig_prv: Option<Ed25519PrivateKey>,
39 pub_key: Option<X25519PublicKey>,
40 sig_pub: Option<Ed25519PublicKey>,
41 hash: [u8; 16],
42}
43
44impl Identity {
45 pub fn new(rng: &mut dyn Rng) -> Self {
46 let prv = X25519PrivateKey::generate(rng);
47 let sig_prv = Ed25519PrivateKey::generate(rng);
48
49 let pub_key = prv.public_key();
50 let sig_pub = sig_prv.public_key();
51
52 let mut pub_bytes = [0u8; 64];
53 pub_bytes[..32].copy_from_slice(&pub_key.public_bytes());
54 pub_bytes[32..].copy_from_slice(&sig_pub.public_bytes());
55
56 let hash = truncated_hash(&pub_bytes);
57
58 Identity {
59 prv: Some(prv),
60 sig_prv: Some(sig_prv),
61 pub_key: Some(pub_key),
62 sig_pub: Some(sig_pub),
63 hash,
64 }
65 }
66
67 pub fn from_private_key(prv_bytes: &[u8; 64]) -> Self {
68 let x_prv_bytes: [u8; 32] = prv_bytes[..32].try_into().unwrap();
69 let ed_seed: [u8; 32] = prv_bytes[32..].try_into().unwrap();
70
71 let prv = X25519PrivateKey::from_bytes(&x_prv_bytes);
72 let sig_prv = Ed25519PrivateKey::from_bytes(&ed_seed);
73
74 let pub_key = prv.public_key();
75 let sig_pub = sig_prv.public_key();
76
77 let mut pub_bytes = [0u8; 64];
78 pub_bytes[..32].copy_from_slice(&pub_key.public_bytes());
79 pub_bytes[32..].copy_from_slice(&sig_pub.public_bytes());
80
81 let hash = truncated_hash(&pub_bytes);
82
83 Identity {
84 prv: Some(prv),
85 sig_prv: Some(sig_prv),
86 pub_key: Some(pub_key),
87 sig_pub: Some(sig_pub),
88 hash,
89 }
90 }
91
92 pub fn from_public_key(pub_bytes: &[u8; 64]) -> Self {
93 let x_pub_bytes: [u8; 32] = pub_bytes[..32].try_into().unwrap();
94 let ed_pub_bytes: [u8; 32] = pub_bytes[32..].try_into().unwrap();
95
96 let pub_key = X25519PublicKey::from_bytes(&x_pub_bytes);
97 let sig_pub = Ed25519PublicKey::from_bytes(&ed_pub_bytes);
98
99 let hash = truncated_hash(pub_bytes);
100
101 Identity {
102 prv: None,
103 sig_prv: None,
104 pub_key: Some(pub_key),
105 sig_pub: Some(sig_pub),
106 hash,
107 }
108 }
109
110 pub fn get_private_key(&self) -> Option<[u8; 64]> {
111 match (&self.prv, &self.sig_prv) {
112 (Some(prv), Some(sig_prv)) => {
113 let mut result = [0u8; 64];
114 result[..32].copy_from_slice(&prv.private_bytes());
115 result[32..].copy_from_slice(&sig_prv.private_bytes());
116 Some(result)
117 }
118 _ => None,
119 }
120 }
121
122 pub fn get_public_key(&self) -> Option<[u8; 64]> {
123 match (&self.pub_key, &self.sig_pub) {
124 (Some(pub_key), Some(sig_pub)) => {
125 let mut result = [0u8; 64];
126 result[..32].copy_from_slice(&pub_key.public_bytes());
127 result[32..].copy_from_slice(&sig_pub.public_bytes());
128 Some(result)
129 }
130 _ => None,
131 }
132 }
133
134 pub fn hash(&self) -> &[u8; 16] {
135 &self.hash
136 }
137
138 pub fn encrypt(&self, plaintext: &[u8], rng: &mut dyn Rng) -> Result<Vec<u8>, CryptoError> {
139 let pub_key = self.pub_key.as_ref().ok_or(CryptoError::NoPublicKey)?;
140
141 let ephemeral = X25519PrivateKey::generate(rng);
142 let ephemeral_pub_bytes = ephemeral.public_key().public_bytes();
143 let shared_key = ephemeral.exchange(pub_key);
144
145 let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
146 .map_err(CryptoError::HkdfError)?;
147
148 let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
149 let ciphertext = token.encrypt(plaintext, rng);
150
151 let mut result = Vec::with_capacity(32 + ciphertext.len());
152 result.extend_from_slice(&ephemeral_pub_bytes);
153 result.extend_from_slice(&ciphertext);
154 Ok(result)
155 }
156
157 pub fn encrypt_deterministic(
159 &self,
160 plaintext: &[u8],
161 ephemeral_prv: &[u8; 32],
162 iv: &[u8; 16],
163 ) -> Result<Vec<u8>, CryptoError> {
164 let pub_key = self.pub_key.as_ref().ok_or(CryptoError::NoPublicKey)?;
165
166 let ephemeral = X25519PrivateKey::from_bytes(ephemeral_prv);
167 let ephemeral_pub_bytes = ephemeral.public_key().public_bytes();
168 let shared_key = ephemeral.exchange(pub_key);
169
170 let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
171 .map_err(CryptoError::HkdfError)?;
172
173 let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
174 let ciphertext = token.encrypt_with_iv(plaintext, iv);
175
176 let mut result = Vec::with_capacity(32 + ciphertext.len());
177 result.extend_from_slice(&ephemeral_pub_bytes);
178 result.extend_from_slice(&ciphertext);
179 Ok(result)
180 }
181
182 pub fn decrypt(&self, ciphertext_token: &[u8]) -> Result<Vec<u8>, CryptoError> {
183 let prv = self.prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
184
185 if ciphertext_token.len() <= KEYSIZE / 8 / 2 {
186 return Err(CryptoError::InvalidCiphertext);
187 }
188
189 let peer_pub_bytes: [u8; 32] = ciphertext_token[..32].try_into().unwrap();
190 let peer_pub = X25519PublicKey::from_bytes(&peer_pub_bytes);
191 let ciphertext = &ciphertext_token[32..];
192
193 let shared_key = prv.exchange(&peer_pub);
194
195 let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
196 .map_err(CryptoError::HkdfError)?;
197
198 let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
199 token.decrypt(ciphertext).map_err(CryptoError::TokenError)
200 }
201
202 pub fn sign(&self, message: &[u8]) -> Result<[u8; 64], CryptoError> {
203 let sig_prv = self.sig_prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
204 Ok(sig_prv.sign(message))
205 }
206
207 pub fn verify(&self, signature: &[u8; 64], message: &[u8]) -> bool {
208 match &self.sig_pub {
209 Some(sig_pub) => sig_pub.verify(signature, message),
210 None => false,
211 }
212 }
213}
214
215fn truncated_hash(data: &[u8]) -> [u8; 16] {
216 let full = sha256::sha256(data);
217 let mut result = [0u8; 16];
218 result.copy_from_slice(&full[..16]);
219 result
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225 use crate::FixedRng;
226
227 #[test]
228 fn test_identity_key_roundtrip() {
229 let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
230 let id = Identity::new(&mut rng);
231 let prv_bytes = id.get_private_key().unwrap();
232 let id2 = Identity::from_private_key(&prv_bytes);
233 assert_eq!(id.get_public_key().unwrap(), id2.get_public_key().unwrap());
234 }
235
236 #[test]
237 fn test_identity_hash() {
238 let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
239 let id = Identity::new(&mut rng);
240 let pub_key = id.get_public_key().unwrap();
241 let expected_hash = truncated_hash(&pub_key);
242 assert_eq!(*id.hash(), expected_hash);
243 }
244
245 #[test]
246 fn test_identity_encrypt_decrypt_roundtrip() {
247 let mut rng = FixedRng::new(&(0..128).collect::<Vec<u8>>());
248 let id = Identity::new(&mut rng);
249 let plaintext = b"Hello, Reticulum! This is a test of the encrypt/decrypt pipeline.";
250 let mut rng2 = FixedRng::new(&(128..255).collect::<Vec<u8>>());
251 let ciphertext = id.encrypt(plaintext, &mut rng2).unwrap();
252 let decrypted = id.decrypt(&ciphertext).unwrap();
253 assert_eq!(decrypted, plaintext);
254 }
255
256 #[test]
257 fn test_identity_sign_verify() {
258 let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
259 let id = Identity::new(&mut rng);
260 let msg = b"Sign this message";
261 let sig = id.sign(msg).unwrap();
262 assert!(id.verify(&sig, msg));
263 assert!(!id.verify(&sig, b"Wrong message"));
264 }
265}