import copy import os import sys import types import pytest import yaml from unittest.mock import patch from cryptography.fernet import Fernet from talemate.util.encryption import ( ENC_PREFIX, _key_file_path, _keyring_available, decrypt_sensitive_values, decrypt_value, encrypt_sensitive_values, encrypt_value, get_fernet, reset_fernet, ) # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture() def fresh_fernet(tmp_path): """Use a temporary key file for each test, with keyring disabled.""" key_file = tmp_path / "encryption.key" with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=False), ): reset_fernet() yield key_file reset_fernet() @pytest.fixture() def mock_keyring_module(): """A module-like mock keyring backed by a dict.""" store = {} mod = types.ModuleType("keyring") def get_password(service, username): return store.get((service, username)) def set_password(service, username, password): store[(service, username)] = password def delete_password(service, username): if (service, username) not in store: raise Exception("not found") del store[(service, username)] mod.get_password = get_password mod.set_password = set_password mod.delete_password = delete_password mod.store = store # for test assertions return mod @pytest.fixture() def fresh_fernet_with_keyring(tmp_path, mock_keyring_module): """Use a temporary key file AND mock keyring for each test.""" key_file = tmp_path / "encryption.key" with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=True), patch("talemate.util.encryption.keyring", mock_keyring_module), ): reset_fernet() yield key_file, mock_keyring_module reset_fernet() # --------------------------------------------------------------------------- # Round-trip encrypt/decrypt # --------------------------------------------------------------------------- class TestEncryptDecryptRoundTrip: def test_basic(self, fresh_fernet): original = "sk-ant-api03-abc123" encrypted = encrypt_value(original) assert encrypted.startswith(ENC_PREFIX) assert decrypt_value(encrypted) == original def test_unicode(self, fresh_fernet): original = "key-with-unicode-\u00e9\u00e8" encrypted = encrypt_value(original) assert decrypt_value(encrypted) == original def test_long_key(self, fresh_fernet): original = "sk-" + "a" * 500 encrypted = encrypt_value(original) assert decrypt_value(encrypted) == original # --------------------------------------------------------------------------- # Passthrough / edge cases # --------------------------------------------------------------------------- class TestPassthrough: def test_empty_string_encrypt(self, fresh_fernet): assert encrypt_value("") == "" def test_empty_string_decrypt(self, fresh_fernet): assert decrypt_value("") == "" def test_none_decrypt(self, fresh_fernet): assert decrypt_value(None) is None def test_plaintext_passthrough_on_decrypt(self, fresh_fernet): """Legacy plaintext values (no ENC: prefix) return unchanged.""" assert decrypt_value("sk-plain-key-123") == "sk-plain-key-123" def test_already_encrypted_not_double_encrypted(self, fresh_fernet): original = "sk-abc" encrypted = encrypt_value(original) double = encrypt_value(encrypted) assert double == encrypted # --------------------------------------------------------------------------- # Lost / corrupted key handling # --------------------------------------------------------------------------- class TestInvalidTokenHandling: def test_lost_key_returns_none(self, fresh_fernet): """When the key file is replaced, decrypt returns None and regenerates.""" encrypted = encrypt_value("sk-secret") # Simulate key loss by writing a new key fresh_fernet.write_bytes(Fernet.generate_key()) reset_fernet() result = decrypt_value(encrypted) assert result is None def test_corrupted_key_file_handled(self, fresh_fernet): """A corrupted key file triggers regeneration on get_fernet().""" fresh_fernet.write_bytes(b"not-a-valid-fernet-key") reset_fernet() # Should not raise — regenerates the key f = get_fernet() assert f is not None # And encryption works with the new key encrypted = encrypt_value("sk-test") assert encrypted.startswith(ENC_PREFIX) assert decrypt_value(encrypted) == "sk-test" # --------------------------------------------------------------------------- # Dict-walking logic # --------------------------------------------------------------------------- class TestDictWalking: def test_top_level_api_key(self, fresh_fernet): data = {"openai": {"api_key": "sk-123"}} encrypt_sensitive_values(data) assert data["openai"]["api_key"].startswith(ENC_PREFIX) decrypt_sensitive_values(data) assert data["openai"]["api_key"] == "sk-123" def test_client_api_key(self, fresh_fernet): data = { "clients": { "MyClient": { "api_key": "sk-client-key", "name": "MyClient", "type": "openai", } } } encrypt_sensitive_values(data) assert data["clients"]["MyClient"]["api_key"].startswith(ENC_PREFIX) assert data["clients"]["MyClient"]["name"] == "MyClient" def test_override_api_key(self, fresh_fernet): data = { "clients": { "Proxy": { "override_api_key": "sk-proxy-key", "override_base_url": "https://proxy.example.com", } } } encrypt_sensitive_values(data) assert data["clients"]["Proxy"]["override_api_key"].startswith(ENC_PREFIX) assert ( data["clients"]["Proxy"]["override_base_url"] == "https://proxy.example.com" ) def test_agent_action_config_api_key_encrypted(self, fresh_fernet): """AgentActionConfig api_key values (dicts with 'value') are encrypted.""" data = { "agents": { "tts": { "actions": { "openai_compatible": { "config": {"api_key": {"value": "sk-real-secret-key"}} } } } } } encrypt_sensitive_values(data) inner = data["agents"]["tts"]["actions"]["openai_compatible"]["config"][ "api_key" ]["value"] assert inner.startswith(ENC_PREFIX) decrypt_sensitive_values(data) inner = data["agents"]["tts"]["actions"]["openai_compatible"]["config"][ "api_key" ]["value"] assert inner == "sk-real-secret-key" def test_agent_action_config_empty_api_key_not_touched(self, fresh_fernet): """AgentActionConfig api_key with empty value stays empty.""" data = { "agents": { "tts": { "actions": { "openai_compatible": {"config": {"api_key": {"value": ""}}} } } } } original = copy.deepcopy(data) encrypt_sensitive_values(data) assert data == original def test_none_api_key_not_touched(self, fresh_fernet): data = {"openai": {"api_key": None}} encrypt_sensitive_values(data) assert data["openai"]["api_key"] is None def test_empty_string_api_key_passthrough(self, fresh_fernet): data = {"clients": {"c1": {"api_key": ""}}} encrypt_sensitive_values(data) assert data["clients"]["c1"]["api_key"] == "" def test_full_round_trip(self, fresh_fernet): original = { "openai": {"api_key": "sk-openai-123"}, "anthropic": {"api_key": "sk-ant-456"}, "clients": { "c1": { "api_key": "sk-c1", "override_api_key": "sk-override", "name": "c1", }, "c2": { "api_key": None, "name": "c2", }, }, "agents": { "tts": { "actions": { "elevenlabs": { "config": {"api_key": {"value": "elevenlabs.api_key"}} }, "openai_compatible": { "config": {"api_key": {"value": "sk-real-key"}} }, } } }, } data = copy.deepcopy(original) encrypt_sensitive_values(data) # Verify encryption happened assert data["openai"]["api_key"].startswith(ENC_PREFIX) assert data["anthropic"]["api_key"].startswith(ENC_PREFIX) assert data["clients"]["c1"]["api_key"].startswith(ENC_PREFIX) assert data["clients"]["c1"]["override_api_key"].startswith(ENC_PREFIX) # None stays None assert data["clients"]["c2"]["api_key"] is None # Dict-type api_key inner values are encrypted assert data["agents"]["tts"]["actions"]["elevenlabs"]["config"]["api_key"][ "value" ].startswith(ENC_PREFIX) assert data["agents"]["tts"]["actions"]["openai_compatible"]["config"][ "api_key" ]["value"].startswith(ENC_PREFIX) decrypt_sensitive_values(data) assert data["openai"]["api_key"] == "sk-openai-123" assert data["anthropic"]["api_key"] == "sk-ant-456" assert data["clients"]["c1"]["api_key"] == "sk-c1" assert data["clients"]["c1"]["override_api_key"] == "sk-override" assert data["clients"]["c2"]["api_key"] is None assert ( data["agents"]["tts"]["actions"]["elevenlabs"]["config"]["api_key"]["value"] == "elevenlabs.api_key" ) assert ( data["agents"]["tts"]["actions"]["openai_compatible"]["config"]["api_key"][ "value" ] == "sk-real-key" ) # --------------------------------------------------------------------------- # Env variable store (sensitive map sections) # --------------------------------------------------------------------------- class TestEnvVariableStore: def test_all_values_encrypted(self, fresh_fernet): """Every value in the env store is sensitive, regardless of its name.""" data = {"env": {"KIMI_API_KEY": "sk-kimi-123", "SOME_TOKEN": "tok-456"}} encrypt_sensitive_values(data) assert data["env"]["KIMI_API_KEY"].startswith(ENC_PREFIX) assert data["env"]["SOME_TOKEN"].startswith(ENC_PREFIX) decrypt_sensitive_values(data) assert data["env"]["KIMI_API_KEY"] == "sk-kimi-123" assert data["env"]["SOME_TOKEN"] == "tok-456" def test_plaintext_passthrough_on_decrypt(self, fresh_fernet): """Hand-edited plaintext env values load without migration.""" data = {"env": {"MY_KEY": "plain-value"}} decrypt_sensitive_values(data) assert data["env"]["MY_KEY"] == "plain-value" def test_name_colliding_with_sensitive_field_name(self, fresh_fernet): """A var literally named api_key round-trips despite the walk also matching it by field name.""" data = {"env": {"api_key": "sk-collide"}} encrypt_sensitive_values(data) assert data["env"]["api_key"].startswith(ENC_PREFIX) # not double-encrypted assert decrypt_value(data["env"]["api_key"]) == "sk-collide" decrypt_sensitive_values(data) assert data["env"]["api_key"] == "sk-collide" def test_missing_or_empty_section(self, fresh_fernet): encrypt_sensitive_values({}) data = {"env": {}} encrypt_sensitive_values(data) assert data["env"] == {} def test_lost_key_drops_entries(self, fresh_fernet): """Undecryptable env values are dropped instead of surviving as None (the schema requires string values).""" data = {"env": {"MY_KEY": "sk-secret", "PLAIN": "still-here"}} encrypt_sensitive_values(data) # PLAIN got encrypted too; keep a plaintext entry alongside data["env"]["PLAIN"] = "still-here" fresh_fernet.write_bytes(Fernet.generate_key()) reset_fernet() decrypt_sensitive_values(data) assert "MY_KEY" not in data["env"] assert data["env"]["PLAIN"] == "still-here" def test_yaml_round_trip(self, fresh_fernet): data = {"env": {"KIMI_API_KEY": "sk-kimi-secret"}} encrypt_sensitive_values(data) yaml_str = yaml.dump(data) assert "sk-kimi-secret" not in yaml_str loaded = yaml.safe_load(yaml_str) decrypt_sensitive_values(loaded) assert loaded["env"]["KIMI_API_KEY"] == "sk-kimi-secret" # --------------------------------------------------------------------------- # YAML round-trip (simulates save / load cycle) # --------------------------------------------------------------------------- class TestYamlRoundTrip: def test_encrypted_not_in_yaml(self, fresh_fernet): data = {"openai": {"api_key": "sk-test-123"}} encrypt_sensitive_values(data) yaml_str = yaml.dump(data) assert "sk-test-123" not in yaml_str assert ENC_PREFIX in yaml_str def test_load_from_encrypted_yaml(self, fresh_fernet): data = {"openai": {"api_key": "sk-test-123"}} encrypt_sensitive_values(data) yaml_str = yaml.dump(data) loaded = yaml.safe_load(yaml_str) decrypt_sensitive_values(loaded) assert loaded["openai"]["api_key"] == "sk-test-123" def test_load_from_plaintext_yaml(self, fresh_fernet): """Legacy plaintext config files work without any migration step.""" yaml_str = "openai:\n api_key: sk-legacy-key\n" loaded = yaml.safe_load(yaml_str) decrypt_sensitive_values(loaded) assert loaded["openai"]["api_key"] == "sk-legacy-key" # --------------------------------------------------------------------------- # Key file path # --------------------------------------------------------------------------- class TestKeyFilePath: def test_default_path(self): """Default path is TALEMATE_ROOT/secrets/encryption.key.""" with patch.dict(os.environ, {}, clear=True): path = _key_file_path() assert path.name == "encryption.key" assert path.parent.name == "secrets" def test_env_var_override(self): """TALEMATE_ENCRYPTION_KEY_DIR overrides the default location.""" with patch.dict(os.environ, {"TALEMATE_ENCRYPTION_KEY_DIR": "/custom/secrets"}): path = _key_file_path() assert str(path) == "/custom/secrets/encryption.key" # --------------------------------------------------------------------------- # File permissions (Linux only) # --------------------------------------------------------------------------- @pytest.mark.skipif(sys.platform == "win32", reason="Unix permissions only") class TestFilePermissions: def test_key_file_permissions(self, fresh_fernet): get_fernet() # triggers key generation stat = fresh_fernet.stat() assert oct(stat.st_mode & 0o777) == "0o600" # --------------------------------------------------------------------------- # Keyring storage # --------------------------------------------------------------------------- class TestKeyringStorage: def test_new_key_stored_in_keyring(self, fresh_fernet_with_keyring): """When no key exists, a new key is generated and stored in keyring.""" key_file, mock_kr = fresh_fernet_with_keyring get_fernet() assert ("talemate", "encryption_key") in mock_kr.store assert not key_file.exists() def test_encrypt_decrypt_with_keyring(self, fresh_fernet_with_keyring): """Full round-trip works when key is in keyring only.""" key_file, mock_kr = fresh_fernet_with_keyring encrypted = encrypt_value("sk-secret-123") assert encrypted.startswith(ENC_PREFIX) assert decrypt_value(encrypted) == "sk-secret-123" assert not key_file.exists() def test_keyring_read_on_subsequent_load(self, fresh_fernet_with_keyring): """After storing in keyring, resetting and re-loading reads from keyring.""" key_file, mock_kr = fresh_fernet_with_keyring encrypted = encrypt_value("sk-abc") reset_fernet() assert decrypt_value(encrypted) == "sk-abc" assert not key_file.exists() # --------------------------------------------------------------------------- # File-to-keyring migration # --------------------------------------------------------------------------- class TestFileMigrationToKeyring: def test_existing_file_migrated_to_keyring(self, fresh_fernet_with_keyring): """When a key file exists and keyring is available, key migrates.""" key_file, mock_kr = fresh_fernet_with_keyring # Pre-create a key file (simulating pre-keyring installation) key = Fernet.generate_key() key_file.write_bytes(key) get_fernet() assert ("talemate", "encryption_key") in mock_kr.store assert mock_kr.store[("talemate", "encryption_key")] == key.decode("utf-8") assert not key_file.exists() def test_migration_preserves_encrypted_data(self, fresh_fernet_with_keyring): """Data encrypted with the file-based key can still be decrypted after migration.""" key_file, mock_kr = fresh_fernet_with_keyring # Create a key file and encrypt some data using it key = Fernet.generate_key() key_file.write_bytes(key) # First load: triggers migration encrypted = encrypt_value("sk-migrate-me") reset_fernet() # Second load: key now comes from keyring assert not key_file.exists() assert decrypt_value(encrypted) == "sk-migrate-me" # --------------------------------------------------------------------------- # Keyring fallback to file # --------------------------------------------------------------------------- class TestKeyringFallback: def test_keyring_store_fails_falls_back_to_file(self, tmp_path): """When keyring.set_password raises, key is written to file.""" key_file = tmp_path / "encryption.key" failing_mod = types.ModuleType("keyring") failing_mod.get_password = lambda s, u: None failing_mod.set_password = lambda s, u, p: (_ for _ in ()).throw( Exception("no keyring") ) failing_mod.delete_password = lambda s, u: None with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=True), patch("talemate.util.encryption.keyring", failing_mod), ): reset_fernet() f = get_fernet() assert f is not None assert key_file.exists() def test_keyring_unavailable_uses_file(self, fresh_fernet): """When keyring is unavailable, behaves exactly like the old code.""" get_fernet() assert fresh_fernet.exists() # --------------------------------------------------------------------------- # Keyring edge cases # --------------------------------------------------------------------------- class TestKeyringEdgeCases: def test_keyring_disappears_after_store(self, tmp_path, mock_keyring_module): """ If keyring was used to store the key but is later unavailable, the user gets a new key (data loss on encrypted values, but no crash). """ key_file = tmp_path / "encryption.key" # First run: keyring works, key stored there with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=True), patch("talemate.util.encryption.keyring", mock_keyring_module), ): reset_fernet() encrypted = encrypt_value("sk-important") reset_fernet() # Second run: keyring unavailable, no file exists with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=False), ): reset_fernet() result = decrypt_value(encrypted) assert result is None def test_corrupted_keyring_value(self, tmp_path, mock_keyring_module): """If keyring returns a corrupted key, regeneration occurs gracefully.""" key_file = tmp_path / "encryption.key" mock_keyring_module.store[("talemate", "encryption_key")] = ( "not-a-valid-fernet-key" ) with ( patch("talemate.util.encryption._key_file_path", return_value=key_file), patch("talemate.util.encryption._keyring_available", return_value=True), patch("talemate.util.encryption.keyring", mock_keyring_module), ): reset_fernet() f = get_fernet() assert f is not None # A new valid key should now be in the keyring new_key = mock_keyring_module.store[("talemate", "encryption_key")] assert new_key != "not-a-valid-fernet-key" # --------------------------------------------------------------------------- # TALEMATE_DISABLE_KEYRING env var # --------------------------------------------------------------------------- class TestDisableKeyringEnvVar: def test_env_var_disables_keyring(self): """TALEMATE_DISABLE_KEYRING=1 forces file-based storage.""" with patch.dict(os.environ, {"TALEMATE_DISABLE_KEYRING": "1"}): assert _keyring_available() is False