xep_0454: use streaming API from CipherContext

Signed-off-by: Maxime “pep” Buquet <pep@bouah.net>

1 files changed, 12 insertions, 8 deletions

diff --git a/slixmpp/plugins/xep_0454/__init__.py b/slixmpp/plugins/xep_0454/__init__.py
index 2d3cb8f6..9e9c3b02 100644
--- a/slixmpp/plugins/xep_0454/__init__.py
+++ b/slixmpp/plugins/xep_0454/__init__.py
@@ -57,14 +57,18 @@ class XEP_0454(BasePlugin):
             modes.GCM(aes_gcm_iv),
         ).encryptor()
 
-        # TODO: use streaming API from CipherContext
-        if input_file:
-            plain = input_file.read()
-        else:
-            with filename.open(mode='rb') as file:
-                plain = file.read()
-
-        payload = aes_gcm.update(plain) + aes_gcm.tag + aes_gcm.finalize()
+        if input_file is None:
+            input_file = open(filename, 'rb')
+
+        payload = b''
+        while True:
+            buf = input_file.read(4096)
+            if not buf:
+                break
+            payload += aes_gcm.update(buf)
+
+        aes_gcm.finalize()
+        payload += aes_gcm.tag
         fragment = aes_gcm_iv.hex() + aes_gcm_key.hex()
         return (payload, fragment)