Update app.py

app.py

@@ -1,4 +1,10 @@
 import os
+import json
+import hashlib
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+import base64
 import openai
 import gradio as gr
 from epub2txt import epub2txt
@@ -10,52 +16,192 @@ class GUI:
         self.prompt = os.getenv("prompt", "Summarize the following text:")
         self.client = None
         self.api_key = os.getenv("POE_API_KEY")
+        self.current_user = None
+        self.keys_file = "user_keys.json"
         
         with gr.Blocks(title="ePub Summarizer") as demo:
             with gr.Row():
-                gr.Markdown(scale=2).attach_load_event(self.hello, None)
-                gr.LoginButton()
-                gr.LogoutButton()
+                welcome_md = gr.Markdown(scale=2)
+                login_btn = gr.LoginButton()
+                logout_btn = gr.LogoutButton()
             
-            # API Key input section
-            with gr.Row(visible=not self.api_key):
-                with gr.Column():
-                    gr.Markdown("""
-                    ### Poe API Key Required
-                    
-                    To use this tool, you need a Poe API key:
-                    
-                    1. Visit [https://poe.com/api_key](https://poe.com/api_key)
-                    2. If you don't have an account, create one first
-                    3. Generate a new API key or copy your existing one
-                    4. Paste it in the field below
-                    
-                    **Note**: Your API key is only stored temporarily during this session for security.
-                    """)
-                    api_key_input = gr.Textbox(
-                        label="Poe API Key",
-                        placeholder="Enter your Poe API key here...",
-                        type="password"
+            # API Key input section (shown after login)
+            api_key_section = gr.Column(visible=False)
+            with api_key_section:
+                gr.Markdown("""
+                ### Poe API Key Setup
+                
+                To use this tool, you need a Poe API key:
+                
+                1. Visit [https://poe.com/api_key](https://poe.com/api_key)
+                2. If you don't have an account, create one first
+                3. Generate a new API key or copy your existing one
+                4. Paste it in the field below
+                
+                **Security**: Your API key will be encrypted and tied to your login account.
+                """)
+                api_key_input = gr.Textbox(
+                    label="Poe API Key",
+                    placeholder="Enter your Poe API key here...",
+                    type="password"
+                )
+                with gr.Row():
+                    remember_key = gr.Checkbox(
+                        label="Remember my API key (encrypted storage)", 
+                        value=True,
+                        info="Your key will be encrypted and saved for future sessions"
                     )
                     api_key_btn = gr.Button("Set API Key", variant="primary")
+                clear_key_btn = gr.Button("Clear Saved Key", variant="secondary", visible=False)
             
             out = gr.Markdown()
-            inp = gr.File(file_types=['.epub'], visible=bool(self.api_key))
-            
-            # Set up event handlers
-            if not self.api_key:
-                api_key_btn.click(
-                    self.set_api_key, 
-                    inputs=[api_key_input], 
-                    outputs=[out, inp]
-                )
-            else:
-                self._initialize_client()
+            inp = gr.File(file_types=['.epub'], visible=False)
             
+            # Event handlers
+            welcome_md.attach_load_event(self.on_load, None, [welcome_md, api_key_section, inp, clear_key_btn])
+            login_btn.login(self.on_login, None, [welcome_md, api_key_section, inp, clear_key_btn])
+            logout_btn.logout(self.on_logout, None, [welcome_md, api_key_section, inp, clear_key_btn])
+            api_key_btn.click(
+                self.set_api_key, 
+                inputs=[api_key_input, remember_key], 
+                outputs=[out, inp, api_key_section, clear_key_btn]
+            )
+            clear_key_btn.click(
+                self.clear_saved_key,
+                outputs=[out, api_key_section, clear_key_btn, api_key_input]
+            )
             inp.change(self.process, inp, out)
             
         demo.queue(concurrency_count=2).launch()
 
+    def _get_user_id(self, profile):
+        """Generate a unique user ID from profile"""
+        if not profile:
+            return None
+        # Use email as primary identifier, fallback to username
+        identifier = profile.email if hasattr(profile, 'email') and profile.email else profile.username
+        return hashlib.sha256(identifier.encode()).hexdigest()
+
+    def _generate_key_from_user(self, user_id):
+        """Generate encryption key from user ID"""
+        # Use user ID as salt for key derivation
+        salt = user_id.encode()[:32].ljust(32, b'0')  # Ensure 32 bytes
+        kdf = PBKDF2HMAC(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=salt,
+            iterations=100000,
+        )
+        key = base64.urlsafe_b64encode(kdf.derive(user_id.encode()))
+        return Fernet(key)
+
+    def _load_user_keys(self):
+        """Load encrypted user keys from file"""
+        try:
+            if os.path.exists(self.keys_file):
+                with open(self.keys_file, 'r') as f:
+                    return json.load(f)
+        except Exception as e:
+            print(f"Error loading user keys: {e}")
+        return {}
+
+    def _save_user_keys(self, keys_data):
+        """Save encrypted user keys to file"""
+        try:
+            with open(self.keys_file, 'w') as f:
+                json.dump(keys_data, f)
+        except Exception as e:
+            print(f"Error saving user keys: {e}")
+
+    def _get_saved_key(self, user_id):
+        """Retrieve and decrypt user's API key"""
+        try:
+            keys_data = self._load_user_keys()
+            if user_id in keys_data:
+                cipher_suite = self._generate_key_from_user(user_id)
+                encrypted_key = base64.urlsafe_b64decode(keys_data[user_id].encode())
+                return cipher_suite.decrypt(encrypted_key).decode()
+        except Exception as e:
+            print(f"Error retrieving saved key: {e}")
+        return None
+
+    def _save_encrypted_key(self, user_id, api_key):
+        """Encrypt and save user's API key"""
+        try:
+            keys_data = self._load_user_keys()
+            cipher_suite = self._generate_key_from_user(user_id)
+            encrypted_key = cipher_suite.encrypt(api_key.encode())
+            keys_data[user_id] = base64.urlsafe_b64encode(encrypted_key).decode()
+            self._save_user_keys(keys_data)
+            return True
+        except Exception as e:
+            print(f"Error saving encrypted key: {e}")
+            return False
+
+    def _delete_saved_key(self, user_id):
+        """Delete user's saved API key"""
+        try:
+            keys_data = self._load_user_keys()
+            if user_id in keys_data:
+                del keys_data[user_id]
+                self._save_user_keys(keys_data)
+                return True
+        except Exception as e:
+            print(f"Error deleting saved key: {e}")
+        return False
+
+    def on_load(self, profile: gr.OAuthProfile | None = None):
+        """Handle initial page load"""
+        if profile is None:
+            return (
+                gr.update(value='# ePub Summarization Tool\n\nPlease login to access the tool.'),
+                gr.update(visible=False),  # api_key_section
+                gr.update(visible=False),  # inp
+                gr.update(visible=False)   # clear_key_btn
+            )
+        return self.on_login(profile)
+
+    def on_login(self, profile: gr.OAuthProfile | None = None):
+        """Handle user login"""
+        if not profile:
+            return self.on_logout()
+        
+        self.current_user = self._get_user_id(profile)
+        user_name = profile.name if hasattr(profile, 'name') else profile.username
+        
+        # Check if user has a saved API key
+        saved_key = self._get_saved_key(self.current_user)
+        if saved_key:
+            self.api_key = saved_key
+            if self._initialize_client():
+                return (
+                    gr.update(value=f'# ePub Summarization Tool\n\nWelcome back {user_name}! ✅ Your saved API key is loaded and ready.'),
+                    gr.update(visible=False),  # api_key_section
+                    gr.update(visible=True),   # inp
+                    gr.update(visible=True)    # clear_key_btn
+                )
+        
+        # No saved key or failed to initialize
+        return (
+            gr.update(value=f'# ePub Summarization Tool\n\nWelcome {user_name}! Please set up your Poe API key below.'),
+            gr.update(visible=True),   # api_key_section
+            gr.update(visible=False),  # inp
+            gr.update(visible=False)   # clear_key_btn
+        )
+
+    def on_logout(self, profile: gr.OAuthProfile | None = None):
+        """Handle user logout"""
+        self.current_user = None
+        self.api_key = None
+        self.client = None
+        
+        return (
+            gr.update(value='# ePub Summarization Tool\n\nPlease login to access the tool.'),
+            gr.update(visible=False),  # api_key_section
+            gr.update(visible=False),  # inp
+            gr.update(visible=False)   # clear_key_btn
+        )
+
     def _initialize_client(self):
         """Initialize the Poe API client"""
         try:
@@ -68,10 +214,23 @@ class GUI:
             print(f"Error initializing Poe client: {e}")
             return False
 
-    def set_api_key(self, api_key):
+    def set_api_key(self, api_key, remember_key):
         """Set and validate the API key"""
+        if not self.current_user:
+            return (
+                gr.update(value="❌ Please login first."), 
+                gr.update(visible=False),
+                gr.update(visible=True),
+                gr.update(visible=False)
+            )
+        
         if not api_key or not api_key.strip():
-            return gr.update(value="⚠️ Please enter a valid API key."), gr.update(visible=False)
+            return (
+                gr.update(value="⚠️ Please enter a valid API key."), 
+                gr.update(visible=False),
+                gr.update(visible=True),
+                gr.update(visible=False)
+            )
         
         self.api_key = api_key.strip()
         
@@ -83,33 +242,80 @@ class GUI:
                     messages=[{"role": "user", "content": "Hello"}],
                     max_tokens=10
                 )
+                
+                # Save key if user wants to remember it
+                if remember_key:
+                    if self._save_encrypted_key(self.current_user, self.api_key):
+                        success_msg = "✅ API key validated and saved successfully! You can now upload an ePub file."
+                    else:
+                        success_msg = "✅ API key validated successfully! (Note: Failed to save for future use)"
+                else:
+                    success_msg = "✅ API key validated successfully! You can now upload an ePub file."
+                
                 return (
-                    gr.update(value="✅ API key validated successfully! You can now upload an ePub file."), 
-                    gr.update(visible=True)
+                    gr.update(value=success_msg), 
+                    gr.update(visible=True),   # inp
+                    gr.update(visible=False),  # api_key_section
+                    gr.update(visible=remember_key)  # clear_key_btn
                 )
             except Exception as e:
                 error_msg = str(e)
                 if "401" in error_msg or "unauthorized" in error_msg.lower():
                     return (
                         gr.update(value="❌ Invalid API key. Please check your key and try again."), 
+                        gr.update(visible=False),
+                        gr.update(visible=True),
                         gr.update(visible=False)
                     )
                 elif "quota" in error_msg.lower() or "limit" in error_msg.lower():
                     return (
                         gr.update(value="⚠️ API key valid but quota exceeded. Please check your Poe account."), 
+                        gr.update(visible=False),
+                        gr.update(visible=True),
                         gr.update(visible=False)
                     )
                 else:
                     return (
                         gr.update(value=f"❌ API connection error: {error_msg}"), 
+                        gr.update(visible=False),
+                        gr.update(visible=True),
                         gr.update(visible=False)
                     )
         else:
             return (
                 gr.update(value="❌ Failed to initialize API client."), 
+                gr.update(visible=False),
+                gr.update(visible=True),
                 gr.update(visible=False)
             )
 
+    def clear_saved_key(self):
+        """Clear the user's saved API key"""
+        if not self.current_user:
+            return (
+                gr.update(value="❌ No user logged in."),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(value="")
+            )
+        
+        if self._delete_saved_key(self.current_user):
+            self.api_key = None
+            self.client = None
+            return (
+                gr.update(value="✅ Saved API key cleared. Please enter your API key below."),
+                gr.update(visible=True),   # api_key_section
+                gr.update(visible=False),  # clear_key_btn
+                gr.update(value="")        # clear input
+            )
+        else:
+            return (
+                gr.update(value="⚠️ Failed to clear saved key."),
+                gr.update(visible=True),
+                gr.update(visible=True),
+                gr.update(value="")
+            )
+
     def get_model_response(self, text: str) -> str:
         """
         Get response from Poe API for the given text
@@ -129,7 +335,7 @@ class GUI:
 
     def process(self, file, profile: gr.OAuthProfile | None = None):
         if profile is None:
-            return gr.update(value='Login to access the tool.')
+            return gr.update(value='⚠️ Please login to access the tool.')
 
         if not self.client:
             return gr.update(value='⚠️ Please set your Poe API key first.')
@@ -209,15 +415,6 @@ class GUI:
         except Exception as e:
             yield gr.update(value=f"Error processing file: {str(e)}")
 
-    def hello(self, profile: gr.OAuthProfile | None = None):
-        if profile is None:
-            return '# ePub Summarization Tool\n\nLogin to access the tool.'
-        
-        if not self.api_key:
-            return f"# ePub Summarization Tool\n\nWelcome {profile.name}!\n\nPlease set your Poe API key below to get started."
-        
-        return f"# ePub Summarization Tool\n\nWelcome {profile.name}! Ready to summarize ePub files."
-
 # Run the application
 if __name__ == "__main__":
     GUI()