1 files changed, 36 insertions, 10 deletions

diff --git a/src/zipcrypto.rs b/src/zipcrypto.rs
index 91d4039..c3696e4 100644
--- a/src/zipcrypto.rs
+++ b/src/zipcrypto.rs
@@ -6,7 +6,8 @@
 use std::num::Wrapping;
 
 /// A container to hold the current key state
-struct ZipCryptoKeys {
+#[derive(Clone, Copy)]
+pub(crate) struct ZipCryptoKeys {
     key_0: Wrapping<u32>,
     key_1: Wrapping<u32>,
     key_2: Wrapping<u32>,
@@ -49,6 +50,13 @@ impl ZipCryptoKeys {
     fn crc32(crc: Wrapping<u32>, input: u8) -> Wrapping<u32> {
         (crc >> 8) ^ Wrapping(CRCTABLE[((crc & Wrapping(0xff)).0 as u8 ^ input) as usize])
     }
+    pub(crate) fn derive(password: &[u8]) -> ZipCryptoKeys {
+        let mut keys = ZipCryptoKeys::new();
+        for byte in password.iter() {
+            keys.update(*byte);
+        }
+        keys
+    }
 }
 
 /// A ZipCrypto reader with unverified password
@@ -70,17 +78,10 @@ impl<R: std::io::Read> ZipCryptoReader<R> {
     /// would be impossible to decrypt files that were encrypted with a
     /// password byte sequence that is unrepresentable in UTF-8.
     pub fn new(file: R, password: &[u8]) -> ZipCryptoReader<R> {
-        let mut result = ZipCryptoReader {
+        ZipCryptoReader {
             file,
-            keys: ZipCryptoKeys::new(),
-        };
-
-        // Key the cipher by updating the keys with the password.
-        for byte in password.iter() {
-            result.keys.update(*byte);
+            keys: ZipCryptoKeys::derive(password),
         }
-
-        result
     }
 
     /// Read the ZipCrypto header bytes and validate the password.
@@ -122,6 +123,31 @@ impl<R: std::io::Read> ZipCryptoReader<R> {
         Ok(Some(ZipCryptoReaderValid { reader: self }))
     }
 }
+pub(crate) struct ZipCryptoWriter<W> {
+    pub(crate) writer: W,
+    pub(crate) buffer: Vec<u8>,
+    pub(crate) keys: ZipCryptoKeys,
+}
+impl<W: std::io::Write> ZipCryptoWriter<W> {
+    pub(crate) fn finish(mut self, crc32: u32) -> std::io::Result<W> {
+        self.buffer[11] = (crc32 >> 24) as u8;
+        for byte in self.buffer.iter_mut() {
+            *byte = self.keys.encrypt_byte(*byte);
+        }
+        self.writer.write_all(&self.buffer)?;
+        self.writer.flush()?;
+        Ok(self.writer)
+    }
+}
+impl<W: std::io::Write> std::io::Write for ZipCryptoWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.buffer.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
 
 /// A ZipCrypto reader with verified password
 pub struct ZipCryptoReaderValid<R> {