Add suggested changes

This commit is contained in:
hawkeye116477 2026-08-16 12:26:50 +02:00
parent abe96b5eea
commit 557b16af7f
No known key found for this signature in database
GPG key ID: 66367B184988BC50
2 changed files with 78 additions and 68 deletions

View file

@ -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

View file

@ -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