Update app.py

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 from cryptography.fernet import Fernet
-from datetime import datetime, timedelta
+from datetime import datetime
 import sqlite3
 import bcrypt
 import uuid
@@ -10,16 +10,18 @@ import logging
 import threading
 import schedule
 import time
-from dataclasses import dataclass
-from typing import Optional, Dict, List
-import json
+from dotenv import load_dotenv
+import os
+
+# Load environment variables
+load_dotenv()
 
 class EmailConfig:
-    SMTP_SERVER = "smtp.gmail.com"
-    SMTP_PORT = 465
-    SMTP_USERNAME = "Mail"  # Replace with your email
-    SMTP_PASSWORD = "lxau wkjx pzhu krqs"     # Replace with your password
-    SENDER_EMAIL = "shaheerawan001@gmail.com"   # Replace with sender email
+    SMTP_SERVER = os.getenv("SMTP_SERVER", "smtp.gmail.com")
+    SMTP_PORT = int(os.getenv("SMTP_PORT", "465"))
+    SMTP_USERNAME = os.getenv("SMTP_USERNAME")
+    SMTP_PASSWORD = os.getenv("SMTP_PASSWORD")
+    SENDER_EMAIL = os.getenv("SENDER_EMAIL")
 
 class Database:
     def __init__(self):
@@ -87,22 +89,14 @@ class DigitalHeir:
     def _check_accounts(self):
         conn = self.db.conn
         cursor = conn.cursor()
-        cursor.execute("""
-            SELECT * FROM social_accounts 
-            WHERE status = 'pending'
-        """)
+        cursor.execute("SELECT * FROM social_accounts WHERE status = 'pending'")
         accounts = cursor.fetchall()
         
         for account in accounts:
             days_inactive = (datetime.now() - datetime.fromisoformat(account[7])).days
-            if days_inactive >= account[6]:  # inactivity_threshold
+            if days_inactive >= account[6]:
                 self._notify_heir(account)
-                
-                cursor.execute("""
-                    UPDATE social_accounts 
-                    SET status = 'inherited' 
-                    WHERE id = ?
-                """, (account[0],))
+                cursor.execute("UPDATE social_accounts SET status = 'inherited' WHERE id = ?", (account[0],))
                 conn.commit()
     
     def _notify_heir(self, account):
@@ -112,12 +106,7 @@ class DigitalHeir:
         
         if heir:
             password = self.decrypt_data(account[4])
-            self._send_email(
-                heir[3],  # heir email
-                account[2],  # platform
-                account[3],  # username
-                password
-            )
+            self._send_email(heir[3], account[2], account[3], password)
     
     def _send_email(self, recipient: str, platform: str, username: str, password: str):
         msg = MIMEText(f"""
@@ -149,17 +138,15 @@ def register_user(username: str, password: str, confirm_password: str):
     if password != confirm_password:
         st.error("Passwords do not match")
         return
-        
+    
     conn = app.db.conn
     cursor = conn.cursor()
     
-    # Check if username exists
     cursor.execute("SELECT * FROM users WHERE username = ?", (username,))
     if cursor.fetchone():
         st.error("Username already exists")
         return
-        
-    # Hash password
+    
     salt = bcrypt.gensalt()
     hashed_password = bcrypt.hashpw(password.encode('utf-8'), salt)
     
@@ -179,29 +166,11 @@ def register_user(username: str, password: str, confirm_password: str):
     except Exception as e:
         st.error(f"Registration failed: {str(e)}")
 
-def show_auth():
-    tab1, tab2 = st.tabs(["Login", "Register"])
-    
-    with tab1:
-        with st.form("login_form"):
-            username = st.text_input("Username")
-            password = st.text_input("Password", type="password")
-            if st.form_submit_button("Login"):
-                login_user(username, password)
-    
-    with tab2:
-        with st.form("register_form"):
-            new_username = st.text_input("Username")
-            new_password = st.text_input("Password", type="password")
-            confirm_password = st.text_input("Confirm Password", type="password")
-            if st.form_submit_button("Register"):
-                register_user(new_username, new_password, confirm_password)
-
 def login_user(username: str, password: str):
     if not username or not password:
         st.error("Username and password are required")
         return
-        
+    
     conn = app.db.conn
     cursor = conn.cursor()
     
@@ -219,18 +188,50 @@ def login_user(username: str, password: str):
     else:
         st.error("Invalid username or password")
 
-def show_account_management():
-    st.title("Digital Asset Management")
+def save_heir(name, email, relationship):
+    if not name or not email:
+        st.error("Name and email are required")
+        return
     
-    with st.form("add_account"):
-        platform = st.selectbox("Platform", ["Facebook", "Twitter", "Instagram", "LinkedIn"])
-        username = st.text_input("Username/Email")
-        password = st.text_input("Password", type="password")
-        heir = st.selectbox("Select Heir", get_user_heirs())
-        inactive_days = st.number_input("Inactivity Period (days)", min_value=30, value=90)
-        
-        if st.form_submit_button("Add Account"):
-            save_social_account(platform, username, password, heir, inactive_days)
+    conn = app.db.conn
+    cursor = conn.cursor()
+    
+    try:
+        cursor.execute("""
+            INSERT INTO heirs (id, user_id, name, email, relationship, created_at)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """, (
+            str(uuid.uuid4()),
+            st.session_state.current_user['id'],
+            name,
+            email,
+            relationship,
+            datetime.now().isoformat()
+        ))
+        conn.commit()
+        st.success("Heir added successfully!")
+        st.rerun()
+    except Exception as e:
+        st.error(f"Failed to add heir: {str(e)}")
+
+def get_user_heirs():
+    conn = app.db.conn
+    cursor = conn.cursor()
+    cursor.execute("""
+        SELECT id, name, email, relationship 
+        FROM heirs 
+        WHERE user_id = ?
+    """, (st.session_state.current_user['id'],))
+    
+    return [
+        {
+            'id': heir[0],
+            'name': heir[1],
+            'email': heir[2],
+            'relationship': heir[3]
+        }
+        for heir in cursor.fetchall()
+    ]
 
 def save_social_account(platform, username, password, heir, inactive_days):
     encrypted_password = app.encrypt_data(password)
@@ -259,25 +260,17 @@ def save_social_account(platform, username, password, heir, inactive_days):
     except Exception as e:
         st.error(f"Failed to save account: {str(e)}")
 
-def get_user_heirs():
-    conn = app.db.conn
-    cursor = conn.cursor()
-    cursor.execute("""
-        SELECT id, name, email, relationship 
-        FROM heirs 
-        WHERE user_id = ?
-    """, (st.session_state.current_user['id'],))
+def show_heir_management():
+    st.sidebar.title("Heir Management")
     
-    # Convert tuple to dictionary
-    heirs = []
-    for heir in cursor.fetchall():
-        heirs.append({
-            'id': heir[0],
-            'name': heir[1],
-            'email': heir[2],
-            'relationship': heir[3]
-        })
-    return heirs
+    with st.sidebar.form("add_heir"):
+        heir_name = st.text_input("Heir Name")
+        heir_email = st.text_input("Heir Email")
+        heir_relationship = st.selectbox("Relationship", 
+            ["Family", "Friend", "Legal Representative", "Other"])
+        
+        if st.form_submit_button("Add Heir"):
+            save_heir(heir_name, heir_email, heir_relationship)
 
 def show_account_management():
     st.title("Digital Asset Management")
@@ -286,7 +279,7 @@ def show_account_management():
     if not heirs:
         st.warning("Please add heirs before adding accounts")
         return
-        
+    
     with st.form("add_account"):
         platform = st.selectbox("Platform", ["Facebook", "Twitter", "Instagram", "LinkedIn"])
         username = st.text_input("Username/Email")
@@ -299,13 +292,65 @@ def show_account_management():
         if st.form_submit_button("Add Account"):
             save_social_account(platform, username, password, heir, inactive_days)
 
+def show_auth():
+    tab1, tab2 = st.tabs(["Login", "Register"])
+    
+    with tab1:
+        with st.form("login_form"):
+            username = st.text_input("Username")
+            password = st.text_input("Password", type="password")
+            if st.form_submit_button("Login"):
+                login_user(username, password)
+    
+    with tab2:
+        with st.form("register_form"):
+            new_username = st.text_input("Username")
+            new_password = st.text_input("Password", type="password")
+            confirm_password = st.text_input("Confirm Password", type="password")
+            if st.form_submit_button("Register"):
+                register_user(new_username, new_password, confirm_password)
+
+def show_heirs_list():
+    heirs = get_user_heirs()
+    
+    if not heirs:
+        st.info("No heirs added yet")
+        return
+    
+    for heir in heirs:
+        with st.expander(f"{heir['name']} - {heir['relationship']}"):
+            st.write(f"Email: {heir['email']}")
+            if st.button("Remove", key=f"remove_{heir['id']}"):
+                remove_heir(heir['id'])
+
+def remove_heir(heir_id):
+    conn = app.db.conn
+    cursor = conn.cursor()
+    
+    try:
+        cursor.execute("DELETE FROM heirs WHERE id = ?", (heir_id,))
+        conn.commit()
+        st.success("Heir removed successfully!")
+        st.rerun()
+    except sqlite3.Error:
+        st.error("Error removing heir")
+
 def main():
     st.set_page_config(page_title="Digital Inheritance System", layout="wide")
     
     if 'current_user' not in st.session_state:
         show_auth()
     else:
-        show_account_management()
+        st.title("Digital Inheritance System")
+        show_heir_management()
+        
+        tab1, tab2 = st.tabs(["Manage Accounts", "View Heirs"])
+        
+        with tab1:
+            show_account_management()
+        
+        with tab2:
+            show_heirs_list()
 
 app = DigitalHeir()