diff --git a/legendary/lfs/lgndry.py b/legendary/lfs/lgndry.py index 875f67a..eb33aaf 100644 --- a/legendary/lfs/lgndry.py +++ b/legendary/lfs/lgndry.py @@ -148,7 +148,7 @@ 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, 'user.json')) as lock: + with LockedJSONData(os.path.join(self.path, 'current_user.json')) as lock: try: yield lock finally: diff --git a/legendary/lfs/utils.py b/legendary/lfs/utils.py index 7de9e21..3fcc5b0 100644 --- a/legendary/lfs/utils.py +++ b/legendary/lfs/utils.py @@ -1,14 +1,18 @@ - +import base64 import hashlib import json import logging import os import shutil from collections.abc import Iterator +from contextlib import suppress from pathlib import Path from sys import stdout from time import perf_counter +import keyring +from Cryptodome.Cipher import AES +from Cryptodome.Util.Padding import pad, unpad from filelock import FileLock from legendary.lfs.wine_helpers import case_insensitive_file_search @@ -165,6 +169,71 @@ def clean_filename(filename): def get_dir_size(path): return sum(f.stat().st_size for f in Path(path).glob('**/*') if f.is_file()) +def get_service_for_keyring(current_user_info): + service_name = "legendary" + if os.name == 'nt': + service_name = f"legendary/{current_user_info['account_id']}" + return service_name + +def remove_encryption_key(current_user_info): + with suppress(keyring.errors.PasswordDeleteError): + keyring.delete_password(get_service_for_keyring(current_user_info), current_user_info['account_id']) + +def get_encryption_key(current_user_info): + final_key = "" + key = "" + try: + key = keyring.get_password(get_service_for_keyring(current_user_info), current_user_info['account_id']) + except Exception: + if current_user_info['account_id'] is not None and current_user_info["key"] is not None: + final_key = hashlib.sha256((current_user_info['account_id'] + current_user_info["key"]).encode("utf-8")).digest() + if key is not None and final_key == "": + final_key = base64.b64decode(key.encode('utf-8')) + return final_key + +def decrypt_file(path, current_user_info): + try: + key = get_encryption_key(current_user_info) + if key is None or len(key) != 32: + return "" + encrypted_data = None + with open(path, "rb") as encrypted_file_content: + encrypted_data = encrypted_file_content.read() + iv_cipher = AES.new(key, AES.MODE_ECB) + iv = iv_cipher.decrypt(encrypted_data[:16]) + cipher = AES.new(key, AES.MODE_CBC, iv) + decrypted_data = unpad(cipher.decrypt(encrypted_data[16:]), AES.block_size).decode("utf-8") + json_decrypted_data = json.loads(decrypted_data) + except Exception as ex: + logger.warn(f'Failed to decrypt data with {ex!r}') + decrypted_data = None + json_decrypted_data = None + return json_decrypted_data + +def encrypt_to_file(path, current_user_info, data): + 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) + 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') + encrypted_data = cipher.encrypt(pad(input_data, AES.block_size)) + encrypted_iv = iv_cipher.encrypt(cipher.iv) + with open(path, 'wb') as f: + f.write(encrypted_iv + encrypted_data) + return current_user_info + class LockedJSONData(FileLock): def __init__(self, lock_file: str): @@ -172,6 +241,7 @@ class LockedJSONData(FileLock): self._file_path = lock_file self._data = None + self._user_data = None self._initial_data = None def __enter__(self): @@ -179,17 +249,53 @@ class LockedJSONData(FileLock): if os.path.exists(self._file_path): with open(self._file_path, 'r', encoding='utf-8') as f: - self._data = json.load(f) - self._initial_data = self._data + 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) return self def __exit__(self, exc_type, exc_val, exc_tb): super().__exit__(exc_type, exc_val, exc_tb) - if self._data != self._initial_data: - if self._data is not None: + 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: with open(self._file_path, 'w', encoding='utf-8') as f: - json.dump(self._data, f, indent=2, sort_keys=True) + json.dump(new_user_data, f, indent=2, sort_keys=True) else: if os.path.exists(self._file_path): os.remove(self._file_path) @@ -205,4 +311,11 @@ 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 diff --git a/pyproject.toml b/pyproject.toml index c564d53..289b59d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "requests", "filelock", "pycryptodomex", + "keyring" ] requires-python = ">= 3.10" authors = [