lm_rs/
installation_key.rs

1use anyhow::Result;
2use base64::{engine::general_purpose::STANDARD, Engine as _};
3use p256::{
4    ecdsa::{signature::Signer, Signature, SigningKey},
5    elliptic_curve::rand_core::OsRng,
6    pkcs8::EncodePublicKey,
7    SecretKey,
8};
9use serde::{Deserialize, Serialize};
10use sha2::{Digest, Sha256};
11use std::time::{SystemTime, UNIX_EPOCH};
12use uuid::Uuid;
13
14/// Installation key containing cryptographic material for authentication
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct InstallationKey {
17    /// 32-byte secret for proof generation
18    #[serde(
19        serialize_with = "serialize_bytes_as_base64",
20        deserialize_with = "deserialize_base64_as_bytes"
21    )]
22    pub secret: Vec<u8>,
23    /// ECDSA private key on P-256 curve  
24    #[serde(
25        serialize_with = "serialize_signing_key_as_base64",
26        deserialize_with = "deserialize_base64_as_signing_key"
27    )]
28    pub private_key: SigningKey,
29    /// Installation ID (UUID)
30    pub installation_id: String,
31}
32
33impl InstallationKey {
34    /// Get the public key in base64-encoded DER format
35    pub fn public_key_b64(&self) -> String {
36        let verifying_key = *self.private_key.verifying_key();
37        // Use DER encoding to match Python implementation
38        let public_key_der = verifying_key.to_public_key_der().unwrap();
39        STANDARD.encode(public_key_der.as_bytes())
40    }
41
42    /// Get the base string: installation_id.sha256(public_key_der_bytes)
43    pub fn base_string(&self) -> String {
44        let verifying_key = *self.private_key.verifying_key();
45        // Use DER encoding to match Python implementation
46        let public_key_der = verifying_key.to_public_key_der().unwrap();
47        let mut hasher = Sha256::new();
48        hasher.update(public_key_der.as_bytes());
49        let pub_hash = hasher.finalize();
50        let pub_hash_b64 = STANDARD.encode(pub_hash);
51        format!("{}.{}", self.installation_id, pub_hash_b64)
52    }
53}
54
55/// Generate installation key from installation ID following the Python pattern
56pub fn generate_installation_key(installation_id: String) -> Result<InstallationKey> {
57    // Generate ECDSA private key on P-256 curve
58    let secret_key = SecretKey::random(&mut OsRng);
59    let signing_key = SigningKey::from(secret_key);
60    let verifying_key = *signing_key.verifying_key();
61
62    // Get public key bytes in DER format to match Python implementation
63    let public_key_der = verifying_key.to_public_key_der().unwrap();
64    let pub_b64 = STANDARD.encode(public_key_der.as_bytes());
65
66    // Create installation hash
67    let mut hasher = Sha256::new();
68    hasher.update(installation_id.as_bytes());
69    let inst_hash = hasher.finalize();
70    let inst_hash_b64 = STANDARD.encode(inst_hash);
71
72    // Create triple: installation_id.pub_b64.inst_hash_b64
73    let triple = format!("{}.{}.{}", installation_id, pub_b64, inst_hash_b64);
74
75    // Generate 32-byte secret from triple
76    let mut secret_hasher = Sha256::new();
77    secret_hasher.update(triple.as_bytes());
78    let secret_bytes = secret_hasher.finalize();
79
80    Ok(InstallationKey {
81        secret: secret_bytes.to_vec(),
82        private_key: signing_key,
83        installation_id,
84    })
85}
86
87/// Generate a new random installation ID (UUID v4)
88pub fn generate_installation_id() -> String {
89    Uuid::new_v4().to_string().to_lowercase()
90}
91
92/// La Marzocco's custom proof generation algorithm (Y5.e equivalent)
93pub fn generate_request_proof(base_string: &str, secret32: &[u8]) -> Result<String> {
94    if secret32.len() != 32 {
95        return Err(anyhow::anyhow!("secret must be 32 bytes"));
96    }
97
98    // Use a fixed-size stack buffer to avoid heap allocation
99    let mut work = [0u8; 32];
100    work.copy_from_slice(secret32);
101
102    for byte_val in base_string.as_bytes() {
103        let idx = (*byte_val as usize) % 32;
104        let shift_idx = (idx + 1) % 32;
105        let shift_amount = work[shift_idx] & 7; // 0-7 bit shift
106
107        // XOR then rotate left
108        let xor_result = byte_val ^ work[idx];
109        let rotated = if shift_amount == 0 {
110            xor_result
111        } else {
112            (xor_result << shift_amount) | (xor_result >> (8 - shift_amount))
113        };
114        work[idx] = rotated;
115    }
116
117    let mut hasher = Sha256::new();
118    hasher.update(work);
119    let result = hasher.finalize();
120    Ok(STANDARD.encode(result))
121}
122
123/// Generate extra headers for normal API calls after authentication
124pub fn generate_extra_request_headers(
125    installation_key: &InstallationKey,
126) -> Result<Vec<(String, String)>> {
127    // Generate nonce and timestamp
128    let nonce = Uuid::new_v4().to_string().to_lowercase();
129    let timestamp = SystemTime::now()
130        .duration_since(UNIX_EPOCH)?
131        .as_millis()
132        .to_string();
133
134    // Create proof using Y5.e algorithm: installation_id.nonce.timestamp
135    let proof_input = format!(
136        "{}.{}.{}",
137        installation_key.installation_id, nonce, timestamp
138    );
139    let proof = generate_request_proof(&proof_input, &installation_key.secret)?;
140
141    // Create signature data: installation_id.nonce.timestamp.proof
142    let signature_data = format!("{}.{}", proof_input, proof);
143
144    // Sign with ECDSA
145    let signature: Signature = installation_key.private_key.sign(signature_data.as_bytes());
146    let signature_b64 = STANDARD.encode(signature.to_der());
147
148    // Return headers
149    Ok(vec![
150        (
151            "X-App-Installation-Id".to_string(),
152            installation_key.installation_id.clone(),
153        ),
154        ("X-Timestamp".to_string(), timestamp),
155        ("X-Nonce".to_string(), nonce),
156        ("X-Request-Signature".to_string(), signature_b64),
157    ])
158}
159
160/// Serde helper functions for serialization
161fn serialize_bytes_as_base64<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
162where
163    S: serde::Serializer,
164{
165    let b64 = STANDARD.encode(bytes);
166    serializer.serialize_str(&b64)
167}
168
169fn deserialize_base64_as_bytes<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
170where
171    D: serde::Deserializer<'de>,
172{
173    let s = String::deserialize(deserializer)?;
174    STANDARD.decode(s).map_err(serde::de::Error::custom)
175}
176
177fn serialize_signing_key_as_base64<S>(
178    signing_key: &SigningKey,
179    serializer: S,
180) -> Result<S::Ok, S::Error>
181where
182    S: serde::Serializer,
183{
184    let bytes = signing_key.to_bytes();
185    let b64 = STANDARD.encode(bytes);
186    serializer.serialize_str(&b64)
187}
188
189fn deserialize_base64_as_signing_key<'de, D>(deserializer: D) -> Result<SigningKey, D::Error>
190where
191    D: serde::Deserializer<'de>,
192{
193    let s = String::deserialize(deserializer)?;
194    let bytes = STANDARD.decode(s).map_err(serde::de::Error::custom)?;
195    let secret_key = SecretKey::from_slice(&bytes).map_err(serde::de::Error::custom)?;
196    Ok(SigningKey::from(secret_key))
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    #[test]
204    fn test_installation_id_generation() {
205        let id1 = generate_installation_id();
206        let id2 = generate_installation_id();
207
208        // IDs should be different
209        assert_ne!(id1, id2);
210
211        // IDs should be valid UUIDs (36 characters with dashes)
212        assert_eq!(id1.len(), 36);
213        assert_eq!(id2.len(), 36);
214        assert!(id1.contains('-'));
215        assert!(id2.contains('-'));
216    }
217
218    #[test]
219    fn test_installation_key_generation() {
220        let installation_id = "test-installation-id".to_string();
221        let key = generate_installation_key(installation_id.clone()).unwrap();
222
223        assert_eq!(key.installation_id, installation_id);
224        assert_eq!(key.secret.len(), 32);
225
226        // Test that we can get public key
227        let pub_key_b64 = key.public_key_b64();
228        assert!(!pub_key_b64.is_empty());
229
230        // Test base string format
231        let base_string = key.base_string();
232        assert!(base_string.starts_with(&installation_id));
233        assert!(base_string.contains('.'));
234    }
235
236    #[test]
237    fn test_request_proof_generation() {
238        let secret = vec![0u8; 32]; // All zeros for testing
239        let base_string = "test.base.string";
240
241        let proof = generate_request_proof(base_string, &secret).unwrap();
242        assert!(!proof.is_empty());
243
244        // Should be base64 encoded SHA256 hash (44 characters)
245        assert_eq!(proof.len(), 44);
246    }
247
248    #[test]
249    fn test_request_proof_error_on_wrong_secret_size() {
250        let secret = vec![0u8; 31]; // Wrong size
251        let base_string = "test";
252
253        let result = generate_request_proof(base_string, &secret);
254        assert!(result.is_err());
255        assert!(result.unwrap_err().to_string().contains("32 bytes"));
256    }
257
258    #[test]
259    fn test_extra_request_headers_generation() {
260        let installation_id = "test-id".to_string();
261        let key = generate_installation_key(installation_id.clone()).unwrap();
262
263        let headers = generate_extra_request_headers(&key).unwrap();
264
265        // Should have 4 headers
266        assert_eq!(headers.len(), 4);
267
268        // Check header names
269        let header_names: Vec<String> = headers.iter().map(|(k, _)| k.clone()).collect();
270        assert!(header_names.contains(&"X-App-Installation-Id".to_string()));
271        assert!(header_names.contains(&"X-Timestamp".to_string()));
272        assert!(header_names.contains(&"X-Nonce".to_string()));
273        assert!(header_names.contains(&"X-Request-Signature".to_string()));
274
275        // Check installation ID matches
276        let installation_id_header = headers
277            .iter()
278            .find(|(k, _)| k == "X-App-Installation-Id")
279            .unwrap();
280        assert_eq!(installation_id_header.1, installation_id);
281    }
282
283    #[test]
284    fn test_cross_language_compatibility() {
285        use base64::{engine::general_purpose::STANDARD, Engine as _};
286        use p256::{ecdsa::SigningKey, pkcs8::DecodePrivateKey};
287
288        // Test data generated from Python implementation
289        let installation_id = "test-installation-id-12345";
290        let expected_secret = "liWQoHaK8Sg+auapkWedGzGs/8HDt6JCIP3tw2c8WYA=";
291        let private_key_der = "MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQg25BwIri/hrmMdfAvHgbFk9TQ/nmA70OYEdmrFuhbux+hRANCAASSLyaAtUj6jGO/F2VQJMN9XNGQkNMZNktiENHlVgMKaTrTMXuR/dyQU/4boce+LqcHoOBjZVDYz5JsXyKM6qyE";
292        let expected_public_key_der = "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEki8mgLVI+oxjvxdlUCTDfVzRkJDTGTZLYhDR5VYDCmk60zF7kf3ckFP+G6HHvi6nB6DgY2VQ2M+SbF8ijOqshA==";
293        let expected_base_string =
294            "test-installation-id-12345.eZhZZDD3ciI13+s7zV9QlgLW9Eo+lDKGJAKUn8SpAtA=";
295
296        // Decode the private key and reconstruct InstallationKey
297        let private_key_bytes = STANDARD.decode(private_key_der).unwrap();
298        let signing_key = SigningKey::from_pkcs8_der(&private_key_bytes).unwrap();
299        let secret_bytes = STANDARD.decode(expected_secret).unwrap();
300
301        let key = InstallationKey {
302            secret: secret_bytes,
303            private_key: signing_key,
304            installation_id: installation_id.to_string(),
305        };
306
307        // Test that our implementation produces the same results
308        assert_eq!(key.public_key_b64(), expected_public_key_der);
309        assert_eq!(key.base_string(), expected_base_string);
310
311        // Test proof generation with same input as Python
312        let test_proof_input = "test-installation-id-12345.test-nonce.1234567890";
313        let expected_proof = "DXFZPKXiSDRih9U43F+YhJGn2DRt05XLLY9W9dtGl6g=";
314        let proof = generate_request_proof(test_proof_input, &key.secret).unwrap();
315        assert_eq!(proof, expected_proof);
316    }
317
318    #[test]
319    fn test_installation_key_serialization() {
320        let installation_id = "test-id".to_string();
321        let key = generate_installation_key(installation_id.clone()).unwrap();
322
323        // Test JSON serialization
324        let json = serde_json::to_string(&key).unwrap();
325        assert!(!json.is_empty());
326
327        // Test deserialization
328        let deserialized: InstallationKey = serde_json::from_str(&json).unwrap();
329        assert_eq!(deserialized.installation_id, key.installation_id);
330        assert_eq!(deserialized.secret, key.secret);
331
332        // Verify keys work the same
333        assert_eq!(deserialized.public_key_b64(), key.public_key_b64());
334        assert_eq!(deserialized.base_string(), key.base_string());
335    }
336}
lm_rs/installation_key.rs

lm_rs/
installation_key.rs
