Merge pull request #11 from FragmentScreen/fix/keyring-alt

Keyring Alternative

Author: Lui-Instruct

fGOaria/classes/oauth.py

@@ -2,6 +2,7 @@
 from .client_oauth import ClientOauth
 from ..utils.imports_config import *
 from .token import Token
+from ..utils.encryption import TokenEncryption
 
 
 class OAuth :
@@ -16,6 +17,7 @@ def __init__(self) -> None:
         self.client_secret =  os.getenv('ARIA_CLIENT_SECRET')
         self.token_str_key = os.getenv('ARIA_KEYRING')
         self.refresh_grant = os.getenv('ARIA_CONNECTION_REFRESH_GRANT')
+        self.token_encryption = TokenEncryption()
 
 
     # LOGIN
@@ -79,19 +81,34 @@ def refresh_token(self, token : Token) -> Union[Token, None] :
     #  KEYRING STORAGE
 
     def get_keyring_token_data(self) -> Union[dict, None] :
-        token_data_str = keyring.get_password(self.token_str_key, '')
-        if not token_data_str :
-            raise Exception(' Either the password entered is incorrect, or no access token is stored.')
-        return json.loads(token_data_str)
+        try:
+            token_data_str = keyring.get_password(self.token_str_key, '')
+            if not token_data_str :
+                token_data = self.token_encryption.decrypt_token(self.password)
+                if token_data:
+                    return token_data
+                raise Exception(' Either the password entered is incorrect, or no access token is stored.')
+            return json.loads(token_data_str)
+        except Exception as e:
+            token_data = self.token_encryption.decrypt_token(self.password)
+            if token_data:
+                return token_data
+            raise Exception('Failed to retrieve token from both keyring and fallback storage.')
 
     def set_token_keyring_data(self, token : Token) -> None :
         try :
             click.echo('Attempting to store Token...')
             token_json = json.dumps(token.to_dict())
-            keyring.set_password(self.token_str_key, '', token_json)
-            print_with_spaces('Token data successfully stored in keyring.')
-        except key_err.PasswordSetError as e :
-            logging.error(f"Error setting keyring: {e}")
+            try:
+                keyring.set_password(self.token_str_key, '', token_json)
+                print_with_spaces('Token data successfully stored in keyring.')
+            except key_err.PasswordSetError as e :
+                logging.warning(f"Keyring storage failed, using fallback encryption: {e}")
+                self.token_encryption.encrypt_token(token.to_dict(), self.password)
+                print_with_spaces('Token data successfully stored using fallback encryption.')
+        except Exception as e:
+            logging.error(f"Error storing token: {e}")
+            raise Exception("Failed to store token in both keyring and fallback storage.")
 
 
     # GET DATA OBJECTS

fGOaria/utils/encryption.py

@@ -0,0 +1,88 @@
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+import base64
+import os
+import json
+import logging
+from pathlib import Path
+import platform
+
+class TokenEncryption:
+    def __init__(self):
+        self.token_file = Path.home() / '.aria_token'
+        self.salt_file = Path.home() / '.aria_salt'
+        self._ensure_salt_exists()
+        logging.info(f"TokenEncryption initialized on {platform.system()}")
+        
+    def _ensure_salt_exists(self):
+        """Ensure the salt file exists and contains a valid salt."""
+        try:
+            if not self.salt_file.exists():
+                salt = os.urandom(16)
+                self.salt_file.write_bytes(salt)
+                logging.info(f"Created new salt file at {self.salt_file}")
+            else:
+                salt = self.salt_file.read_bytes()
+                if len(salt) != 16:
+                    salt = os.urandom(16)
+                    self.salt_file.write_bytes(salt)
+                    logging.info("Regenerated invalid salt file")
+        except Exception as e:
+            logging.error(f"Error managing salt file: {str(e)}")
+            raise Exception(f"Failed to manage salt file: {str(e)}")
+    
+    def _get_key(self, password: str) -> bytes:
+        """Derive an encryption key from the password using PBKDF2."""
+        try:
+            salt = self.salt_file.read_bytes()
+            kdf = PBKDF2HMAC(
+                algorithm=hashes.SHA256(),
+                length=32,
+                salt=salt,
+                iterations=100000,
+            )
+            key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
+            return key
+        except Exception as e:
+            logging.error(f"Error deriving encryption key: {str(e)}")
+            raise Exception(f"Failed to derive encryption key: {str(e)}")
+    
+    def encrypt_token(self, token_data: dict, password: str) -> None:
+        """Encrypt and store token data."""
+        try:
+            key = self._get_key(password)
+            f = Fernet(key)
+            encrypted_data = f.encrypt(json.dumps(token_data).encode())
+            self.token_file.write_bytes(encrypted_data)
+            logging.info(f"Successfully encrypted and stored token at {self.token_file}")
+        except Exception as e:
+            logging.error(f"Error encrypting token: {str(e)}")
+            raise Exception(f"Failed to encrypt token: {str(e)}")
+    
+    def decrypt_token(self, password: str) -> dict:
+        """Decrypt and retrieve token data."""
+        if not self.token_file.exists():
+            logging.info("No token file found")
+            return None
+            
+        try:
+            key = self._get_key(password)
+            f = Fernet(key)
+            encrypted_data = self.token_file.read_bytes()
+            decrypted_data = f.decrypt(encrypted_data)
+            logging.info("Successfully decrypted token")
+            return json.loads(decrypted_data)
+        except Exception as e:
+            logging.error(f"Error decrypting token: {str(e)}")
+            return None
+    
+    def clear_token(self) -> None:
+        """Remove stored token data."""
+        try:
+            if self.token_file.exists():
+                self.token_file.unlink()
+                logging.info("Successfully cleared token file")
+        except Exception as e:
+            logging.error(f"Error clearing token file: {str(e)}")
+            raise Exception(f"Failed to clear token file: {str(e)}") 

requirements.txt

@@ -1,6 +1,8 @@
 click==8.1.7
+requests==2.31.0
 keyring==24.3.0
-PyYAML==6.0.1
+python-dotenv==1.0.0
 questionary==2.0.1
-Requests==2.31.0
+PyYAML==6.0.1
+cryptography==42.0.2
 setuptools==68.2.2

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name='fandanGO-aria',
-    version='2.2.0',
+    version='2.2.1',
     description="ARIA connection for managing access and metadata deposition, primarily through fandanGO",
     long_description=long_description,
     long_description_content_type='text/markdown',
@@ -22,6 +22,7 @@
         'python-dotenv',
         'questionary',
         'PyYAML',
+        'cryptography',
     ],
     entry_points={
         'console_scripts': [