diff --git a/src/firebolt/utils/token_storage.py b/src/firebolt/utils/token_storage.py index 131c76f8e43..a3cb3adc987 100644 --- a/src/firebolt/utils/token_storage.py +++ b/src/firebolt/utils/token_storage.py @@ -3,7 +3,8 @@ from json import JSONDecodeError from json import dump as json_dump from json import load as json_load -from os import makedirs, path, urandom +from os import makedirs, path, scandir, unlink, urandom +from re import fullmatch from time import time from typing import Optional @@ -27,23 +28,18 @@ def generate_salt() -> str: return b64encode(urandom(16)).decode("ascii") -def generate_file_name(username: str, password: str) -> str: - """Generate unique file name based on username and password. - - Username and password values are not exposed. +def generate_file_name(username: str) -> str: + """Generate a cache file name from the account or client identity. Args: username (str): Username - password (str): Password Returns: - str: File name 64 characters long + str: Versioned cache file name """ - username_hash = sha256(username.encode("utf-8")).hexdigest()[:32] - password_hash = sha256(password.encode("utf-8")).hexdigest()[:32] - - return f"{username_hash}{password_hash}.json" + username_hash = sha256(username.encode("utf-8")).hexdigest() + return f"token-v2-{username_hash}.json" class TokenSecureStorage: @@ -59,14 +55,25 @@ class TokenSecureStorage: def __init__(self, username: str, password: str): self._data_dir = user_data_dir(appname=APPNAME) makedirs(self._data_dir, exist_ok=True) + self._remove_legacy_cache_files() - self._token_file = path.join( - self._data_dir, generate_file_name(username, password) - ) + self._token_file = path.join(self._data_dir, generate_file_name(username)) self.salt = self._get_salt() self.encrypter = FernetEncrypter(self.salt, username, password) + def _remove_legacy_cache_files(self) -> None: + # Legacy names expose password hashes even after their tokens expire. + with scandir(self._data_dir) as entries: + for entry in entries: + if fullmatch(r"[0-9a-f]{64}\.json", entry.name) and entry.is_file( + follow_symlinks=False + ): + try: + unlink(entry.path) + except FileNotFoundError: + pass + def _get_salt(self) -> str: """Get salt from the file if exists, or generate a new one. diff --git a/tests/unit/common/test_token_storage.py b/tests/unit/common/test_token_storage.py index 4e441e877fa..e2ef8dfa6b1 100644 --- a/tests/unit/common/test_token_storage.py +++ b/tests/unit/common/test_token_storage.py @@ -1,7 +1,6 @@ import os from unittest.mock import patch -from appdirs import user_config_dir from pyfakefs.fake_filesystem import FakeFilesystem from firebolt.utils.token_storage import ( @@ -99,13 +98,68 @@ def test_token_storage_json_broken(fs: FakeFilesystem): """ settings = {"username": "username", "password": "password"} - data_dir = os.path.join(user_config_dir(), "firebolt") - fs.create_dir(data_dir) - fs.create_file(os.path.join(data_dir, "token.json"), contents="{Not a valid json") + storage = TokenSecureStorage(**settings) + fs.create_file(storage._token_file, contents="{Not a valid json") assert TokenSecureStorage(**settings).get_cached_token() is None +@patch("firebolt.utils.token_storage.time", return_value=0) +def test_cache_filename_does_not_depend_on_password( + _mock_time, fs: FakeFilesystem +) -> None: + original = TokenSecureStorage(username="username", password="password") + original.cache_token("token", 1) + + changed = TokenSecureStorage( + username="username", + password="different-password", # legit:ignore-secrets + ) + assert changed._token_file == original._token_file + assert changed.get_cached_token() is None + assert original.get_cached_token() == "token" + + +@patch("firebolt.utils.token_storage.time", return_value=0) +def test_legacy_cache_is_discarded_without_migration( + _mock_time, fs: FakeFilesystem +) -> None: + storage = TokenSecureStorage(username="username", password="password") + storage.cache_token("legacy-token", 1) + legacy_file = os.path.join(storage._data_dir, "a" * 64 + ".json") + os.rename(storage._token_file, legacy_file) + + fresh = TokenSecureStorage(username="username", password="password") + assert not os.path.exists(legacy_file) + assert fresh.get_cached_token() is None + fresh.cache_token("fresh-token", 1) + assert fresh.get_cached_token() == "fresh-token" + + +@patch("firebolt.utils.token_storage.time", return_value=0) +def test_cleanup_preserves_new_cache_and_unrelated_files( + _mock_time, fs: FakeFilesystem +) -> None: + storage = TokenSecureStorage(username="username", password="password") + storage.cache_token("current-token", 1) + legacy_files = [ + os.path.join(storage._data_dir, character * 64 + ".json") + for character in ("a", "b") + ] + for filename in legacy_files: + fs.create_file(filename, contents='{"expiration": 0}') + unrelated_file = os.path.join(storage._data_dir, "notes.json") + fs.create_file(unrelated_file, contents="unrelated") + unrelated_dir = os.path.join(storage._data_dir, "c" * 64 + ".json") + fs.create_dir(unrelated_dir) + + fresh = TokenSecureStorage(username="username", password="password") + assert all(not os.path.exists(filename) for filename in legacy_files) + assert os.path.isfile(unrelated_file) + assert os.path.isdir(unrelated_dir) + assert fresh.get_cached_token() == "current-token" + + @patch("firebolt.utils.token_storage.time", return_value=0) def test_multiple_tokens(fs: FakeFilesystem) -> None: """ diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4daff6954a8..05a16329c25 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -41,6 +41,17 @@ def global_fake_fs(request) -> None: yield +@fixture(autouse=True) +def isolate_token_cache(monkeypatch, request) -> None: + if "nofakefs" not in request.keywords: + return + cache_dir = request.getfixturevalue("tmp_path") + monkeypatch.setattr( + "firebolt.utils.token_storage.user_data_dir", + lambda appname: str(cache_dir / appname), + ) + + @fixture(autouse=True) def clear_cache() -> None: _firebolt_cache.clear()