[bug] fix loose ends on _crypto and tests

3 files changed, 18 insertions, 58 deletions

diff --git a/client/src/leap/soledad/client/_crypto.py b/client/src/leap/soledad/client/_crypto.py
index c152502d..8dbdc094 100644
--- a/client/src/leap/soledad/client/_crypto.py
+++ b/client/src/leap/soledad/client/_crypto.py
@@ -307,48 +307,6 @@ class BlobEncryptor(object):
         return defer.succeed(result)
 
 
-class CryptoStreamBodyProducer(FileBodyProducer):
-
-    """
-    A BodyProducer that gets the tag from the last 16 bytes before closing the
-    fd.
-    """
-    _tag = None
-
-    @property
-    def tag(self):
-        # XXX this is a bit tricky. If you call this
-        # before the end of the stream, you will ruin everything
-        if not self._tag:
-            self._writeTag()
-        return self._tag
-
-    def _writeTag(self):
-        fd = self._inputFile
-        fd.seek(-16, os.SEEK_END)
-        self._tag = fd.read(16)
-        fd.seek(0)
-
-    def stopProducing(self):
-        self._writeTag()
-        self._inputFile.close()
-        self._task.stop()
-
-    def _writeloop(self, consumer):
-        """
-        Return an iterator which reads one chunk of bytes from the input file
-        and writes them to the consumer for each time it is iterated.
-        """
-        while True:
-            bytes = self._inputFile.read(self._readSize)
-            if not bytes:
-                self._writeTag()
-                self._inputFile.close()
-                break
-            consumer.write(base64.urlsafe_b64decode(bytes))
-            yield None
-
-
 # TODO maybe rename this to just Decryptor, since it will be used by blobs
 # and non blobs in soledad.
 class BlobDecryptor(object):
@@ -381,14 +339,13 @@ class BlobDecryptor(object):
         self.result = result or BytesIO()
         sym_key = _get_sym_key_for_doc(doc_info.doc_id, secret)
 
-        self._aes = AESWriter(sym_key, iv, self.result, tag=None)
+        self._aes = AESWriter(sym_key, iv, self.result, tag=self.tag)
         self._aes.authenticate(preamble)
         if start_stream:
             self._start_stream()
 
     def _start_stream(self):
-        self._producer = CryptoStreamBodyProducer(self.fd, readSize=2**16)
-        self._producer.armor = self.armor
+        self._producer = FileBodyProducer(self.fd, readSize=2**16)
 
     def _consume_preamble(self):
 
@@ -397,6 +354,8 @@ class BlobDecryptor(object):
             parts = self.fd.getvalue().split()
             encoded_preamble = parts[0]
             preamble = base64.urlsafe_b64decode(encoded_preamble)
+            ciphertext = base64.urlsafe_b64decode(parts[1])
+            self.tag, ciphertext = ciphertext[-16:], ciphertext[:-16]
 
         except (TypeError, ValueError):
             raise InvalidBlob
@@ -433,21 +392,18 @@ class BlobDecryptor(object):
             raise InvalidBlob('invalid doc id')
 
         self.fd.seek(0)
-        tail = ''.join(parts[1:])
-        self.fd.write(tail)
-        self.fd.seek(len(tail))
+        self.fd.write(ciphertext)
+        self.fd.seek(len(ciphertext))
         self.fd.truncate()
         self.fd.seek(0)
         return preamble, iv
 
     def _end_stream(self):
         try:
-            self._aes.end()[1]
+            self._aes.end()
         except InvalidTag:
             raise InvalidBlob('Invalid Tag. Blob authentication failed.')
         fd = self.result
-        fd.seek(-16, os.SEEK_END)
-        fd.truncate()
         fd.seek(0)
         return self.result
 
diff --git a/client/src/leap/soledad/client/http_target/fetch.py b/client/src/leap/soledad/client/http_target/fetch.py
index 8676ceed..cf4984d1 100644
--- a/client/src/leap/soledad/client/http_target/fetch.py
+++ b/client/src/leap/soledad/client/http_target/fetch.py
@@ -115,7 +115,7 @@ class HTTPDocFetcher(object):
     def __atomic_doc_parse(self, doc_info, content, total):
         doc = SoledadDocument(doc_info['id'], doc_info['rev'], content)
         if is_symmetrically_encrypted(content):
-            content = yield self._crypto.decrypt_doc(doc)
+            content = (yield self._crypto.decrypt_doc(doc)).getvalue()
         elif old_crypto.is_symmetrically_encrypted(doc):
             content = self._deprecated_crypto.decrypt_doc(doc)
         doc.set_json(content)
diff --git a/testing/tests/client/test_crypto.py b/testing/tests/client/test_crypto.py
index 399fdc99..e268428d 100644
--- a/testing/tests/client/test_crypto.py
+++ b/testing/tests/client/test_crypto.py
@@ -105,7 +105,9 @@ class BlobTestCase(unittest.TestCase):
             secret='A' * 96)
 
         encrypted = yield blob.encrypt()
-        preamble, ciphertext = _crypto._split(encrypted.getvalue())
+        preamble, ciphertext = encrypted.getvalue().split()
+        preamble = base64.urlsafe_b64decode(preamble)
+        ciphertext = base64.urlsafe_b64decode(ciphertext)
         ciphertext = ciphertext[:-16]
 
         assert len(preamble) == _crypto.PACMAN.size
@@ -140,7 +142,7 @@ class BlobTestCase(unittest.TestCase):
             self.doc_info, ciphertext,
             secret='A' * 96)
         decrypted = yield decryptor.decrypt()
-        assert decrypted == snowden1
+        assert decrypted.getvalue() == snowden1
 
     @defer.inlineCallbacks
     def test_encrypt_and_decrypt(self):
@@ -157,7 +159,7 @@ class BlobTestCase(unittest.TestCase):
         doc2 = SoledadDocument('id1', '1')
         doc2.set_json(encrypted)
         assert _crypto.is_symmetrically_encrypted(encrypted)
-        decrypted = yield crypto.decrypt_doc(doc2)
+        decrypted = (yield crypto.decrypt_doc(doc2)).getvalue()
         assert len(decrypted) != 0
         assert json.loads(decrypted) == payload
 
@@ -172,7 +174,9 @@ class BlobTestCase(unittest.TestCase):
 
         encrypted = yield crypto.encrypt_doc(doc1)
         encdict = json.loads(encrypted)
-        preamble, raw = _crypto._split(str(encdict['raw']))
+        preamble, raw = str(encdict['raw']).split()
+        preamble = base64.urlsafe_b64decode(preamble)
+        raw = base64.urlsafe_b64decode(raw)
         # mess with tag
         messed = raw[:-16] + '0' * 16
 
@@ -269,7 +273,7 @@ class PreambleTestCase(unittest.TestCase):
         preamble = self.blob._encode_preamble()
         unpacked = _crypto.PACMAN.unpack(preamble)
         size = unpacked[7]
-        assert size == len(snowden1)
+        assert size == _crypto._ceiling(len(snowden1))
 
     @defer.inlineCallbacks
     def test_preamble_can_come_without_size(self):
@@ -291,7 +295,7 @@ class PreambleTestCase(unittest.TestCase):
         cleartext = yield _crypto.BlobDecryptor(
             self.doc_info, BytesIO(ciphertext),
             secret='A' * 96).decrypt()
-        assert cleartext == self.cleartext.getvalue()
+        assert cleartext.getvalue() == self.cleartext.getvalue()
         warnings = self.flushWarnings()
         assert len(warnings) == 1
         assert 'legacy document without size' in warnings[0]['message']