Update app.py

app.py

@@ -1,14 +1,10 @@
 import gradio as gr
 import os
 import time
-import sqlite3
-import secrets
-from pathlib import Path
-from functools import partial
 from helper import (
     create_user, verify_user, get_user_files, upload_file, download_file,
     delete_file, empty_vault, is_admin, get_all_accounts, delete_account,
-    encrypt_file, decrypt_file, create_tables
+    encrypt_file, decrypt_file, is_rate_limited, is_account_locked, record_login_attempt
 )
 
 # Constants
@@ -19,13 +15,6 @@ INACTIVITY_TIMEOUT = 300  # 5 minutes in seconds
 current_user = None
 last_activity_time = time.time()
 
-# Initialize database
-db_path = Path("database/grimvault.db")
-db_path.parent.mkdir(exist_ok=True)
-conn = sqlite3.connect(str(db_path))
-create_tables(conn)
-conn.close()
-
 def check_inactivity():
     global current_user, last_activity_time
     if current_user and time.time() - last_activity_time > INACTIVITY_TIMEOUT:
@@ -42,12 +31,19 @@ def login(username, password):
     if check_inactivity():
         return gr.update(visible=True), gr.update(visible=False), "You have been logged out due to inactivity."
     
-    user = verify_user(username, password)
-    if user:
+    if is_account_locked(username):
+        return gr.update(visible=True), gr.update(visible=False), "Account is locked. Please try again later."
+    
+    if is_rate_limited(username):
+        return gr.update(visible=True), gr.update(visible=False), "Too many login attempts. Please try again later."
+    
+    if verify_user(username, password):
         current_user = username
         update_activity()
+        record_login_attempt(username, True)
         return gr.update(visible=False), gr.update(visible=True), f"Welcome, {username}!"
     else:
+        record_login_attempt(username, False)
         return gr.update(visible=True), gr.update(visible=False), "Invalid username or password."
 
 def logout():
@@ -75,8 +71,8 @@ def upload(files):
         if file.size > MAX_FILE_SIZE:
             results.append(f"File {file.name} exceeds the 5GB limit.")
         else:
-            encrypted_file = encrypt_file(file.name, file.read())
-            result = upload_file(current_user, file.name, encrypted_file)
+            encrypted_content = encrypt_file(file.name, file.read())
+            result = upload_file(current_user, file.name, encrypted_content)
             results.append(result)
     
     update_activity()
@@ -100,9 +96,9 @@ def download(filename):
     if not current_user:
         return None
     
-    file_content = download_file(current_user, filename)
-    if file_content:
-        decrypted_content = decrypt_file(filename, file_content)
+    encrypted_content = download_file(current_user, filename)
+    if encrypted_content:
+        decrypted_content = decrypt_file(filename, encrypted_content)
         update_activity()
         return decrypted_content
     else: