diff --git a/self_encryption/__init__.py b/self_encryption/__init__.py index e5e7ec7ea..ae3927db4 100644 --- a/self_encryption/__init__.py +++ b/self_encryption/__init__.py @@ -1,5 +1,16 @@ try: - from ._self_encryption import * + from ._self_encryption import ( + PyDataMap as DataMap, + PyXorName as XorName, + EncryptResult, + encrypt_from_file, + decrypt_from_storage, + streaming_decrypt_from_storage, + MIN_CHUNK_SIZE, + MIN_ENCRYPTABLE_BYTES, + MAX_CHUNK_SIZE, + COMPRESSION_QUALITY, + ) from .cli import cli except ImportError as e: import sys @@ -7,8 +18,15 @@ raise __all__ = [ - 'cli', + 'DataMap', + 'XorName', + 'EncryptResult', 'encrypt_from_file', 'decrypt_from_storage', 'streaming_decrypt_from_storage', + 'MIN_CHUNK_SIZE', + 'MIN_ENCRYPTABLE_BYTES', + 'MAX_CHUNK_SIZE', + 'COMPRESSION_QUALITY', + 'cli', ] \ No newline at end of file diff --git a/self_encryption/cli.py b/self_encryption/cli.py index 4695094cb..274ffda56 100644 --- a/self_encryption/cli.py +++ b/self_encryption/cli.py @@ -12,17 +12,12 @@ from typing import Optional import sys -from self_encryption import ( - DataMap, - EncryptedChunk, - XorName, - encrypt, +from ._self_encryption import ( + PyDataMap as DataMap, + PyXorName as XorName, encrypt_from_file, - decrypt, decrypt_from_storage, - shrink_data_map, streaming_decrypt_from_storage, - verify_chunk, ) def print_error(message: str): diff --git a/src/lib.rs b/src/lib.rs index f0dc01342..f2bb9436c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -68,7 +68,6 @@ unused_comparisons, unused_features, unused_parens, - while_true, )] #![cfg_attr(not(feature = "python"), deny(warnings))] #![warn( @@ -532,7 +531,7 @@ where // Create a buffered reader with a reasonable buffer size let mut reader = BufReader::with_capacity(1024 * 1024, file); - + // Read all chunks first to get their hashes let mut chunks = Vec::with_capacity(num_chunks); for chunk_index in 0..num_chunks { @@ -540,7 +539,7 @@ where let chunk_size = end - start; let mut chunk_data = vec![0u8; chunk_size]; reader.read_exact(&mut chunk_data)?; - + let hash = XorName::from_content(&chunk_data); chunks.push(crate::chunk::RawChunk { index: chunk_index, @@ -548,20 +547,20 @@ where hash, }); } - + // Process chunks and store them immediately let data_map = encrypt::encrypt_stream(chunks.clone())?; - + // Now encrypt and store each chunk let src_hashes: Vec<_> = chunks.iter().map(|c| c.hash).collect(); - + for chunk in chunks { let pki = get_pad_key_and_iv(chunk.index, &src_hashes); let encrypted_content = encrypt::encrypt_chunk(chunk.data, pki)?; let hash = XorName::from_content(&encrypted_content); chunk_store(hash, encrypted_content)?; } - + // Shrink the data map and store additional chunks if needed let (shrunk_data_map, _) = shrink_data_map(data_map, |hash, content| { chunk_store(hash, content)?; @@ -624,10 +623,7 @@ where .iter() .map(|info| { let content = chunk_cache.get(&info.dst_hash).ok_or_else(|| { - Error::Generic(format!( - "Chunk not found for hash: {:?}", - info.dst_hash - )) + Error::Generic(format!("Chunk not found for hash: {:?}", info.dst_hash)) })?; Ok(EncryptedChunk { content: content.clone(), @@ -756,10 +752,10 @@ mod tests { std::fs::write(&file_path, small_data)?; let store = |_: XorName, _: Bytes| -> Result<()> { Ok(()) }; - + let result = streaming_encrypt_from_file(&file_path, store); assert!(result.is_err()); - + Ok(()) } @@ -773,7 +769,7 @@ mod tests { let storage = Arc::new(Mutex::new(HashMap::new())); let storage_clone = storage.clone(); - + let store = move |hash: XorName, content: Bytes| -> Result<()> { let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec()); Ok(()) @@ -781,7 +777,7 @@ mod tests { let data_map = streaming_encrypt_from_file(&file_path, store)?; assert!(data_map.chunk_identifiers.len() >= 3); - + Ok(()) } @@ -797,7 +793,7 @@ mod tests { let storage = Arc::new(Mutex::new(HashMap::new())); let storage_clone = storage.clone(); - + let store = move |hash: XorName, content: Bytes| -> Result<()> { let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec()); Ok(()) @@ -806,16 +802,22 @@ mod tests { // First get the number of chunks directly let bytes = Bytes::from(large_data.clone()); let (num_chunks, _) = chunk::batch_chunks(bytes); - assert!(num_chunks > 3, - "Should have more than 3 chunks before shrinking. Got: {}", num_chunks); + assert!( + num_chunks > 3, + "Should have more than 3 chunks before shrinking. Got: {}", + num_chunks + ); // Now test the streaming encryption let data_map = streaming_encrypt_from_file(&file_path, store)?; - + // After shrinking, should be exactly 3 chunks - assert_eq!(data_map.chunk_identifiers.len(), 3, - "Final data map should have exactly 3 chunks after shrinking"); - + assert_eq!( + data_map.chunk_identifiers.len(), + 3, + "Final data map should have exactly 3 chunks after shrinking" + ); + // Verify chunk indices are sequential let mut prev_index = None; for chunk_info in &data_map.chunk_identifiers { @@ -824,7 +826,7 @@ mod tests { } prev_index = Some(chunk_info.index); } - + Ok(()) } @@ -851,7 +853,7 @@ mod tests { let result = streaming_encrypt_from_file(&file_path, store); assert!(result.is_err()); - + Ok(()) } @@ -865,7 +867,7 @@ mod tests { let storage = Arc::new(Mutex::new(HashMap::new())); let storage_clone = storage.clone(); - + let store = move |hash: XorName, content: Bytes| -> Result<()> { let _ = storage_clone.lock().unwrap().insert(hash, content.to_vec()); Ok(()) @@ -873,7 +875,7 @@ mod tests { // Encrypt the file let data_map = streaming_encrypt_from_file(&file_path, store)?; - + // Convert stored chunks to EncryptedChunk format let stored = storage.lock().unwrap(); let encrypted_chunks: Vec<_> = stored @@ -886,7 +888,7 @@ mod tests { // Decrypt and verify let decrypted = decrypt(&data_map, &encrypted_chunks)?; assert_eq!(&original_data[..], &decrypted[..]); - + Ok(()) } @@ -921,10 +923,11 @@ mod tests { // Create a very large data map (12 chunks) let original_map = create_dummy_data_map(100000); - let (shrunk_map, _shrink_chunks) = shrink_data_map(original_map.clone(), |hash, content| { - let _ = storage_clone.lock().unwrap().insert(hash, content); - Ok(()) - })?; + let (shrunk_map, _shrink_chunks) = + shrink_data_map(original_map.clone(), |hash, content| { + let _ = storage_clone.lock().unwrap().insert(hash, content); + Ok(()) + })?; // Verify multiple levels of shrinking occurred assert!(shrunk_map.child().unwrap() > 0); @@ -942,22 +945,28 @@ mod tests { fn test_encryption_algorithm_consistency() -> Result<()> { // Create deterministic test data let test_data = vec![42u8; MIN_ENCRYPTABLE_BYTES * 2]; // Repeating value for predictability - + // First encryption let storage1 = Arc::new(Mutex::new(HashMap::new())); let storage1_clone = storage1.clone(); - + let store1 = move |hash: XorName, content: Bytes| -> Result<()> { - let _ = storage1_clone.lock().unwrap().insert(hash, content.to_vec()); + let _ = storage1_clone + .lock() + .unwrap() + .insert(hash, content.to_vec()); Ok(()) }; // Second encryption let storage2 = Arc::new(Mutex::new(HashMap::new())); let storage2_clone = storage2.clone(); - + let store2 = move |hash: XorName, content: Bytes| -> Result<()> { - let _ = storage2_clone.lock().unwrap().insert(hash, content.to_vec()); + let _ = storage2_clone + .lock() + .unwrap() + .insert(hash, content.to_vec()); Ok(()) }; @@ -965,7 +974,7 @@ mod tests { let temp_dir = tempfile::TempDir::new()?; let file_path1 = temp_dir.path().join("test1.bin"); let file_path2 = temp_dir.path().join("test2.bin"); - + std::fs::write(&file_path1, &test_data)?; std::fs::write(&file_path2, &test_data)?; @@ -983,7 +992,7 @@ mod tests { // Compare stored chunks let stored1 = storage1.lock().unwrap(); let stored2 = storage2.lock().unwrap(); - + assert_eq!( stored1.len(), stored2.len(), @@ -992,10 +1001,11 @@ mod tests { // Compare each chunk's content for (hash, content1) in stored1.iter() { - let content2 = stored2.get(hash).expect("Chunk should exist in both storages"); + let content2 = stored2 + .get(hash) + .expect("Chunk should exist in both storages"); assert_eq!( - content1, - content2, + content1, content2, "Encrypted chunks should be identical for same input" ); }