From c19ee21fe75e6e1149a1bd52b2f14c90bef88b8b Mon Sep 17 00:00:00 2001
From: Maximiliano Sandoval R <msandova@gnome.org>
Date: Fri, 3 Nov 2023 23:59:10 +0100
Subject: [PATCH] Change master seed on each save

Fixes: https://github.com/libkeepass/pykeepass/issues/219
---
 pykeepass/kdbx_parsing/common.py | 19 +++++++++++++++----
 pykeepass/kdbx_parsing/kdbx.py   | 29 ++++++++++++++++++++++++++---
 pykeepass/kdbx_parsing/kdbx3.py  | 17 +++++++++--------
 pykeepass/kdbx_parsing/kdbx4.py  | 17 +++++++++--------
 pykeepass/pykeepass.py           | 12 ++++++------
 tests/tests.py                   | 22 ++++++++++++++++++++++
 6 files changed, 87 insertions(+), 29 deletions(-)

diff --git a/pykeepass/kdbx_parsing/common.py b/pykeepass/kdbx_parsing/common.py
index 5b547366..1f014db8 100644
--- a/pykeepass/kdbx_parsing/common.py
+++ b/pykeepass/kdbx_parsing/common.py
@@ -1,9 +1,10 @@
 from Cryptodome.Cipher import AES, ChaCha20, Salsa20
 from .twofish import Twofish
+from Cryptodome.Random import get_random_bytes
 from Cryptodome.Util import Padding as CryptoPadding
 import hashlib
 from construct import (
-    Adapter, BitStruct, BitsSwapped, Container, Flag, Padding, ListContainer, Mapping, GreedyBytes, Int32ul, Switch
+    Adapter, BitStruct, BitsSwapped, Bytes, Container, Flag, Padding, ListContainer, Mapping, GreedyBytes, Int32ul, Switch, stream_write
 )
 from lxml import etree
 from copy import deepcopy
@@ -20,6 +21,16 @@
 log = logging.getLogger(__name__)
 
 
+class RandomBytes(Bytes):
+    """Same as Bytes, but generate random bytes when building"""
+
+    def _build(self, obj, stream, context, path):
+        length = self.length(context) if callable(self.length) else self.length
+        data = get_random_bytes(length)
+        stream_write(stream, data, length, path)
+        return data
+
+
 class HeaderChecksumError(Exception):
     pass
 
@@ -167,7 +178,7 @@ def compute_master(context):
 
     # combine the transformed key with the header master seed to find the master_key
     master_key = hashlib.sha256(
-        context._.header.value.dynamic_header.master_seed.data +
+        context._.header.dynamic_header.master_seed.data +
         context.transformed_key).digest()
     return master_key
 
@@ -296,7 +307,7 @@ class DecryptedPayload(Adapter):
     def _decode(self, payload_data, con, path):
         cipher = self.get_cipher(
             con.master_key,
-            con._.header.value.dynamic_header.encryption_iv.data
+            con._.header.dynamic_header.encryption_iv.data
         )
         payload_data = cipher.decrypt(payload_data)
         # FIXME: Construct ugliness.  Fixes #244.  First 32 bytes of decrypted kdbx3 payload
@@ -316,7 +327,7 @@ def _encode(self, payload_data, con, path):
         payload_data = self.pad(payload_data)
         cipher = self.get_cipher(
             con.master_key,
-            con._.header.value.dynamic_header.encryption_iv.data
+            con._.header.dynamic_header.encryption_iv.data
         )
         payload_data = cipher.encrypt(payload_data)
 
diff --git a/pykeepass/kdbx_parsing/kdbx.py b/pykeepass/kdbx_parsing/kdbx.py
index 3a43f7cf..fd79c0a5 100644
--- a/pykeepass/kdbx_parsing/kdbx.py
+++ b/pykeepass/kdbx_parsing/kdbx.py
@@ -1,15 +1,38 @@
-from construct import Struct, Switch, Bytes, Int16ul, RawCopy, Check, this
+from construct import Struct, Switch, Bytes, Int16ul, RawCopy, Check, this, stream_seek, stream_tell, stream_read, Subconstruct
 from .kdbx3 import DynamicHeader as DynamicHeader3
 from .kdbx3 import Body as Body3
 from .kdbx4 import DynamicHeader as DynamicHeader4
 from .kdbx4 import Body as Body4
 
+
+class Copy(Subconstruct):
+    """Same as RawCopy, but don't create parent container when parsing.
+    Instead store data in ._data attribute of subconstruct, and never rebuild from data
+    """
+
+    def _parse(self, stream, context, path):
+        offset1 = stream_tell(stream, path)
+        obj = self.subcon._parsereport(stream, context, path)
+        offset2 = stream_tell(stream, path)
+        stream_seek(stream, offset1, 0, path)
+        obj._data = stream_read(stream, offset2 - offset1, path)
+        return obj
+
+    def _build(self, obj, stream, context, path):
+        offset1 = stream_tell(stream, path)
+        obj = self.subcon._build(obj, stream, context, path)
+        offset2 = stream_tell(stream, path)
+        stream_seek(stream, offset1, 0, path)
+        obj._data = stream_read(stream, offset2 - offset1, path)
+        return obj
+
+
 # verify file signature
 def check_signature(ctx):
     return ctx.sig1 == b'\x03\xd9\xa2\x9a' and ctx.sig2 == b'\x67\xFB\x4B\xB5'
 
 KDBX = Struct(
-    "header" / RawCopy(
+    "header" / Copy(
         Struct(
             "sig1" / Bytes(4),
             "sig2" / Bytes(4),
@@ -25,7 +48,7 @@ def check_signature(ctx):
         )
     ),
     "body" / Switch(
-        this.header.value.major_version,
+        this.header.major_version,
         {3: Body3,
          4: Body4
          }
diff --git a/pykeepass/kdbx_parsing/kdbx3.py b/pykeepass/kdbx_parsing/kdbx3.py
index 87db0f31..7e441d69 100644
--- a/pykeepass/kdbx_parsing/kdbx3.py
+++ b/pykeepass/kdbx_parsing/kdbx3.py
@@ -10,7 +10,7 @@
 from .common import (
     aes_kdf, AES256Payload, ChaCha20Payload, TwoFishPayload, Concatenated,
     DynamicDict, compute_key_composite, Decompressed, Reparsed,
-    compute_master, CompressionFlags, XML, CipherId, ProtectedStreamId, Unprotect
+    compute_master, CompressionFlags, XML, CipherId, ProtectedStreamId, Unprotect, RandomBytes
 )
 
 
@@ -33,8 +33,8 @@ def compute_transformed(context):
             keyfile=context._._.keyfile
         )
         transformed_key = aes_kdf(
-            context._.header.value.dynamic_header.transform_seed.data,
-            context._.header.value.dynamic_header.transform_rounds.data,
+            context._.header.dynamic_header.transform_seed.data,
+            context._.header.dynamic_header.transform_rounds.data,
             key_composite
         )
 
@@ -67,6 +67,7 @@ def compute_transformed(context):
             {'compression_flags': CompressionFlags,
              'cipher_id': CipherId,
              'transform_rounds': Int64ul,
+             'master_seed': RandomBytes(32),
              'protected_stream_id': ProtectedStreamId
              },
             default=GreedyBytes
@@ -130,16 +131,16 @@ def compute_transformed(context):
         # validate payload decryption
         "cred_check" / Checksum(
             Bytes(32),
-            lambda this: this._._.header.value.dynamic_header.stream_start_bytes.data,
+            lambda this: this._._.header.dynamic_header.stream_start_bytes.data,
             this,
             # exception=CredentialsError
         ),
         "xml" / Unprotect(
-            this._._.header.value.dynamic_header.protected_stream_id.data,
-            this._._.header.value.dynamic_header.protected_stream_key.data,
+            this._._.header.dynamic_header.protected_stream_id.data,
+            this._._.header.dynamic_header.protected_stream_key.data,
             XML(
                 IfThenElse(
-                    this._._.header.value.dynamic_header.compression_flags.data.compression,
+                    this._._.header.dynamic_header.compression_flags.data.compression,
                     Decompressed(Concatenated(PayloadBlocks)),
                     Concatenated(PayloadBlocks)
                 )
@@ -157,7 +158,7 @@ def compute_transformed(context):
     "payload" / If(this._._.decrypt,
         UnpackedPayload(
             Switch(
-                this._.header.value.dynamic_header.cipher_id.data,
+                this._.header.dynamic_header.cipher_id.data,
                 {'aes256': AES256Payload(GreedyBytes),
                  'chacha20': ChaCha20Payload(GreedyBytes),
                  'twofish': TwoFishPayload(GreedyBytes),
diff --git a/pykeepass/kdbx_parsing/kdbx4.py b/pykeepass/kdbx_parsing/kdbx4.py
index 41ce7a89..18e15cdb 100644
--- a/pykeepass/kdbx_parsing/kdbx4.py
+++ b/pykeepass/kdbx_parsing/kdbx4.py
@@ -12,7 +12,7 @@
 )
 from .common import (
     aes_kdf, Concatenated, AES256Payload, ChaCha20Payload, TwoFishPayload,
-    DynamicDict, compute_key_composite, Reparsed, Decompressed,
+    DynamicDict, RandomBytes, compute_key_composite, Reparsed, Decompressed,
     compute_master, CompressionFlags, XML, CipherId, ProtectedStreamId, Unprotect
 )
 
@@ -34,7 +34,7 @@ def compute_transformed(context):
         password=context._._.password,
         keyfile=context._._.keyfile
     )
-    kdf_parameters = context._.header.value.dynamic_header.kdf_parameters.data.dict
+    kdf_parameters = context._.header.dynamic_header.kdf_parameters.data.dict
 
     if context._._.transformed_key is not None:
         transformed_key = context._._.transformed_key
@@ -73,12 +73,12 @@ def compute_header_hmac_hash(context):
         hashlib.sha512(
             b'\xff' * 8 +
             hashlib.sha512(
-                context._.header.value.dynamic_header.master_seed.data +
+                context._.header.dynamic_header.master_seed.data +
                 context.transformed_key +
                 b'\x01'
             ).digest()
         ).digest(),
-        context._.header.data,
+        context._.header._data,
         hashlib.sha256
     ).digest()
 
@@ -140,6 +140,7 @@ def compute_header_hmac_hash(context):
             this.id,
             {'compression_flags': CompressionFlags,
              'kdf_parameters': VariantDictionary,
+             'master_seed': RandomBytes(32),
              'cipher_id': CipherId
              },
             default=GreedyBytes
@@ -165,7 +166,7 @@ def compute_payload_block_hash(this):
         hashlib.sha512(
             struct.pack('<Q', this._index) +
             hashlib.sha512(
-                this._._.header.value.dynamic_header.master_seed.data +
+                this._._.header.dynamic_header.master_seed.data +
                 this._.transformed_key + b'\x01'
             ).digest()
         ).digest(),
@@ -200,7 +201,7 @@ def compute_payload_block_hash(this):
 ))
 
 DecryptedPayload = Switch(
-    this._.header.value.dynamic_header.cipher_id.data,
+    this._.header.dynamic_header.cipher_id.data,
     {'aes256': AES256Payload(EncryptedPayload),
      'chacha20': ChaCha20Payload(EncryptedPayload),
      'twofish': TwoFishPayload(EncryptedPayload)
@@ -256,7 +257,7 @@ def compute_payload_block_hash(this):
     "sha256" / Checksum(
         Bytes(32),
         lambda data: hashlib.sha256(data).digest(),
-        this._.header.data,
+        this._.header._data,
         # exception=HeaderChecksumError,
     ),
     "cred_check" / If(this._._.decrypt,
@@ -270,7 +271,7 @@ def compute_payload_block_hash(this):
     "payload" / If(this._._.decrypt,
         UnpackedPayload(
             IfThenElse(
-                this._.header.value.dynamic_header.compression_flags.data.compression,
+                this._.header.dynamic_header.compression_flags.data.compression,
                 Decompressed(DecryptedPayload),
                 DecryptedPayload
             )
diff --git a/pykeepass/pykeepass.py b/pykeepass/pykeepass.py
index b32c12d5..dbc890fa 100644
--- a/pykeepass/pykeepass.py
+++ b/pykeepass/pykeepass.py
@@ -184,15 +184,15 @@ def version(self):
         """tuple: Length 2 tuple of ints containing major and minor versions.
         Generally (3, 1) or (4, 0)."""
         return (
-            self.kdbx.header.value.major_version,
-            self.kdbx.header.value.minor_version
+            self.kdbx.header.major_version,
+            self.kdbx.header.minor_version
         )
 
     @property
     def encryption_algorithm(self):
         """str: encryption algorithm used by database during decryption.
         Can be one of 'aes256', 'chacha20', or 'twofish'."""
-        return self.kdbx.header.value.dynamic_header.cipher_id.data
+        return self.kdbx.header.dynamic_header.cipher_id.data
 
     @property
     def kdf_algorithm(self):
@@ -201,7 +201,7 @@ def kdf_algorithm(self):
         if self.version == (3, 1):
             return 'aeskdf'
         elif self.version == (4, 0):
-            kdf_parameters = self.kdbx.header.value.dynamic_header.kdf_parameters.data.dict
+            kdf_parameters = self.kdbx.header.dynamic_header.kdf_parameters.data.dict
             if kdf_parameters['$UUID'].value == kdf_uuids['argon2']:
                 return 'argon2'
             elif kdf_parameters['$UUID'].value == kdf_uuids['argon2id']:
@@ -221,9 +221,9 @@ def database_salt(self):
        credentials which are used in extension to current keyfile."""
 
        if self.version == (3, 1):
-            return self.kdbx.header.value.dynamic_header.transform_seed.data
+            return self.kdbx.header.dynamic_header.transform_seed.data
 
-       kdf_parameters = self.kdbx.header.value.dynamic_header.kdf_parameters.data.dict
+       kdf_parameters = self.kdbx.header.dynamic_header.kdf_parameters.data.dict
        return kdf_parameters['S'].value
 
     @property
diff --git a/tests/tests.py b/tests/tests.py
index 348eeb38..b21d7ff5 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -1230,6 +1230,28 @@ def test_open_no_decrypt(self):
 
             self.assertEqual(kp.database_salt, salt)
 
+    def test_master_key_differs(self):
+        databases = [
+            # 'test3.kdbx',
+            'test4.kdbx',
+        ]
+        keyfiles = [
+            # 'test3.key',
+            'test4.key',
+        ]
+        for database, keyfile in zip(databases, keyfiles):
+            path = os.path.join(base_dir, database)
+            keyfile = os.path.join(base_dir, keyfile)
+            kp = PyKeePass(path, password='password', keyfile=keyfile)
+            master_seed = kp.kdbx.header.dynamic_header.master_seed.data
+            stream = BytesIO()
+            kp.save(stream)
+            stream.seek(0)
+            new_kp = PyKeePass(stream, password='password', keyfile=keyfile)
+            new_master_seed = new_kp.kdbx.header.dynamic_header.master_seed.data
+
+            self.assertNotEqual(master_seed, new_master_seed)
+
 if __name__ == '__main__':
     unittest.main()