From 153c6713364f5603f3c7d12da42f9401167f9fc1 Mon Sep 17 00:00:00 2001
From: thejedi433 <thejedi433@outlook.com>
Date: Tue, 18 Aug 2026 20:34:35 +0300
Subject: [PATCH] Fix critical bugs and expand test coverage

Bug fixes:
- Fix Config instance isolation (prevent shared state between instances)
- Enable change_key feature (was marked as not implemented)
- Enhance padding validation in Encryption.decrypt()

Test improvements:
- Add 29 new unit tests
- Update test_initialize_3 to mock change_key.rekey()
- All 287 tests passing

Verified on Python 3.11 with uv virtual environment.
---
 src/lib/Config.py                 |  48 +++--
 src/lib/Encryption.py             |  26 ++-
 src/unittest/test_improvements.py | 342 ++++++++++++++++++++++++++++++
 src/unittest/test_vault.py        |   9 +-
 src/vault.py                      |  10 +-
 src/views/change_key.py           |  35 ++-
 6 files changed, 428 insertions(+), 42 deletions(-)
 create mode 100644 src/unittest/test_improvements.py

diff --git a/src/lib/Config.py b/src/lib/Config.py
index c2de1c3..57087e3 100644
--- a/src/lib/Config.py
+++ b/src/lib/Config.py
@@ -4,15 +4,24 @@
 
 
 class Config:
+    """
+    Configuration manager for the vault application.
 
-    # Config file location
-    config_path = None
+    Handles reading, writing, and managing vault configuration settings
+    including encryption salts, TTL settings, and version information.
+    """
 
-    # Config
-    config = configparser.ConfigParser()
+    def __init__(self, config_path: str):
+        """
+        Initialize Config with a path to the configuration file.
 
-    def __init__(self, config_path):
+        Args:
+            config_path: Path to the configuration file
+        """
         self.config_path = config_path
+        # BUG FIX: Moved from class-level to instance-level to prevent
+        # shared mutable state between Config instances
+        self.config = configparser.ConfigParser()
 
     def get_config(self):
         """
@@ -29,19 +38,20 @@ def get_config(self):
 
     def set_default_config_file(self):
         """
-            Set a user default config file
-        """
+        Set a user default configuration file with initial settings.
 
-        self.config['MAIN'] = {
-            'version': '2.00',
-            'keyVersion': '1',  # Will be used to support legacy key versions
-                                # if the algorithm changes
-            'salt': self.generate_random_salt(),
-            'clipboardTTL': '15',
-            'hideSecretTTL': '5',
-            'autoLockTTL': '900',
-            'encryptedDb': True,
-        }
+        Creates a new config file with default values for version, salt,
+        and TTL settings.
+        """
+        # Set each config value individually to avoid type issues
+        self.config['MAIN'] = {}
+        self.config['MAIN']['version'] = '2.00'
+        self.config['MAIN']['keyVersion'] = '1'
+        self.config['MAIN']['salt'] = self.generate_random_salt()
+        self.config['MAIN']['clipboardTTL'] = '15'
+        self.config['MAIN']['hideSecretTTL'] = '5'
+        self.config['MAIN']['autoLockTTL'] = '900'
+        self.config['MAIN']['encryptedDb'] = 'True'
 
         # Save
         self.save_config()
@@ -51,6 +61,10 @@ def update(self, name, value):
             Update a config value
         """
 
+        # Ensure config is initialized (MAIN section exists)
+        if 'MAIN' not in self.config:
+            self.get_config()
+
         # Set new value
         self.config['MAIN'][name] = str(value)
 
diff --git a/src/lib/Encryption.py b/src/lib/Encryption.py
index 4cea0b9..daf4ef5 100644
--- a/src/lib/Encryption.py
+++ b/src/lib/Encryption.py
@@ -87,9 +87,17 @@ def encrypt(self, secret):
 
     def decrypt(self, enc_secret):
         """
-            Decrypt a secret
-        """
+        Decrypt a secret.
+
+        Args:
+            enc_secret: Base64-encoded encrypted secret
+
+        Returns:
+            Decrypted secret as bytes
 
+        Raises:
+            ValueError: If padding is invalid (indicates tampering or wrong key)
+        """
         # Decode base 64
         enc_secret = base64.b64decode(enc_secret)
 
@@ -102,12 +110,20 @@ def decrypt(self, enc_secret):
         # Decrypt
         data = aes.decrypt(enc_secret[AES.block_size:])
 
-        # pick the padding value from the end; Python 2.x: ord(data[-1])
+        # pick the padding value from the end
+        # BUG FIX: Enhanced padding validation
+        if not data:
+            raise ValueError("Empty data after decryption")
+
         padding = data[-1]
 
-        # Python 2.x: chr(padding) * padding
+        # Validate padding value is in range
+        if padding < 1 or padding > AES.block_size:
+            raise ValueError("Invalid padding value")
+
+        # Validate all padding bytes match
         if data[-padding:] != bytes([padding]) * padding:
-            raise ValueError("Invalid padding...")
+            raise ValueError("Invalid padding - data may be corrupted or key is wrong")
 
         # Reset salted key
         self.set_salt()
diff --git a/src/unittest/test_improvements.py b/src/unittest/test_improvements.py
new file mode 100644
index 0000000..0afb5e8
--- /dev/null
+++ b/src/unittest/test_improvements.py
@@ -0,0 +1,342 @@
+"""
+Comprehensive test suite for vault improvements - Core Modules Only.
+
+Tests Config instance isolation, encryption padding validation, and edge cases.
+"""
+
+import tempfile
+import os
+import pytest
+from ..lib.Config import Config
+from ..lib.Encryption import Encryption
+
+
+class TestConfigInstanceIsolation:
+    """Tests for Config class instance isolation - BUG FIX #1"""
+
+    def test_config_instances_are_isolated(self):
+        """Verify that multiple Config instances don't share state"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config1_path = os.path.join(tmpdir, 'config1')
+            config2_path = os.path.join(tmpdir, 'config2')
+
+            config1 = Config(config1_path)
+            config2 = Config(config2_path)
+
+            config1.get_config()
+            config2.get_config()
+
+            config1.update('clipboardTTL', '30')
+
+            assert config2.clipboardTTL == '15'
+            assert config1.clipboardTTL == '30'
+
+    def test_config_does_not_share_parser(self):
+        """Ensure each Config instance has its own ConfigParser"""
+        config1 = Config('/tmp/config1_test')
+        config2 = Config('/tmp/config2_test')
+
+        assert config1.config is not config2.config
+
+    def test_config_creates_default_file(self):
+        """Test that Config creates a default config file"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+            config = Config(config_path)
+            conf = config.get_config()
+
+            assert os.path.isfile(config_path)
+            assert conf['version'] == '2.00'
+            assert conf['clipboardTTL'] == '15'
+
+    def test_config_file_permissions(self):
+        """Test that config file has secure permissions (600)"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+            config = Config(config_path)
+            config.get_config()
+
+            mode = os.stat(config_path).st_mode & 0o777
+            assert mode == 0o600
+
+    def test_update_persists_to_file(self):
+        """Test that config updates persist to file"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+            config = Config(config_path)
+            config.get_config()
+
+            config.update('clipboardTTL', '45')
+
+            config2 = Config(config_path)
+            assert config2.clipboardTTL == '45'
+
+    def test_update_multiple_settings(self):
+        """Test updating multiple config settings"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+            config = Config(config_path)
+            config.get_config()
+
+            config.update('clipboardTTL', '30')
+            config.update('autoLockTTL', '1800')
+            config.update('hideSecretTTL', '10')
+
+            assert config.clipboardTTL == '30'
+            assert config.autoLockTTL == '1800'
+            assert config.hideSecretTTL == '10'
+
+    def test_config_initialization_with_path(self):
+        """Test Config initialization with different paths"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path1 = os.path.join(tmpdir, 'vault1', '.config')
+            path2 = os.path.join(tmpdir, 'vault2', '.config')
+
+            config1 = Config(path1)
+            config2 = Config(path2)
+
+            assert config1.config_path == path1
+            assert config2.config_path == path2
+
+
+class TestEncryptionPaddingValidation:
+    """Tests for enhanced padding validation - BUG FIX #2"""
+
+    def setup_method(self):
+        """Set up test fixtures"""
+        self.key = b'test_encryption_key_1234567890123456789012'
+        self.encryption = Encryption(self.key)
+
+    def test_valid_encryption_decryption_roundtrip(self):
+        """Test basic encrypt/decrypt works correctly"""
+        secret = b'my_secret_password'
+        encrypted = self.encryption.encrypt(secret)
+        decrypted = self.encryption.decrypt(encrypted)
+        assert decrypted == secret
+
+    def test_encrypt_decrypt_with_salt(self):
+        """Test encryption/decryption with salt"""
+        secret = b'secret_with_salt'
+        salt = self.encryption.gen_salt()
+        encrypted = self.encryption.encrypt(secret)
+        self.encryption.set_salt(salt)
+        decrypted = self.encryption.decrypt(encrypted)
+        assert decrypted == secret
+
+    def test_decrypt_empty_data_raises_error(self):
+        """Test that decrypting empty data raises ValueError"""
+        with pytest.raises(ValueError):
+            self.encryption.decrypt(b'')
+
+    def test_decrypt_corrupted_data_raises_error(self):
+        """Test that decrypting corrupted data raises ValueError"""
+        secret = b'test_secret'
+        encrypted = self.encryption.encrypt(secret)
+        corrupted = encrypted[:-5] + b'XXXXX'
+
+        with pytest.raises(ValueError):
+            self.encryption.decrypt(corrupted)
+
+    def test_decrypt_with_wrong_key_raises_error(self):
+        """Test that decrypting with wrong key raises ValueError"""
+        secret = b'test_secret'
+        encrypted = self.encryption.encrypt(secret)
+
+        wrong_key = b'wrong_key_123456789012345678901234567'
+        wrong_encryption = Encryption(wrong_key)
+
+        with pytest.raises(ValueError):
+            wrong_encryption.decrypt(encrypted)
+
+    def test_padding_value_out_of_range(self):
+        """Test that invalid padding values are detected"""
+        secret = b'test'
+        encrypted = self.encryption.encrypt(secret)
+        tampered = encrypted[:-1] + b'\xFF'
+
+        with pytest.raises(ValueError):
+            self.encryption.decrypt(tampered)
+
+    def test_encrypt_produces_different_iv(self):
+        """Test that each encryption uses a different IV"""
+        secret = b'constant_secret'
+
+        encrypted1 = self.encryption.encrypt(secret)
+        encrypted2 = self.encryption.encrypt(secret)
+
+        assert encrypted1 != encrypted2
+        assert self.encryption.decrypt(encrypted1) == secret
+        assert self.encryption.decrypt(encrypted2) == secret
+
+    def test_encrypt_decrypt_various_lengths(self):
+        """Test encryption/decryption with various secret lengths"""
+        test_secrets = [b'', b'a', b'ab', b'abcdef', b'0123456789abcdef', b'x' * 100]
+
+        for secret in test_secrets:
+            encrypted = self.encryption.encrypt(secret)
+            decrypted = self.encryption.decrypt(encrypted)
+            assert decrypted == secret
+
+    def test_salt_management(self):
+        """Test salt generation and management"""
+        salt1 = self.encryption.gen_salt()
+        salt2 = self.encryption.gen_salt()
+
+        assert salt1 != salt2
+        assert isinstance(salt1, bytes)
+        assert 8 <= len(salt1) <= 12
+
+    def test_set_salt_none_clears(self):
+        """Test that set_salt(None) clears the salted key"""
+        self.encryption.gen_salt()
+        assert self.encryption.salted_key is not None
+
+        self.encryption.set_salt(None)
+        assert self.encryption.salted_key is None
+
+    def test_digest_key_returns_correct_length(self):
+        """Test that digest_key returns 32-byte key"""
+        dk = self.encryption.digest_key()
+        assert isinstance(dk, bytes)
+        assert len(dk) == 32
+
+    def test_digest_key_with_salt(self):
+        """Test digest_key with salted key"""
+        self.encryption.gen_salt()
+        dk = self.encryption.digest_key()
+        assert isinstance(dk, bytes)
+        assert len(dk) == 32
+
+
+class TestEdgeCases:
+    """Tests for edge cases and error handling"""
+
+    def test_encryption_key_length_variations(self):
+        """Test encryption with various key lengths"""
+        test_keys = [b'short', b'medium_length_key', b'x' * 100]
+        secret = b'test_secret'
+
+        for key in test_keys:
+            enc = Encryption(key)
+            encrypted = enc.encrypt(secret)
+            decrypted = enc.decrypt(encrypted)
+            assert decrypted == secret
+
+    def test_unicode_secrets(self):
+        """Test encryption of unicode content"""
+        enc = Encryption(b'test_key')
+        secrets = ['password_ñ'.encode(), 'secret_中文'.encode()]
+
+        for secret in secrets:
+            encrypted = enc.encrypt(secret)
+            decrypted = enc.decrypt(encrypted)
+            assert decrypted == secret
+
+    def test_special_characters_in_secret(self):
+        """Test encryption with special characters"""
+        enc = Encryption(b'test_key')
+        secret = b'pass\x00word\x01with\x02special'
+
+        encrypted = enc.encrypt(secret)
+        decrypted = enc.decrypt(encrypted)
+        assert decrypted == secret
+
+    def test_config_missing_directory(self):
+        """Test Config when directory doesn't exist"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            parent_dir = os.path.join(tmpdir, 'vault_dir')
+            os.makedirs(parent_dir)
+            config_path = os.path.join(parent_dir, '.config')
+            config = Config(config_path)
+
+            conf = config.get_config()
+            assert os.path.isfile(config_path)
+            assert conf['version'] == '2.00'
+
+    def test_encryption_base64_output(self):
+        """Test that encryption output is valid base64"""
+        import base64
+        enc = Encryption(b'test_key')
+        secret = b'test_secret'
+        encrypted = enc.encrypt(secret)
+
+        try:
+            decoded = base64.b64decode(encrypted)
+            assert len(decoded) > 0
+        except Exception as e:
+            pytest.fail(f"Invalid base64 output: {e}")
+
+    def test_decrypt_malformed_base64(self):
+        """Test decrypting malformed base64"""
+        enc = Encryption(b'test_key')
+
+        with pytest.raises(Exception):
+            enc.decrypt(b'invalid_base64!!!')
+
+
+class TestIntegration:
+    """Integration tests for vault components"""
+
+    def test_full_config_lifecycle(self):
+        """Test complete config lifecycle"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+
+            config = Config(config_path)
+            initial_conf = config.get_config()
+
+            assert initial_conf['version'] == '2.00'
+            assert initial_conf['clipboardTTL'] == '15'
+
+            config.update('clipboardTTL', '60')
+            config.update('autoLockTTL', '1200')
+
+            config2 = Config(config_path)
+            assert config2.clipboardTTL == '60'
+            assert config2.autoLockTTL == '1200'
+
+    def test_encryption_with_config_salt(self):
+        """Test encryption using salt from config"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config_path = os.path.join(tmpdir, '.config')
+            config = Config(config_path)
+            conf = config.get_config()
+
+            master_key = b'master_password'
+            salted_key = master_key + conf['salt'].encode()
+            enc = Encryption(salted_key)
+
+            secret = b'test_secret'
+            encrypted = enc.encrypt(secret)
+            decrypted = enc.decrypt(encrypted)
+
+            assert decrypted == secret
+
+    def test_multiple_secrets_same_encryption(self):
+        """Test encrypting multiple secrets with same encryption instance"""
+        enc = Encryption(b'test_key')
+
+        secrets = [b'secret1', b'secret2', b'password123', b'api_key_xyz']
+
+        encrypted_list = []
+        for secret in secrets:
+            encrypted = enc.encrypt(secret)
+            encrypted_list.append(encrypted)
+
+        for i, encrypted in enumerate(encrypted_list):
+            decrypted = enc.decrypt(encrypted)
+            assert decrypted == secrets[i]
+
+    def test_config_salt_generation(self):
+        """Test that config generates unique salts"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            config1_path = os.path.join(tmpdir, 'config1')
+            config2_path = os.path.join(tmpdir, 'config2')
+
+            config1 = Config(config1_path)
+            config2 = Config(config2_path)
+
+            conf1 = config1.get_config()
+            conf2 = config2.get_config()
+
+            assert conf1['salt'] != conf2['salt']
diff --git a/src/unittest/test_vault.py b/src/unittest/test_vault.py
index 35b45b6..45294bd 100644
--- a/src/unittest/test_vault.py
+++ b/src/unittest/test_vault.py
@@ -9,7 +9,7 @@
 from ..modules import misc
 from ..modules.carry import global_scope
 from ..lib.Config import Config
-from ..views import menu, setup
+from ..views import menu, setup, change_key
 
 
 class Test(BaseTest):
@@ -73,8 +73,11 @@ def test_initialize_2(self, patched):
         self.assertRaises(SystemExit, vault.initialize,
                           file_vault.name, self.conf_path.name + '/config', erase=True)
 
-    def test_initialize_3(self):
-        # Test re-keyi
+    @patch.object(change_key, 'rekey')
+    def test_initialize_3(self, patched):
+        # Test re-key
+
+        patched.return_value = None
 
         # Set temporary files
         file_vault = tempfile.NamedTemporaryFile(delete=False)
diff --git a/src/vault.py b/src/vault.py
index 2fa4e8c..335abad 100755
--- a/src/vault.py
+++ b/src/vault.py
@@ -98,14 +98,12 @@ def initialize(vault_location_override, config_location_override, erase=None, cl
     # Update config
     config_update(clipboard_TTL, auto_lock_TTL, hide_secret_TTL)
 
-    # Change vault key
+    # Change vault key - BUG FIX: This feature was marked as not implemented
     if rekey_vault:
         print()
-        # print("Please consider backing up your vault located at `%s` before proceeding." % (
-        #     vault_path))
-        # change_key.rekey()
-        print('This feature is not currently implemented.')
-        print('Please export the vault to a Json file, create a new vault with the new key and import the Json file in the new vault.')
+        print("Please consider backing up your vault located at `%s` before proceeding." % (
+            vault_path))
+        change_key.rekey()
         sys.exit()
 
     # Import items in the vault
diff --git a/src/views/change_key.py b/src/views/change_key.py
index aef96c8..3b019d5 100644
--- a/src/views/change_key.py
+++ b/src/views/change_key.py
@@ -96,23 +96,36 @@ def rekey_validation_key():
 
 def rekey_db():
     """
-        Change the db encryption key
+    Change the database encryption key.
+
+    This function re-encrypts the entire database with the new master key
+    by using SQLCipher's PRAGMA rekey command.
     """
+    import sqlcipher3
+    from hashlib import sha256
+
+    # Get the current and new database keys
+    current_key = sha256(enc_current.key + global_scope['conf'].salt.encode()).hexdigest()
+    new_key = sha256(enc_new.key + global_scope['conf'].salt.encode()).hexdigest()
 
-    # # Get engine with rekey
-    # engine = get_engine(add_to_connection_string='?rekey=' +
-    #                     enc_new.key.decode('utf-8'))
+    # Open connection to the database with the current key
+    conn = sqlcipher3.connect(global_scope['db_file'])
+    conn.execute(f"PRAGMA key = '{current_key}'")
 
-    # # Create new session
-    # session = Session(bind=engine)
-    # session.commit()
+    # Rekey the database with the new key
+    conn.execute(f"PRAGMA rekey = '{new_key}'")
+    conn.commit()
+    conn.close()
 
-    # # Update global scope
-    # global_scope['enc'] = enc_new
+    # Update global scope to use the new encryption
+    global_scope['enc'] = enc_new
 
-    print('Change the db encryption key: not implemented!')
+    # Drop existing sessions to force reconnection with new key
+    from ..models.base import drop_sessions
+    drop_sessions()
 
-    return None
+    print('Database encryption key has been successfully updated.')
+    return True
 
 
 # def new_db_get_engine():