From 75fe4a15f6c81dab0773db0bbb2de6800be63d6e Mon Sep 17 00:00:00 2001
From: Jan-Michael Brummer <jan-michael.brummer1@volkswagen.de>
Date: Mon, 23 Oct 2023 15:03:18 +0200
Subject: [PATCH] Allow keyfile to be passed as bytes

In order to prevent application to create temporary files,
provide keyfile as bytes.

Fixes: https://github.com/libkeepass/pykeepass/issues/363
---
 pykeepass/kdbx_parsing/common.py | 81 ++++++++++++++++++--------------
 tests/tests.py                   | 23 +++++++++
 2 files changed, 68 insertions(+), 36 deletions(-)

diff --git a/pykeepass/kdbx_parsing/common.py b/pykeepass/kdbx_parsing/common.py
index 5b547366..13d2c31f 100644
--- a/pykeepass/kdbx_parsing/common.py
+++ b/pykeepass/kdbx_parsing/common.py
@@ -105,6 +105,25 @@ def aes_kdf(key, rounds, key_composite):
     return hashlib.sha256(transformed_key).digest()
 
 
+def load_keyfile_composite(key: bytes):
+    try:
+        int(key, 16)
+        is_hex = True
+    except ValueError:
+        is_hex = False
+    # if the length is 32 bytes we assume it is the key
+    if len(key) == 32:
+        keyfile_composite = key
+    # if the length is 64 bytes we assume the key is hex encoded
+    elif len(key) == 64 and is_hex:
+        keyfile_composite = codecs.decode(key, 'hex')
+    # anything else may be a file to hash for the key
+    else:
+        keyfile_composite = hashlib.sha256(key).digest()
+
+    return keyfile_composite
+
+
 def compute_key_composite(password=None, keyfile=None):
     """Compute composite key.
     Used in header verification and payload decryption."""
@@ -116,43 +135,33 @@ def compute_key_composite(password=None, keyfile=None):
         password_composite = b''
     # hash the keyfile
     if keyfile:
-        # try to read XML keyfile
-        try:
-            with open(keyfile, 'r') as f:
-                tree = etree.parse(f).getroot()
-                version = tree.find('Meta/Version').text
-                data_element = tree.find('Key/Data')
-                if version.startswith('1.0'):
-                    keyfile_composite = base64.b64decode(data_element.text)
-                elif version.startswith('2.0'):
-                    # read keyfile data and convert to bytes
-                    keyfile_composite = bytes.fromhex(data_element.text.strip())
-                    # validate bytes against hash
-                    hash = bytes.fromhex(data_element.attrib['Hash'])
-                    hash_computed = hashlib.sha256(keyfile_composite).digest()[:4]
-                    assert hash == hash_computed, "Keyfile has invalid hash"
-        # otherwise, try to read plain keyfile
-        except (etree.XMLSyntaxError, UnicodeDecodeError):
+        if isinstance(keyfile, bytes):
+            keyfile_composite = keyfile_composite = load_keyfile_composite(keyfile)
+        else:
+            # try to read XML keyfile
             try:
-                with open(keyfile, 'rb') as f:
-                    key = f.read()
-
-                    try:
-                        int(key, 16)
-                        is_hex = True
-                    except ValueError:
-                        is_hex = False
-                    # if the length is 32 bytes we assume it is the key
-                    if len(key) == 32:
-                        keyfile_composite = key
-                    # if the length is 64 bytes we assume the key is hex encoded
-                    elif len(key) == 64 and is_hex:
-                        keyfile_composite = codecs.decode(key, 'hex')
-                    # anything else may be a file to hash for the key
-                    else:
-                        keyfile_composite = hashlib.sha256(key).digest()
-            except:
-                raise IOError('Could not read keyfile')
+                with open(keyfile, 'r') as f:
+                    tree = etree.parse(f).getroot()
+                    version = tree.find('Meta/Version').text
+                    data_element = tree.find('Key/Data')
+                    if version.startswith('1.0'):
+                        keyfile_composite = base64.b64decode(data_element.text)
+                    elif version.startswith('2.0'):
+                        # read keyfile data and convert to bytes
+                        keyfile_composite = bytes.fromhex(data_element.text.strip())
+                        # validate bytes against hash
+                        hash = bytes.fromhex(data_element.attrib['Hash'])
+                        hash_computed = hashlib.sha256(keyfile_composite).digest()[:4]
+                        assert hash == hash_computed, "Keyfile has invalid hash"
+            # otherwise, try to read plain keyfile
+            except (etree.XMLSyntaxError, UnicodeDecodeError):
+                try:
+                    with open(keyfile, 'rb') as f:
+                        key = f.read()
+
+                        keyfile_composite = load_keyfile_composite(key)
+                except:
+                    raise IOError('Could not read keyfile')
 
     else:
         keyfile_composite = b''
diff --git a/tests/tests.py b/tests/tests.py
index 348eeb38..1e1e3992 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -1230,6 +1230,29 @@ def test_open_no_decrypt(self):
 
             self.assertEqual(kp.database_salt, salt)
 
+
+    def test_keyfile_as_bytes(self):
+
+        databases = [
+            'test4.kdbx',
+        ]
+        passwords = [
+            'password',
+        ]
+        keyfiles = [
+            base_dir + '/test4.key'
+        ]
+        for database, password, keyfile in zip(databases, passwords, keyfiles):
+            with open(keyfile, "rb") as fh:
+                buf = fh.read()
+
+            kp = PyKeePass(
+                os.path.join(base_dir, database),
+                password,
+                keyfile=buf
+            )
+
+
 if __name__ == '__main__':
     unittest.main()