Update app.py

app.py

@@ -8,124 +8,120 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 # File where encrypted data is stored
 VAULT_FILE = "passwords.json"
-FAILED_ATTEMPTS = 0
-MAX_ATTEMPTS = 5
+# In a real app, the salt would be stored securely or unique per user
+SALT = b'hassan_naseer_permanent_salt_2026'
 
 # --- Core Security Logic ---
 
 def get_cipher(master_password):
-    """Derives a strong AES-256 key from the master password."""
-    salt = b'hassan_naseer_secure_salt_2026' 
     kdf = PBKDF2HMAC(
         algorithm=hashes.SHA256(),
         length=32,
-        salt=salt,
+        salt=SALT,
         iterations=150000,
     )
     key = base64.urlsafe_b64encode(kdf.derive(master_password.encode()))
     return Fernet(key)
 
-def process_vault(action, master_pwd, service=None, secret=None):
-    global FAILED_ATTEMPTS
-    
-    if FAILED_ATTEMPTS >= MAX_ATTEMPTS:
-        return "🛑 SECURITY LOCKOUT: Too many failed attempts. Access disabled."
+def load_vault():
+    if os.path.exists(VAULT_FILE):
+        with open(VAULT_FILE, "r") as f:
+            return json.load(f)
+    return {}
+
+def save_to_vault(vault_data):
+    with open(VAULT_FILE, "w") as f:
+        json.dump(vault_data, f)
+
+# --- App Logic ---
 
+def login(master_pwd):
     if not master_pwd:
-        return "⚠️ Master Password required to access Vault."
+        return gr.update(visible=True), gr.update(visible=False), "⚠️ Enter Password"
+    
+    # We verify the password by attempting to derive the key
+    try:
+        # If the vault exists, we try to decrypt the first key as a test
+        vault = load_vault()
+        if vault:
+            cipher = get_cipher(master_pwd)
+            # Try decrypting the first stored item to verify password
+            test_key = list(vault.keys())[0]
+            cipher.decrypt(vault[test_key].encode())
+        
+        return gr.update(visible=False), gr.update(visible=True), ""
+    except Exception:
+        # If decryption fails, password is wrong (unless vault is empty)
+        if not load_vault():
+             return gr.update(visible=False), gr.update(visible=True), ""
+        return gr.update(visible=True), gr.update(visible=False), "🚫 Incorrect Master Password"
 
+def process_action(action, master_pwd, service=None, secret=None):
     try:
         cipher = get_cipher(master_pwd)
-        
-        if os.path.exists(VAULT_FILE):
-            with open(VAULT_FILE, "r") as f:
-                vault = json.load(f)
-        else:
-            vault = {}
+        vault = load_vault()
 
         if action == "Save":
             if not service or not secret:
-                return "⚠️ Service name and Password are required!"
-            
-            encrypted_data = cipher.encrypt(secret.encode()).decode()
-            vault[service] = encrypted_data
-            
-            with open(VAULT_FILE, "w") as f:
-                json.dump(vault, f)
-            
-            FAILED_ATTEMPTS = 0
-            return f"✅ Encrypted & Saved: {service}"
+                return "⚠️ Fill all fields"
+            vault[service] = cipher.encrypt(secret.encode()).decode()
+            save_to_vault(vault)
+            return f"✅ Saved: {service}"
 
         elif action == "Retrieve":
-            if not service:
-                return "⚠️ Enter a service name to decrypt."
             if service not in vault:
-                return f"❌ No record found for '{service}'."
-            
-            decrypted_data = cipher.decrypt(vault[service].encode()).decode()
+                return "❌ Not found"
+            decrypted = cipher.decrypt(vault[service].encode()).decode()
+            return f"🔑 Password for {service}: {decrypted}"
             
-            FAILED_ATTEMPTS = 0
-            return f"🔑 Password for {service}: {decrypted_data}"
-
     except Exception:
-        FAILED_ATTEMPTS += 1
-        remaining = MAX_ATTEMPTS - FAILED_ATTEMPTS
-        return f"🚫 Access Denied! {remaining} attempts remaining."
+        return "🚫 Critical Error: Check Master Password"
 
 # --- UI Styling ---
-
 custom_css = """
-.gradio-container { max-width: 750px; margin: auto; padding: 20px; }
-.header-box { text-align: center; margin-bottom: 25px; border-bottom: 1px solid #ddd; padding-bottom: 10px; }
-@media (max-width: 768px) {
-    .gradio-container { padding: 10px; }
-    h1 { font-size: 1.5em !important; }
-}
+.gradio-container { max-width: 600px; margin: auto; padding: 20px; }
+.login-box { border: 1px solid #ddd; padding: 30px; border-radius: 15px; background: #fdfdfd; }
+.vault-box { padding: 20px; }
 """
 
-# Fixed: Moved title to gr.Blocks constructor
-with gr.Blocks(title="SecureVault | Hassan Naseer") as demo:
-    gr.HTML("""
-        <div class="header-box">
-            <h1 style="display:inline; color: #2d3436;">🔐 SecureVault Manager</h1>
-            <p style="font-size: 0.9em; color: #636e72;">Developed by Hassan Naseer | AES-256 Encryption</p>
-        </div>
-    """)
-
-    with gr.Row():
-        master_key = gr.Textbox(
-            label="Unlock Vault with Master Password", 
-            type="password", 
-            placeholder="Enter Master Password..."
-        )
-
-    with gr.Tabs():
-        with gr.Tab(label="📥 Add Password"):
-            with gr.Column():
-                svc = gr.Textbox(label="Service", placeholder="e.g., Google, Bank, GitHub")
-                pwd = gr.Textbox(label="Password to Protect", type="password")
-                save_btn = gr.Button("🔒 Encrypt & Store", variant="primary")
-
-        with gr.Tab(label="🔓 Retrieve Password"):
-            with gr.Column():
-                search_svc = gr.Textbox(label="Service Name", placeholder="Enter service to search")
-                get_btn = gr.Button("🔑 Decrypt & Show", variant="secondary")
-
-    status_output = gr.Textbox(label="Operation Status", interactive=False)
-
-    # Event Handlers
-    save_btn.click(
-        fn=lambda m, s, p: process_vault("Save", m, s, p),
-        inputs=[master_key, svc, pwd],
-        outputs=status_output
-    )
+with gr.Blocks(title="SecureVault Pro | Hassan Naseer") as demo:
+    
+    # 1. LOGIN PANEL
+    with gr.Column(visible=True, elem_classes="login-box") as login_panel:
+        gr.Markdown("# 🔐 SecureVault Login")
+        gr.Markdown("Enter your Master Password to decrypt your history.")
+        master_input = gr.Textbox(label="Master Password", type="password")
+        login_btn = gr.Button("Unlock Vault", variant="primary")
+        login_msg = gr.Markdown(color="red")
+
+    # 2. MAIN VAULT PANEL (Hidden by default)
+    with gr.Column(visible=False, elem_classes="vault-box") as vault_panel:
+        gr.HTML(f"""<h2 style='text-align:center;'>Welcome Back, Hassan</h2><p style='text-align:center; color:grey;'>AES-256 Encrypted History Loaded</p>""")
+        
+        with gr.Tabs():
+            with gr.Tab(label="➕ Add New"):
+                svc = gr.Textbox(label="Service")
+                pwd = gr.Textbox(label="Password", type="password")
+                save_btn = gr.Button("Save Encrypted", variant="primary")
 
-    get_btn.click(
-        fn=lambda m, s: process_vault("Retrieve", m, s),
-        inputs=[master_key, search_svc],
-        outputs=status_output
-    )
+            with gr.Tab(label="🔍 History/Retrieve"):
+                search_svc = gr.Textbox(label="Search Service")
+                get_btn = gr.Button("Decrypt & Show", variant="secondary")
+
+        status_output = gr.Textbox(label="Status", interactive=False)
+        logout_btn = gr.Button("🔒 Lock Vault", variant="stop", size="sm")
+
+    # --- Event Handlers ---
+    
+    # Login Flow
+    login_btn.click(login, [master_input], [login_panel, vault_panel, login_msg])
+    
+    # Action Flow
+    save_btn.click(lambda m, s, p: process_action("Save", m, s, p), [master_input, svc, pwd], status_output)
+    get_btn.click(lambda m, s: process_action("Retrieve", m, s), [master_input, search_svc], status_output)
+    
+    # Logout Flow
+    logout_btn.click(lambda: [gr.update(visible=True), gr.update(visible=False), ""], None, [login_panel, vault_panel, master_input])
 
 if __name__ == "__main__":
-    # Fixed: Removed title from here
     demo.launch(css=custom_css)