app.py

import streamlit as st
import hashlib
import os
import base64
import json
from cryptography.fernet import Fernet

# --- File Storage ---
DATA_FILE = "data.json"
LOGINS_FILE = "logins.json"

# --- Initialize session state ---
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False
if "menu" not in st.session_state:
    st.session_state.menu = ["Home", "Sign Up", "Login"]

# --- Load or create JSON data ---
def load_json(file):
    if os.path.exists(file):
        with open(file, "r") as f:
            return json.load(f)
    return {}

def save_json(data, file):
    with open(file, "w") as f:
        json.dump(data, f)

# --- Password Hashing for Signup/Login ---
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

# --- Key Creation from Passkey ---
def get_key_from_passkey(passkey, salt=None):
    if not salt:
        salt = os.urandom(16)
    key = hashlib.pbkdf2_hmac("sha256", passkey.encode(), salt, 100000, dklen=32)
    return base64.urlsafe_b64encode(key), base64.b64encode(salt).decode()

# --- Encryption ---
def encrypt_message(message, passkey):
    key, salt = get_key_from_passkey(passkey)
    f = Fernet(key)
    encrypted = f.encrypt(message.encode()).decode()
    return encrypted, salt

# --- Decryption ---
def decrypt_message(encrypted, passkey, salt):
    try:
        salt_bytes = base64.b64decode(salt)
        key, _ = get_key_from_passkey(passkey, salt_bytes)
        f = Fernet(key)
        return f.decrypt(encrypted.encode()).decode()
    except:
        return None

# --- Load Data ---
stored_data = load_json(DATA_FILE)
logins = load_json(LOGINS_FILE)

# --- Menu Logic ---
def update_menu():
    if st.session_state.logged_in:
        st.session_state.menu = ["Home", "Store Secret", "Get Secret", "Sign Up", "Login"]
    else:
        st.session_state.menu = ["Home", "Sign Up", "Login"]

# --- Streamlit App UI ---
st.title("🔐 My Secret Keeper")

# Update menu according to login state
update_menu()
choice = st.sidebar.selectbox("Menu", st.session_state.menu)

# --- Home Page ---
if choice == "Home":
    st.subheader("Welcome!")
    st.write("Save your secret message using a secret key. Only you can unlock it!")

# --- Store Secret Page ---
elif choice == "Store Secret":
    if not st.session_state.logged_in:
        st.warning("Please log in to store a secret.")
    else:
        st.subheader("🔒 Store Your Secret")
        msg = st.text_area("Your secret:")
        passkey = st.text_input("Passkey:", type="password")
        
        if st.button("Encrypt & Save"):
            if msg and passkey:
                encrypted, salt = encrypt_message(msg, passkey)
                stored_data[encrypted] = salt
                save_json(stored_data, DATA_FILE)
                st.success("✅ Secret saved!")
                st.code(encrypted)
            else:
                st.error("Please enter both secret and passkey.")

# --- Get Secret Page ---
elif choice == "Get Secret":
    if not st.session_state.logged_in:
        st.warning("Please log in to retrieve a secret.")
    else:
        st.subheader("🔓 Get Your Secret")
        encrypted = st.text_area("Paste your encrypted text:")
        passkey = st.text_input("Passkey:", type="password")

        if st.button("Decrypt"):
            salt = stored_data.get(encrypted)
            if salt:
                result = decrypt_message(encrypted, passkey, salt)
                if result:
                    st.success("✅ Here's your secret:")
                    st.code(result)
                else:
                    st.error("❌ Wrong passkey!")
            else:
                st.error("Encrypted text not found.")

# --- Sign Up Page ---
elif choice == "Sign Up":
    st.subheader("👤 Create Account")
    email = st.text_input("Email:")
    password = st.text_input("Password:", type="password")

    if st.button("Sign Up"):
        if email in logins:
            st.error("Email already exists.")
        else:
            logins[email] = hash_password(password)
            save_json(logins, LOGINS_FILE)
            st.success("Account created!")
            st.session_state.logged_in = True
            update_menu()

# --- Login Page ---
elif choice == "Login":
    st.subheader("🔑 Login")
    email = st.text_input("Email:")
    password = st.text_input("Password:", type="password")

    if st.button("Login"):
        if logins.get(email) == hash_password(password):
            st.success("Logged in!")
            st.session_state.logged_in = True
            update_menu()
        else:
            st.error("Wrong email or password.")