[feature] incremental decrypter

2 files changed, 46 insertions, 15 deletions

diff --git a/client/src/leap/soledad/client/_crypto.py b/client/src/leap/soledad/client/_crypto.py
index 923891fc..66b26521 100644
--- a/client/src/leap/soledad/client/_crypto.py
+++ b/client/src/leap/soledad/client/_crypto.py
@@ -328,15 +328,16 @@ class BlobDecryptor(object):
         self.armor = armor
         self._producer = None
         self.result = result or BytesIO()
-        self.sym_key = _get_sym_key_for_doc(doc_info.doc_id, secret)
+        sym_key = _get_sym_key_for_doc(doc_info.doc_id, secret)
         self.size = None
-        self.tag = tag
+        self.tag = None
+
         preamble, iv = self._consume_preamble()
         assert preamble
         assert iv
-        self._aes = AESWriter(self.sym_key, iv, self.result, tag=self.tag)
-        self._aes.authenticate(preamble)
 
+        self._aes = AESWriter(sym_key, iv, self.result, tag=tag or self.tag)
+        self._aes.authenticate(preamble)
         if start_stream:
             self._start_stream()
 
@@ -346,12 +347,18 @@ class BlobDecryptor(object):
     def _consume_preamble(self):
         self.fd.seek(0)
         try:
-            preamble, ciphertext = self.fd.getvalue().split(SEPARATOR, 1)
-            preamble = base64.urlsafe_b64decode(preamble)
-            if self.armor:
-                ciphertext = base64.urlsafe_b64decode(ciphertext)
-            tag, ciphertext = ciphertext[-16:], ciphertext[:-16]
-            self.tag = self.tag or tag
+            parts = self.fd.getvalue().split(SEPARATOR, 1)
+            preamble = base64.urlsafe_b64decode(parts[0])
+            if len(parts) == 2:
+                ciphertext = parts[1]
+                if self.armor:
+                    ciphertext = base64.urlsafe_b64decode(ciphertext)
+                self.tag, ciphertext = ciphertext[-16:], ciphertext[:-16]
+                self.fd.seek(0)
+                self.fd.write(ciphertext)
+                self.fd.seek(len(ciphertext))
+                self.fd.truncate()
+                self.fd.seek(0)
 
         except (TypeError, ValueError):
             raise InvalidBlob
@@ -387,11 +394,6 @@ class BlobDecryptor(object):
         if doc_id != self.doc_id:
             raise InvalidBlob('invalid doc id')
 
-        self.fd.seek(0)
-        self.fd.write(ciphertext)
-        self.fd.seek(len(ciphertext))
-        self.fd.truncate()
-        self.fd.seek(0)
         return preamble, iv
 
     def _end_stream(self):
diff --git a/testing/tests/client/test_crypto.py b/testing/tests/client/test_crypto.py
index 44e1b2a5..11384ad7 100644
--- a/testing/tests/client/test_crypto.py
+++ b/testing/tests/client/test_crypto.py
@@ -143,6 +143,35 @@ class BlobTestCase(unittest.TestCase):
         assert str(decrypted) == snowden1
 
     @defer.inlineCallbacks
+    def test_init_with_preamble_alone(self):
+        ciphertext = yield self.blob.encrypt()
+        preamble = ciphertext.getvalue().split()[0]
+        decryptor = _crypto.BlobDecryptor(
+            self.doc_info, BytesIO(preamble),
+            start_stream=False,
+            secret='A' * 96)
+        assert decryptor._consume_preamble()
+
+    @defer.inlineCallbacks
+    def test_incremental_blob_decryptor(self):
+        ciphertext = yield self.blob.encrypt()
+        preamble, ciphertext = ciphertext.getvalue().split()
+        ciphertext = base64.urlsafe_b64decode(ciphertext)
+
+        decryptor = _crypto.BlobDecryptor(
+            self.doc_info, BytesIO(preamble),
+            start_stream=False,
+            secret='A' * 96,
+            tag=ciphertext[-16:])
+        ciphertext = BytesIO(ciphertext[:-16])
+        chunk = ciphertext.read(10)
+        while chunk:
+            decryptor.write(chunk)
+            chunk = ciphertext.read(10)
+        decrypted = decryptor._end_stream()
+        assert decrypted.getvalue() == snowden1
+
+    @defer.inlineCallbacks
     def test_blob_decryptor(self):
         ciphertext = yield self.blob.encrypt()