Update app.py

app.py

@@ -12,119 +12,93 @@ import gradio as gr
 import os
 import numpy as np
 import librosa
+import base64
 from resemblyzer import VoiceEncoder
 from cryptography.fernet import Fernet
 
-# Initialize voice encoder
+# Initialize Voice Encoder
 encoder = VoiceEncoder()
 
-# Generate or load encryption key
-KEY_FILE = "encryption_key.key"
-if os.path.exists(KEY_FILE):
-    with open(KEY_FILE, "rb") as key_file:
-        encryption_key = key_file.read()
-else:
-    encryption_key = Fernet.generate_key()
-    with open(KEY_FILE, "wb") as key_file:
-        key_file.write(encryption_key)
+# Directory for storing encrypted documents
+encrypted_dir = "encrypted_docs"
+os.makedirs(encrypted_dir, exist_ok=True)
 
-cipher = Fernet(encryption_key)
+# Store user voice embeddings
+db = {}
 
-# Folder to store user embeddings and encrypted documents
-os.makedirs("embeddings", exist_ok=True)
-os.makedirs("encrypted_docs", exist_ok=True)
-
-def load_audio(file, target_sr=16000):
-    """Loads an audio file and resamples it."""
-    audio, sr = librosa.load(file, sr=target_sr)
-    return audio
-
-def extract_embeddings(user_id, files):
-    """Extracts and stores voice embeddings for a user."""
-    embeddings = []
-    for file in files:
-        audio = load_audio(file.name)
-        embedding = encoder.embed_utterance(audio)
-        embeddings.append(embedding)
-
-    np.save(f"embeddings/{user_id}.npy", np.array(embeddings))
+def register_voice(user_id, files):
+    global db
+    embeddings = [encoder.embed_utterance(librosa.load(file.name, sr=16000)[0]) for file in files]
+    db[user_id] = np.mean(embeddings, axis=0)  # Store the mean embedding
     return f"Voice registered successfully for User ID: {user_id}"
 
-def authenticate_user(user_id, test_file):
-    """Authenticates a user based on voice similarity."""
-    if not os.path.exists(f"embeddings/{user_id}.npy"):
-        return False, "No voice data found. Please register first."
-
-    stored_embeddings = np.load(f"embeddings/{user_id}.npy")
-    test_audio = load_audio(test_file.name)
-    test_embedding = encoder.embed_utterance(test_audio)
-
-    similarities = np.dot(stored_embeddings, test_embedding) / (
-        np.linalg.norm(stored_embeddings, axis=1) * np.linalg.norm(test_embedding)
-    )
-    similarity_score = np.mean(similarities)
-
-    if similarity_score > 0.85:
-        return True, f"✅ Access Granted! Similarity Score: {similarity_score:.2f}"
-    return False, f"❌ Access Denied! Similarity Score: {similarity_score:.2f}"
+# Generate an encryption key (For simplicity, using a static key. In production, use per-user keys.)
+key = Fernet.generate_key()
+cipher = Fernet(key)
 
 def encrypt_document(user_id, file):
-    """Encrypts a document and saves it securely."""
-    encrypted_path = f"encrypted_docs/{user_id}_{os.path.basename(file.name)}.enc"
+    if user_id not in db:
+        return "Error: Register your voice first."
     
+    file_path = os.path.join(encrypted_dir, f"{user_id}_{file.name}.enc")
     with open(file.name, "rb") as f:
         encrypted_data = cipher.encrypt(f.read())
-
-    with open(encrypted_path, "wb") as f:
+    with open(file_path, "wb") as f:
         f.write(encrypted_data)
+    
+    return f"Document encrypted and stored securely as {file_path}"
 
-    return f"Document encrypted and stored securely as {encrypted_path}"
-
-def decrypt_document(user_id, test_file, encrypted_file):
-    """Decrypts a document if the user is authenticated via voice."""
-    authenticated, message = authenticate_user(user_id, test_file)
-
-    if authenticated:
-        with open(encrypted_file.name, "rb") as f:
-            decrypted_data = cipher.decrypt(f.read())
-
-        decrypted_path = encrypted_file.name.replace(".enc", "_decrypted")
-        with open(decrypted_path, "wb") as f:
-            f.write(decrypted_data)
-
-        return message, decrypted_path
-    return message, None
+def authenticate_and_decrypt(user_id, voice_file):
+    if user_id not in db:
+        return "Error: No registered voice found. Please register your voice first."
+    
+    # Load and process test voice
+    test_audio = librosa.load(voice_file.name, sr=16000)[0]
+    test_embedding = encoder.embed_utterance(test_audio)
+    similarity = np.dot(db[user_id], test_embedding) / (np.linalg.norm(db[user_id]) * np.linalg.norm(test_embedding))
+    
+    if similarity < 0.8:
+        return "Authentication failed: Voice does not match."
+    
+    # Locate encrypted file
+    user_files = [f for f in os.listdir(encrypted_dir) if f.startswith(user_id)]
+    if not user_files:
+        return "No encrypted document found for this user."
+    
+    file_path = os.path.join(encrypted_dir, user_files[0])
+    with open(file_path, "rb") as f:
+        decrypted_data = cipher.decrypt(f.read())
+    
+    # Save decrypted file temporarily
+    decrypted_file = "decrypted_document.pdf"
+    with open(decrypted_file, "wb") as f:
+        f.write(decrypted_data)
+    
+    return decrypted_file  # Returning file path for download
 
 # Gradio UI
 with gr.Blocks() as app:
-    gr.Markdown("## 🔒 Secure Document System with Voice Authentication")
-
-    # Registration
+    gr.Markdown("# 🔒 Secure Document System with Voice Authentication")
+    
     with gr.Row():
-        user_id_input = gr.Textbox(label="User ID")
-        train_audio = gr.File(label="Upload Voice Samples", file_types=[".wav"], file_count="multiple")
-        train_button = gr.Button("Register Voice")
-    train_output = gr.Textbox()
-    train_button.click(extract_embeddings, inputs=[user_id_input, train_audio], outputs=train_output)
-
-    # Encryption
+        user_id_reg = gr.Textbox(label="User ID")
+        voice_samples = gr.File(label="Upload Voice Samples", file_types=[".wav"], file_count="multiple")
+        reg_button = gr.Button("Register Voice")
+    reg_output = gr.Textbox()
+    reg_button.click(register_voice, inputs=[user_id_reg, voice_samples], outputs=reg_output)
+    
     with gr.Row():
-        doc_user_id = gr.Textbox(label="User ID for Encryption")
-        document_upload = gr.File(label="Upload Document", file_types=[".pdf", ".txt", ".docx"])
+        user_id_enc = gr.Textbox(label="User ID for Encryption")
+        document = gr.File(label="Upload Document", file_types=[".pdf"])  
         encrypt_button = gr.Button("Encrypt Document")
-    encrypt_output = gr.Textbox()
-    encrypt_button.click(encrypt_document, inputs=[doc_user_id, document_upload], outputs=encrypt_output)
-
-    # Authentication & Decryption
+    enc_output = gr.Textbox()
+    encrypt_button.click(encrypt_document, inputs=[user_id_enc, document], outputs=enc_output)
+    
     with gr.Row():
-        auth_user_id = gr.Textbox(label="User ID for Authentication")
-        test_audio = gr.File(label="Upload Voice for Authentication", file_types=[".wav"])
-        encrypted_doc = gr.File(label="Upload Encrypted Document", file_types=[".enc"])
+        user_id_auth = gr.Textbox(label="User ID for Authentication")
+        auth_voice = gr.File(label="Upload Voice for Authentication", file_types=[".wav"])
         decrypt_button = gr.Button("Decrypt Document")
-    decrypt_output = gr.Textbox()
-    decrypted_doc_output = gr.File(label="Decrypted Document")
-    decrypt_button.click(decrypt_document, inputs=[auth_user_id, test_audio, encrypted_doc], outputs=[decrypt_output, decrypted_doc_output])
-
+    dec_output = gr.File(label="Decrypted Document")
+    decrypt_button.click(authenticate_and_decrypt, inputs=[user_id_auth, auth_voice], outputs=dec_output)
+    
 app.launch(share=True)
-
-