Fix composite key computation for BytesIO

pykeepass/kdbx_parsing/common.py

@@ -1,13 +1,13 @@
 import base64
 import codecs
 import hashlib
+import io
 import logging
 import re
 import zlib
 from binascii import Error as BinasciiError
 from collections import OrderedDict
 from copy import deepcopy
-from io import BytesIO
 
 from construct import (
     Adapter,
@@ -128,6 +128,8 @@ def compute_key_composite(password=None, keyfile=None):
     # hash the keyfile
     if keyfile:
         if hasattr(keyfile, "read"):
+            if hasattr(keyfile, "seekable") and keyfile.seekable():
+                keyfile.seek(0)
             keyfile_bytes = keyfile.read()
         else:
             with open(keyfile, 'rb') as f:
@@ -194,7 +196,7 @@ class XML(Adapter):
 
     def _decode(self, data, con, path):
         parser = etree.XMLParser(remove_blank_text=True)
-        return etree.parse(BytesIO(data), parser)
+        return etree.parse(io.BytesIO(data), parser)
 
     def _encode(self, tree, con, path):
         return etree.tostring(tree)

tests/tests.py

@@ -875,6 +875,20 @@ def test_fields(self):
 
 
 class PyKeePassTests3(KDBX3Tests):
+    def test_consecutives_saves_with_stream(self):
+        # https://github.com/libkeepass/pykeepass/pull/388
+        self.setUp()
+
+        with open(base_dir / self.keyfile_tmp, 'rb') as f:
+            keyfile = BytesIO(f.read())
+
+        for _i in range(5):
+            with PyKeePass(
+                base_dir / self.database_tmp,
+                password=self.password,
+                keyfile=keyfile,
+            ) as kp:
+                kp.save()
 
     def test_set_credentials(self):
         self.kp_tmp.password = 'f00bar'