From 6e591b43d60cc6e1552cf2ef8cce6ac6c59a0265 Mon Sep 17 00:00:00 2001
From: Jan-Michael Brummer <jan.brummer@tabos.org>
Date: Tue, 9 Apr 2024 21:58:09 +0200
Subject: [PATCH] Fix composite key computation for BytesIO (#388)

In case a keyfile as BytesIO has been read before, the
next read will be empty. We need to ensure that we
are reading the data from the beginning.

Add seek to start to fix it.

Co-authored-by: Jan-Michael Brummer <jan-michael.brummer1@volkswagen.de>
---
 pykeepass/kdbx_parsing/common.py |  6 ++++--
 tests/tests.py                   | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/pykeepass/kdbx_parsing/common.py b/pykeepass/kdbx_parsing/common.py
index 5542a0bb..ac42d8c0 100644
--- a/pykeepass/kdbx_parsing/common.py
+++ b/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)
diff --git a/tests/tests.py b/tests/tests.py
index 75842192..5bc02b6d 100644
--- a/tests/tests.py
+++ b/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'