Encrypt tokens using a locally generated key

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:

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

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "requests",
     "filelock",
     "pycryptodomex",
+    "keyring"
 ]
 requires-python = ">= 3.10"
 authors = [