data.py

import streamlit as st
import hashlib
import time
import uuid
import json
from cryptography.fernet import Fernet
from datetime import datetime

# Set page configuration
st.set_page_config(
    page_title="VaultX - Secure Data Locker",
    page_icon="🔒",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for a unique UI
def apply_custom_css():
    st.markdown("""
    <style>
        .main-header {
            font-size: 2.5rem;
            font-weight: 700;
            color: #1E3A8A;
            text-align: center;
            margin-bottom: 2rem;
            text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.1);
        }
        .subheader {
            font-size: 1.5rem;
            font-weight: 600;
            color: #3B82F6;
            margin-bottom: 1rem;
        }
        .card {
            background-color: #F3F4F6;
            border-radius: 10px;
            padding: 20px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
            margin-bottom: 20px;
        }
        .info-box {
            background-color: #DBEAFE;
            border-left: 5px solid #3B82F6;
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        .success-box {
            background-color: #D1FAE5;
            border-left: 5px solid #10B981;
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        .warning-box {
            background-color: #FEF3C7;
            border-left: 5px solid #F59E0B;
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        .error-box {
            background-color: #FEE2E2;
            border-left: 5px solid #EF4444;
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        .btn-primary {
            background-color: #3B82F6;
            color: white;
            padding: 10px 20px;
            border-radius: 5px;
            border: none;
            cursor: pointer;
            font-weight: 600;
            transition: all 0.3s ease;
        }
        .btn-primary:hover {
            background-color: #2563EB;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        }
        .sidebar .css-1d391kg {
            background-color: #1E3A8A;
        }
        .sidebar .css-1aumxhk {
            color: white;
        }
        .stButton > button {
            width: 100%;
            background-color: #2563EB;
            color: white;
            font-weight: 600;
            border: none;
            border-radius: 5px;
            padding: 10px 0;
            transition: all 0.3s ease;
        }
        .stButton > button:hover {
            background-color: #1D4ED8;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        }
        .statbox {
            background-color: #F3F4F6;
            border-radius: 10px;
            padding: 15px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
            text-align: center;
        }
        .stat-value {
            font-size: 2rem;
            font-weight: 700;
            color: #1E3A8A;
        }
        .stat-label {
            font-size: 0.9rem;
            color: #6B7280;
        }
    </style>
    """, unsafe_allow_html=True)

# Initialize session state variables if they don't exist
if 'key' not in st.session_state:
    st.session_state.key = Fernet.generate_key()
    st.session_state.cipher = Fernet(st.session_state.key)

if 'stored_data' not in st.session_state:
    st.session_state.stored_data = {}

if 'failed_attempts' not in st.session_state:
    st.session_state.failed_attempts = 0

if 'locked_until' not in st.session_state:
    st.session_state.locked_until = None

if 'authenticated' not in st.session_state:
    st.session_state.authenticated = True

if 'last_activity' not in st.session_state:
    st.session_state.last_activity = time.time()

if 'data_count' not in st.session_state:
    st.session_state.data_count = 0

if 'last_access' not in st.session_state:
    st.session_state.last_access = "Never"

# Activity tracker
def update_activity():
    st.session_state.last_activity = time.time()

# Session timeout (5 minutes)
def check_session_timeout():
    if time.time() - st.session_state.last_activity > 300:  # 5 minutes
        st.session_state.authenticated = False
        return True
    return False

# Function to hash passkey
def hash_passkey(passkey):
    return hashlib.sha256(passkey.encode()).hexdigest()

# Function to encrypt data
def encrypt_data(text, passkey):
    # We use the same key for all encryption but store data with a unique ID
    return st.session_state.cipher.encrypt(text.encode()).decode()

# Function to decrypt data
def decrypt_data(encrypted_text, passkey):
    hashed_passkey = hash_passkey(passkey)
    
    # Check if we're locked out
    if st.session_state.locked_until and time.time() < st.session_state.locked_until:
        remaining_time = int(st.session_state.locked_until - time.time())
        st.error(f"⚠️ Account locked. Try again in {remaining_time} seconds.")
        return None
    
    # Search for the entry with matching encrypted text and passkey
    for key, value in st.session_state.stored_data.items():
        if value["encrypted_text"] == encrypted_text and value["passkey"] == hashed_passkey:
            st.session_state.failed_attempts = 0
            st.session_state.last_access = datetime.now().strftime("%H:%M:%S")
            update_activity()
            return st.session_state.cipher.decrypt(encrypted_text.encode()).decode()
    
    # Increment failed attempts if decryption fails
    st.session_state.failed_attempts += 1
    
    # Lock out after 3 failed attempts
    if st.session_state.failed_attempts >= 3:
        st.session_state.locked_until = time.time() + 30  # Lock for 30 seconds
        st.session_state.authenticated = False
    
    return None

# Function to export data to JSON
def export_data():
    try:
        with open("vault_data.json", "w") as f:
            json.dump(st.session_state.stored_data, f)
        return True
    except Exception as e:
        st.error(f"Failed to export data: {e}")
        return False

# Function to import data from JSON
def import_data():
    try:
        with open("vault_data.json", "r") as f:
            st.session_state.stored_data = json.load(f)
        st.session_state.data_count = len(st.session_state.stored_data)
        return True
    except FileNotFoundError:
        st.warning("No exported data found.")
        return False
    except Exception as e:
        st.error(f"Failed to import data: {e}")
        return False

# Apply custom CSS
apply_custom_css()

# Sidebar navigation
st.sidebar.markdown("<h1 style='text-align: center; color: white;'>🔒 VaultX</h1>", unsafe_allow_html=True)
st.sidebar.markdown("<p style='text-align: center; color: #CBD5E1;'>Secure Data Locker</p>", unsafe_allow_html=True)

# Check session timeout
if check_session_timeout():
    st.warning("⚠️ Session timed out due to inactivity. Please authenticate again.")
    st.session_state.choice = "Login"

# Display user stats in sidebar
st.sidebar.markdown("---")
st.sidebar.markdown("<p style='text-align: center; color: #CBD5E1;'>User Statistics</p>", unsafe_allow_html=True)

col1, col2 = st.sidebar.columns(2)
with col1:
    st.markdown(f"""
    <div class='statbox'>
        <div class='stat-value'>{len(st.session_state.stored_data)}</div>
        <div class='stat-label'>Entries</div>
    </div>
    """, unsafe_allow_html=True)

with col2:
    st.markdown(f"""
    <div class='statbox'>
        <div class='stat-value'>{st.session_state.failed_attempts}</div>
        <div class='stat-label'>Failed Attempts</div>
    </div>
    """, unsafe_allow_html=True)

st.sidebar.markdown(f"""
<div style='text-align: center; color: #CBD5E1; font-size: 0.8rem; margin-top: 10px;'>
    Last Access: {st.session_state.last_access}
</div>
""", unsafe_allow_html=True)

# Navigation options
menu = ["Home", "Store Data", "Retrieve Data", "Data Management", "Login"]
st.sidebar.markdown("---")
choice = st.sidebar.radio("Navigation", menu)

# Lock check
if st.session_state.locked_until and time.time() < st.session_state.locked_until:
    remaining_time = int(st.session_state.locked_until - time.time())
    st.error(f"⚠️ Account locked due to too many failed attempts. Try again in {remaining_time} seconds.")
    choice = "Login"

# Authentication check
if not st.session_state.authenticated and choice != "Login":
    st.warning("⚠️ Please authenticate first.")
    choice = "Login"

# Main content
if choice == "Home":
    update_activity()
    st.markdown("<h1 class='main-header'>🏠 Welcome to VaultX - Secure Data Locker</h1>", unsafe_allow_html=True)
    
    # Split layout for features
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>🔐 Store Securely</h2>", unsafe_allow_html=True)
        st.markdown("""
        <div class='info-box'>
            Store your sensitive data with strong encryption. 
            Each piece of data is protected with a unique passkey.
        </div>
        """, unsafe_allow_html=True)
        if st.button("Go to Store Data"):
            st.session_state.choice = "Store Data"
            st.experimental_rerun()
        st.markdown("</div>", unsafe_allow_html=True)
        
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>🛡️ Security Features</h2>", unsafe_allow_html=True)
        st.markdown("""
        <ul>
            <li>Fernet symmetric encryption</li>
            <li>SHA-256 passkey hashing</li>
            <li>Automatic lockout after failed attempts</li>
            <li>Session timeout protection</li>
        </ul>
        """, unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)
    
    with col2:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>🔍 Retrieve Data</h2>", unsafe_allow_html=True)
        st.markdown("""
        <div class='info-box'>
            Access your stored data by providing the correct passkey.
            Only you can decrypt your sensitive information.
        </div>
        """, unsafe_allow_html=True)
        if st.button("Go to Retrieve Data"):
            st.session_state.choice = "Retrieve Data"
            st.experimental_rerun()
        st.markdown("</div>", unsafe_allow_html=True)
        
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>📊 Data Management</h2>", unsafe_allow_html=True)
        st.markdown("""
        <div class='info-box'>
            Export your encrypted data or import previously exported data.
            Manage all your secure information in one place.
        </div>
        """, unsafe_allow_html=True)
        if st.button("Go to Data Management"):
            st.session_state.choice = "Data Management"
            st.experimental_rerun()
        st.markdown("</div>", unsafe_allow_html=True)

elif choice == "Store Data":
    update_activity()
    st.markdown("<h1 class='main-header'>📂 Store Your Data Securely</h1>", unsafe_allow_html=True)
    
    with st.form("store_data_form"):
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        data_name = st.text_input("Data Name (Optional):", placeholder="e.g., 'My Credit Card', 'Website Password'")
        user_data = st.text_area("Enter Your Data:", placeholder="Enter the sensitive information you want to encrypt")
        passkey = st.text_input("Enter Passkey:", type="password", placeholder="Create a strong, unique passkey")
        confirm_passkey = st.text_input("Confirm Passkey:", type="password", placeholder="Confirm your passkey")
        st.markdown("</div>", unsafe_allow_html=True)
        
        st.markdown("""
        <div class='info-box'>
            <strong>Remember:</strong> Your passkey is the only way to retrieve your data later.
            Make sure to remember it or store it somewhere safe!
        </div>
        """, unsafe_allow_html=True)
        
        submitted = st.form_submit_button("Encrypt & Save")
        
    if submitted:
        if not user_data:
            st.markdown("<div class='error-box'>⚠️ Please enter data to encrypt.</div>", unsafe_allow_html=True)
        elif not passkey:
            st.markdown("<div class='error-box'>⚠️ Please enter a passkey.</div>", unsafe_allow_html=True)
        elif passkey != confirm_passkey:
            st.markdown("<div class='error-box'>⚠️ Passkeys do not match.</div>", unsafe_allow_html=True)
        else:
            hashed_passkey = hash_passkey(passkey)
            encrypted_text = encrypt_data(user_data, passkey)
            
            # Generate a unique ID for this data
            data_id = str(uuid.uuid4())
            
            # Store the data with optional name
            name_to_use = data_name if data_name else f"Data Entry {len(st.session_state.stored_data) + 1}"
            
            st.session_state.stored_data[data_id] = {
                "name": name_to_use,
                "encrypted_text": encrypted_text,
                "passkey": hashed_passkey,
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            
            st.session_state.data_count = len(st.session_state.stored_data)
            
            st.markdown("<div class='success-box'>✅ Data stored securely!</div>", unsafe_allow_html=True)
            
            # Show the encrypted data
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            st.markdown("<h3>Your Encrypted Data</h3>", unsafe_allow_html=True)
            st.code(encrypted_text, language="text")
            st.markdown("</div>", unsafe_allow_html=True)
            
            st.markdown("""
            <div class='warning-box'>
                <strong>Important:</strong> Copy and keep this encrypted text if you want to retrieve your data later.
                You'll need both this text and your passkey to decrypt it.
            </div>
            """, unsafe_allow_html=True)

elif choice == "Retrieve Data":
    update_activity()
    st.markdown("<h1 class='main-header'>🔍 Retrieve Your Data</h1>", unsafe_allow_html=True)
    
    # Show available data entries if any
    if st.session_state.stored_data:
        st.markdown("<div class='info-box'>", unsafe_allow_html=True)
        st.markdown("### Available Data Entries")
        st.markdown("Select from your stored data entries below:")
        st.markdown("</div>", unsafe_allow_html=True)
        
        data_options = ["Select an entry..."] + [f"{value['name']} ({key[:8]}...)" for key, value in st.session_state.stored_data.items()]
        selected_data = st.selectbox("", data_options)
        
        if selected_data != "Select an entry...":
            data_id = selected_data.split("(")[1].split(")")[0].replace("...", "")
            for key, value in st.session_state.stored_data.items():
                if key.startswith(data_id):
                    selected_key = key
                    selected_value = value
                    break
            
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            st.markdown(f"<h3>{selected_value['name']}</h3>", unsafe_allow_html=True)
            st.markdown(f"Created: {selected_value['timestamp']}", unsafe_allow_html=True)
            st.markdown("</div>", unsafe_allow_html=True)
            
            # Set up decrypt form
            with st.form("decrypt_form"):
                st.markdown("<div class='card'>", unsafe_allow_html=True)
                encrypted_text = selected_value['encrypted_text']
                st.markdown("### Encrypted Data")
                st.code(encrypted_text[:50] + "..." if len(encrypted_text) > 50 else encrypted_text, language="text")
                
                passkey = st.text_input("Enter Passkey:", type="password", placeholder="Enter your passkey to decrypt")
                submit_decrypt = st.form_submit_button("Decrypt")
                st.markdown("</div>", unsafe_allow_html=True)
            
            if submit_decrypt:
                if passkey:
                    decrypted_text = decrypt_data(encrypted_text, passkey)
                    
                    if decrypted_text:
                        st.markdown("<div class='success-box'>✅ Data decrypted successfully!</div>", unsafe_allow_html=True)
                        
                        st.markdown("<div class='card'>", unsafe_allow_html=True)
                        st.markdown("### Decrypted Data")
                        st.markdown("```")
                        st.markdown(decrypted_text)
                        st.markdown("```")
                        st.markdown("</div>", unsafe_allow_html=True)
                    else:
                        st.markdown("<div class='error-box'>❌ Incorrect passkey! "
                                   f"Attempts remaining: {3 - st.session_state.failed_attempts}</div>", 
                                   unsafe_allow_html=True)
                        
                        if st.session_state.failed_attempts >= 3:
                            st.markdown("<div class='warning-box'>🔒 Too many failed attempts! "
                                       "Redirecting to Login Page.</div>", 
                                       unsafe_allow_html=True)
                            st.session_state.authenticated = False
                            st.experimental_rerun()
                else:
                    st.markdown("<div class='error-box'>⚠️ Please enter a passkey.</div>", unsafe_allow_html=True)
    else:
        # Or manual decrypt option
        st.markdown("<div class='warning-box'>No data entries found. You can decrypt using the encrypted text.</div>", unsafe_allow_html=True)
        
        with st.form("manual_decrypt_form"):
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            encrypted_text = st.text_area("Enter Encrypted Data:", 
                                         placeholder="Paste the encrypted text you want to decrypt")
            passkey = st.text_input("Enter Passkey:", type="password", 
                                   placeholder="Enter the passkey for this data")
            submit_decrypt = st.form_submit_button("Decrypt")
            st.markdown("</div>", unsafe_allow_html=True)
        
        if submit_decrypt:
            if encrypted_text and passkey:
                decrypted_text = decrypt_data(encrypted_text, passkey)
                
                if decrypted_text:
                    st.markdown("<div class='success-box'>✅ Data decrypted successfully!</div>", unsafe_allow_html=True)
                    
                    st.markdown("<div class='card'>", unsafe_allow_html=True)
                    st.markdown("### Decrypted Data")
                    st.markdown("```")
                    st.markdown(decrypted_text)
                    st.markdown("```")
                    st.markdown("</div>", unsafe_allow_html=True)
                else:
                    st.markdown("<div class='error-box'>❌ Incorrect passkey! "
                               f"Attempts remaining: {3 - st.session_state.failed_attempts}</div>", 
                               unsafe_allow_html=True)
                    
                    if st.session_state.failed_attempts >= 3:
                        st.markdown("<div class='warning-box'>🔒 Too many failed attempts! "
                                   "Redirecting to Login Page.</div>", 
                                   unsafe_allow_html=True)
                        st.session_state.authenticated = False
                        st.experimental_rerun()
            else:
                st.markdown("<div class='error-box'>⚠️ Both fields are required!</div>", unsafe_allow_html=True)

elif choice == "Data Management":
    update_activity()
    st.markdown("<h1 class='main-header'>📊 Data Management</h1>", unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>📤 Export Data</h2>", unsafe_allow_html=True)
        st.markdown("""
        <div class='info-box'>
            Export your encrypted data to a JSON file for backup or transfer.
        </div>
        """, unsafe_allow_html=True)
        
        if st.button("Export Data"):
            if len(st.session_state.stored_data) > 0:
                if export_data():
                    st.markdown("<div class='success-box'>✅ Data exported successfully!</div>", 
                               unsafe_allow_html=True)
            else:
                st.markdown("<div class='warning-box'>⚠️ No data to export.</div>", unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)
    
    with col2:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("<h2 class='subheader'>📥 Import Data</h2>", unsafe_allow_html=True)
        st.markdown("""
        <div class='info-box'>
            Import previously exported data from a JSON file.
        </div>
        """, unsafe_allow_html=True)
        
        if st.button("Import Data"):
            if import_data():
                st.markdown("<div class='success-box'>✅ Data imported successfully!</div>", 
                           unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)
    
    # Data overview
    st.markdown("<div class='card'>", unsafe_allow_html=True)
    st.markdown("<h2 class='subheader'>📋 Data Overview</h2>", unsafe_allow_html=True)
    
    if len(st.session_state.stored_data) > 0:
        for key, value in st.session_state.stored_data.items():
            col1, col2, col3 = st.columns([3, 2, 1])
            with col1:
                st.markdown(f"**{value['name']}**")
            with col2:
                st.markdown(f"Created: {value.get('timestamp', 'Unknown')}")
            with col3:
                if st.button(f"Delete", key=f"delete_{key}"):
                    del st.session_state.stored_data[key]
                    st.session_state.data_count = len(st.session_state.stored_data)
                    st.experimental_rerun()
        
        if st.button("Clear All Data"):
            st.session_state.stored_data = {}
            st.session_state.data_count = 0
            st.experimental_rerun()
    else:
        st.markdown("<div class='warning-box'>No data entries found.</div>", unsafe_allow_html=True)
    
    st.markdown("</div>", unsafe_allow_html=True)

elif choice == "Login":
    st.markdown("<h1 class='main-header'>🔑 Authentication Required</h1>", unsafe_allow_html=True)
    
    # Check if we're locked out
    if st.session_state.locked_until and time.time() < st.session_state.locked_until:
        remaining_time = int(st.session_state.locked_until - time.time())
        st.markdown(f"<div class='error-box'>⚠️ Account locked due to too many failed attempts. "
                   f"Try again in {remaining_time} seconds.</div>", 
                   unsafe_allow_html=True)
    else:
        with st.form("login_form"):
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            login_pass = st.text_input("Enter Master Password:", type="password", 
                                      placeholder="Enter master password to continue")
            submit_login = st.form_submit_button("Login")
            st.markdown("</div>", unsafe_allow_html=True)
            
            st.markdown("<div class='info-box'>", unsafe_allow_html=True)
            st.markdown("**Note:** For demo purposes, the master password is 'admin123'.")
            st.markdown("</div>", unsafe_allow_html=True)
        
        if submit_login:
            if login_pass == "admin123":  # Hardcoded for demo, replace with proper auth
                st.session_state.failed_attempts = 0
                st.session_state.authenticated = True
                st.session_state.locked_until = None
                st.session_state.last_activity = time.time()
                
                st.markdown("<div class='success-box'>✅ Authentication successful!</div>", 
                           unsafe_allow_html=True)
                
                # Redirect to home page
                st.markdown("<meta http-equiv='refresh' content='2; url=/'/>", unsafe_allow_html=True)
            else:
                st.markdown("<div class='error-box'>❌ Incorrect password!</div>", unsafe_allow_html=True)

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #6B7280; font-size: 0.8rem;'>
    VaultX - Secure Data Locker | Developed with ❤️ by Muhammad Yousaf
    <br>© 2025 VaultX. All rights reserved.
</div>
""", unsafe_allow_html=True)