From 557b16af7f5168d8236133fc31d19fa2e197aeda Mon Sep 17 00:00:00 2001 From: hawkeye116477 Date: Sun, 16 Aug 2026 12:26:50 +0200 Subject: [PATCH] Add suggested changes --- legendary/lfs/lgndry.py | 62 +++++++++++++++++++++++++++--- legendary/lfs/utils.py | 84 +++++++++++------------------------------ 2 files changed, 78 insertions(+), 68 deletions(-) diff --git a/legendary/lfs/lgndry.py b/legendary/lfs/lgndry.py index eb33aaf..c434e13 100644 --- a/legendary/lfs/lgndry.py +++ b/legendary/lfs/lgndry.py @@ -1,4 +1,4 @@ - +from hashlib import md5 import json import logging import os @@ -15,6 +15,11 @@ from legendary.utils.aliasing import generate_aliases from legendary.utils.env import is_windows_mac_or_pyi from .utils import LockedJSONData, clean_filename +from .utils import ( + decrypt_file, + encrypt_to_file, + remove_encryption_key +) FILELOCK_DEBUG = False @@ -148,11 +153,48 @@ class LGDLFS: @contextmanager def userdata_lock(self) -> LockedJSONData: """Wrapper around the lock to automatically update user data when it is released""" - with LockedJSONData(os.path.join(self.path, 'current_user.json')) as lock: - try: - yield lock - finally: - self._user_data = lock.data + if not self.config.getboolean('Legendary', 'disable_token_encryption', fallback=False): + with LockedJSONData(lock_file = os.path.join(self.path, 'current_user.json'), save_changes=False) as lock: + current_user_data = None + old_tokens_data = None + migrate_non_encrypted = False + try: + current_user_data = lock.data + non_encrypted_path = os.path.join(self.path, "user.json") + + if current_user_data and (account_id := current_user_data.get('account_id')) is not None: + data_file_path = os.path.join(self.path, f"{md5(account_id.encode('utf-8')).hexdigest()}.enc") + if os.path.exists(data_file_path): + old_tokens_data = decrypt_file(data_file_path, current_user_data) + elif os.path.exists(non_encrypted_path): + migrate_non_encrypted = True + with open(non_encrypted_path, "r", encoding='utf-8') as f: + old_tokens_data = json.load(f) + os.remove(non_encrypted_path) + + if old_tokens_data: + lock.data = old_tokens_data + yield lock + finally: + new_user_data = {} + if lock.data: + if (account_id := lock.data.get('account_id')) is not None: + new_user_data['account_id'] = account_id + if (display_name := lock.data.get('displayName')) is not None: + new_user_data['displayName'] = display_name + + if (old_tokens_data != lock.data or migrate_non_encrypted) and (account_id := new_user_data.get('account_id')) is not None: + new_data_filename = f"{md5(account_id.encode('utf-8')).hexdigest()}.enc" + new_user_data = encrypt_to_file(os.path.join(self.path, new_data_filename), new_user_data, lock.data) + with open(os.path.join(self.path, "current_user.json"), 'w', encoding='utf-8') as f: + json.dump(new_user_data, f, indent=2, sort_keys=True) + self._user_data = lock.data + else: + with LockedJSONData(lock_file = os.path.join(self.path, 'user.json')) as lock: + try: + yield lock + finally: + self._user_data = lock.data @property def userdata(self): @@ -172,6 +214,14 @@ class LGDLFS: def invalidate_userdata(self): with self.userdata_lock as lock: + userdata_file = os.path.join(self.path, 'current_user.json') + if lock.data and (account_id := lock.data.get('account_id')) is not None: + remove_encryption_key(lock.data) + old_data_file = os.path.join(self.path, f"{md5(account_id.encode('utf-8')).hexdigest()}.enc") + if os.path.exists(old_data_file): + os.remove(old_data_file) + if os.path.exists(userdata_file): + os.remove(userdata_file) lock.clear() @property diff --git a/legendary/lfs/utils.py b/legendary/lfs/utils.py index 3fcc5b0..00463c3 100644 --- a/legendary/lfs/utils.py +++ b/legendary/lfs/utils.py @@ -210,21 +210,23 @@ def decrypt_file(path, current_user_info): json_decrypted_data = None return json_decrypted_data -def encrypt_to_file(path, current_user_info, data): +def encrypt_to_file(path, current_user_info, data, store_key_file=False): final_encryption_key = get_encryption_key(current_user_info) if not final_encryption_key: encryption_key = base64.b64encode(os.urandom(32)).decode("utf-8") - try: - service_name = get_service_for_keyring(current_user_info) - k_backend = keyring.core.get_keyring() - if os.name == 'nt': - k_backend.persist = 'local machine' - if service_name is not None: - k_backend.set_password(service_name, current_user_info['account_id'], encryption_key) - except Exception: - current_user_info['key'] = encryption_key - finally: - final_encryption_key = get_encryption_key(current_user_info) + if not store_key_file: + try: + service_name = get_service_for_keyring(current_user_info) + k_backend = keyring.core.get_keyring() + if os.name == 'nt': + k_backend.persist = 'local machine' + if service_name is not None: + k_backend.set_password(service_name, current_user_info['account_id'], encryption_key) + except Exception: + store_key_file = True + if store_key_file: + current_user_info['key'] = encryption_key + final_encryption_key = get_encryption_key(current_user_info) iv_cipher = AES.new(final_encryption_key, AES.MODE_ECB) cipher = AES.new(final_encryption_key, AES.MODE_CBC) input_data = json.dumps(data).encode('utf-8') @@ -236,66 +238,31 @@ def encrypt_to_file(path, current_user_info, data): class LockedJSONData(FileLock): - def __init__(self, lock_file: str): + def __init__(self, lock_file: str, save_changes: bool = True): super().__init__(lock_file + '.lock') self._file_path = lock_file self._data = None - self._user_data = None self._initial_data = None + self._save_changes = save_changes def __enter__(self): super().__enter__() if os.path.exists(self._file_path): with open(self._file_path, 'r', encoding='utf-8') as f: - try: - self._user_data = json.load(f) - self._initial_data = self._user_data - except json.JSONDecodeError: - pass - if self._user_data and (account_id := self._user_data.get('account_id')) is not None: - data_file_path = os.path.join(os.path.dirname(self._file_path), f"{hashlib.md5(account_id.encode('utf-8')).hexdigest()}.enc") - if os.path.exists(data_file_path): - self._data = decrypt_file(data_file_path, self._user_data) - else: - # Migrate non-encrypted data - non_encrypted_path = os.path.join(os.path.dirname(self._file_path), "user.json") - if os.path.exists(non_encrypted_path): - with open(non_encrypted_path, "r", encoding='utf-8') as f: - self._data = json.load(f) - os.remove(non_encrypted_path) + self._data = json.load(f) + self._initial_data = self._data return self def __exit__(self, exc_type, exc_val, exc_tb): super().__exit__(exc_type, exc_val, exc_tb) - if self._user_data is None: - self._user_data = self._data - new_user_data = None - full_old_data = None - old_data_filename = None - if self._initial_data and (initial_account_id := self._initial_data.get('account_id')) is not None: - old_data_filename = f"{hashlib.md5(initial_account_id.encode('utf-8')).hexdigest()}.enc" - - if self._user_data: - new_user_data = {} - if (account_id := self._user_data.get('account_id')) is not None: - new_user_data['account_id'] = account_id - if (display_name := self._user_data.get('displayName')) is not None: - new_user_data['displayName'] = display_name - if old_data_filename: - full_old_data = decrypt_file(os.path.join(os.path.dirname(self._file_path), old_data_filename), self._initial_data) - - if full_old_data != self._data: - if self._user_data and self._data and (account_id := self._user_data.get('account_id')) is not None: - new_data_filename = f"{hashlib.md5(account_id.encode('utf-8')).hexdigest()}.enc" - new_user_data = encrypt_to_file(os.path.join(os.path.dirname(self._file_path), new_data_filename), new_user_data, self._data) - - if self._initial_data != new_user_data: - if new_user_data: + + if self._data != self._initial_data and self._save_changes: + if self._data is not None: with open(self._file_path, 'w', encoding='utf-8') as f: - json.dump(new_user_data, f, indent=2, sort_keys=True) + json.dump(self._data, f, indent=2, sort_keys=True) else: if os.path.exists(self._file_path): os.remove(self._file_path) @@ -311,11 +278,4 @@ class LockedJSONData(FileLock): self._data = new_data def clear(self): - if self._user_data: - if (account_id := self._user_data.get('account_id')) is not None: - remove_encryption_key(self._user_data) - new_data_file = os.path.join(os.path.dirname(self._file_path),f"{hashlib.md5(account_id.encode('utf-8')).hexdigest()}.enc") - if os.path.exists(new_data_file): - os.remove(new_data_file) - self._user_data = None self._data = None