Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,5 @@
-# simple_lookup_secure.py
+# simple_lookup_2fa.py
+from __future__ import annotations
 import json
 from pathlib import Path
 from io import BytesIO
@@ -12,14 +13,14 @@ from argon2 import PasswordHasher
 from argon2.exceptions import VerifyMismatchError
 
 st.set_page_config(page_title="Simple Password Lookup (2FA)", page_icon="🔐", layout="centered")
-st.title("🔐 Simple Password Lookup — protected with 2FA")
+st.title("🔐 Simple Password Lookup — 2FA protected")
 
-# ---------- persistent locations ----------
+# ---------------- Persistence ----------------
 PERSIST_DIR  = Path("/data") if Path("/data").exists() else Path(".")
-PERSIST_FILE = PERSIST_DIR / "creds.xlsx"      # your saved credentials file (uploaded once)
-CONFIG_FILE  = PERSIST_DIR / "auth.json"       # stores password hash + TOTP secret
+PERSIST_FILE = PERSIST_DIR / "creds.xlsx"      # saved credentials file (uploaded once)
+CONFIG_FILE  = PERSIST_DIR / "auth.json"       # stores Argon2 password hash + TOTP secret + label
 
-# ---------- helpers ----------
+# ---------------- Globals / helpers ----------------
 ph = PasswordHasher()
 ALIASES = {
     "name": ["name", "title", "site", "account", "platform", "service", "app"],
@@ -31,6 +32,7 @@ ALIASES = {
 EXPECTED = ["name", "username", "url", "password", "note"]
 
 def standardize_columns(df: pd.DataFrame) -> pd.DataFrame:
+    """Map common header aliases -> EXPECTED, create missing columns, clean strings."""
     df.columns = [str(c).strip().lower() for c in df.columns]
     colmap = {}
     for target, alias_list in ALIASES.items():
@@ -45,6 +47,7 @@ def standardize_columns(df: pd.DataFrame) -> pd.DataFrame:
     out = df[[colmap[c] for c in EXPECTED]].rename(columns={colmap[c]: c for c in colmap})
     for c in EXPECTED:
         out[c] = out[c].astype(str).fillna("").replace("nan", "")
+    # keep rows with at least one meaningful field
     key_cols = ["name", "username", "url", "password"]
     mask = out[key_cols].applymap(lambda x: str(x).strip() != "").any(axis=1)
     return out[mask].reset_index(drop=True)
@@ -70,60 +73,91 @@ def make_qr_png(data: str) -> bytes:
     img.save(buf, format="PNG")
     return buf.getvalue()
 
-# ---------- auth gate (setup or login) ----------
+# ---------------- Session keys ----------------
 if "authed" not in st.session_state:
     st.session_state.authed = False
 if "failures" not in st.session_state:
     st.session_state.failures = 0
+# setup state (used only during first-time 2FA setup)
+if "setup_secret" not in st.session_state:
+    st.session_state.setup_secret = None
+if "setup_label" not in st.session_state:
+    st.session_state.setup_label = None
+if "setup_pw" not in st.session_state:
+    st.session_state.setup_pw = None
+
+# ---------------- Step 1: Upload (only if not saved yet) ----------------
+if not PERSIST_FILE.exists():
+    st.subheader("Step 1 — Upload your Excel/CSV (only once)")
+    up = st.file_uploader("Choose file", type=["xlsx", "xls", "csv"], accept_multiple_files=False)
+    if not up:
+        st.stop()
+    try:
+        df = standardize_columns(read_any(up))
+        PERSIST_FILE.parent.mkdir(parents=True, exist_ok=True)
+        df.to_excel(PERSIST_FILE, index=False)
+        st.success(f"Saved credentials to {PERSIST_FILE}.")
+        st.info("Continue below to set up 2FA.")
+    except Exception as e:
+        st.error(f"Failed to read/save file: {e}")
+        st.stop()
 
+# ---------------- Step 2: 2FA setup (first time) ----------------
 cfg = load_config()
-
 if not cfg:
-    st.subheader("First-time setup")
-    st.write("Create an admin password and pair a TOTP app (Google Authenticator, Microsoft Authenticator, 1Password, etc.).")
-
-    with st.form("setup"):
-        admin_user = st.text_input("Account label (for your authenticator app)", value="SimpleLookup", help="Shown in your authenticator app; can be an email or name.")
-        new_pw = st.text_input("New admin password", type="password")
-        new_pw2 = st.text_input("Confirm password", type="password")
-        submitted = st.form_submit_button("Generate TOTP and Set Password")
-
-    if submitted:
-        if not new_pw or new_pw != new_pw2:
+    st.subheader("Step 2 — Set up 2FA (one-time)")
+    st.write("Create an admin password and pair a TOTP app (Google/Microsoft Authenticator).")
+
+    # Collect label + password; generate secret ONCE and keep it in session
+    with st.form("setup_form"):
+        label = st.text_input("Account label (as shown in your authenticator)", value=st.session_state.setup_label or "SimpleLookup")
+        pw1 = st.text_input("New admin password", type="password", value=st.session_state.setup_pw or "")
+        pw2 = st.text_input("Confirm password", type="password", value=st.session_state.setup_pw or "")
+        gen = st.form_submit_button("Generate QR code")
+    if gen:
+        if not pw1 or pw1 != pw2:
             st.error("Passwords are empty or do not match.")
         else:
-            # Create secret + QR
-            secret = pyotp.random_base32()
-            totp = pyotp.TOTP(secret)
-            # issuer shows up in the authenticator app
-            uri = totp.provisioning_uri(name=admin_user or "SimpleLookup", issuer_name="Simple Password Lookup")
-            qr_png = make_qr_png(uri)
-
-            # Show QR and ask to verify code once
-            st.success("Scan this QR in your authenticator app, then enter a 6-digit code to verify.")
-            st.image(qr_png, caption="Scan in Google Authenticator / Microsoft Authenticator", use_column_width=False)
-            verify_code = st.text_input("Enter a 6-digit code from your app to verify", max_chars=6)
-            if st.button("Verify & Save"):
-                if totp.verify(verify_code, valid_window=1):
-                    hash_pw = ph.hash(new_pw)
-                    save_config({"password_hash": hash_pw, "totp_secret": secret, "label": admin_user})
-                    st.success("2FA setup complete. Please reload and log in.")
-                    st.stop()
-                else:
-                    st.error("Invalid code. Open your authenticator and try a fresh code.")
+            # generate secret only if not already generated
+            if not st.session_state.setup_secret:
+                st.session_state.setup_secret = pyotp.random_base32()
+            st.session_state.setup_label = label
+            st.session_state.setup_pw = pw1
+
+    # If secret staged, show QR and verify input
+    if st.session_state.setup_secret:
+        secret = st.session_state.setup_secret
+        label = st.session_state.setup_label or "SimpleLookup"
+        totp = pyotp.TOTP(secret)
+        uri = totp.provisioning_uri(name=label, issuer_name="Simple Password Lookup")
+        st.image(make_qr_png(uri), caption="Scan in Google/Microsoft Authenticator")
+
+        code = st.text_input("Enter a current 6-digit code to verify", max_chars=6)
+        if st.button("Verify & Save"):
+            if totp.verify(code, valid_window=1):
+                # hash password and save config
+                hash_pw = ph.hash(st.session_state.setup_pw)
+                save_config({"password_hash": hash_pw, "totp_secret": secret, "label": label})
+                # clear setup state
+                st.session_state.setup_secret = None
+                st.session_state.setup_label = None
+                st.session_state.setup_pw = None
+                st.success("2FA setup complete. Please log in below.")
+            else:
+                st.error("Invalid code. Open your authenticator and try a fresh code.")
     st.stop()
-else:
-    # Login form
+
+# ---------------- Step 3: Login (password + 2FA) ----------------
+if not st.session_state.authed:
     st.subheader("Login")
-    with st.form("login"):
+    with st.form("login_form"):
         pw = st.text_input("Admin password", type="password")
         code = st.text_input("Authenticator code (6 digits)", max_chars=6)
         ok = st.form_submit_button("Sign in")
-
     if ok:
         try:
             ph.verify(cfg["password_hash"], pw or "")
-            # allow Argon2 hash upgrade if needed
+            # upgrade hash when needed (argon2 best practice)
             if ph.check_needs_rehash(cfg["password_hash"]):
                 cfg["password_hash"] = ph.hash(pw or "")
                 save_config(cfg)
@@ -139,72 +173,43 @@ else:
         except Exception as e:
             st.session_state.failures += 1
             st.error(f"Login error: {e}")
-
     if not st.session_state.authed:
         if st.session_state.failures >= 5:
-            st.warning("Too many failed attempts. Wait 30 seconds and try again.")
+            st.warning("Too many failed attempts. Wait ~30 seconds and try again.")
         st.stop()
 
-# ---------- past this point: authenticated ----------
+# ---------------- Authenticated area ----------------
 st.success("Authenticated ✅")
 
-# -------------- load creds or upload once --------------
-if "creds" not in st.session_state:
-    if PERSIST_FILE.exists():
-        try:
-            st.session_state.creds = standardize_columns(pd.read_excel(PERSIST_FILE))
-            st.info(f"Loaded saved credentials from: {PERSIST_FILE}")
-        except Exception as e:
-            st.error(f"Found {PERSIST_FILE} but failed to read it: {e}")
-            st.session_state.creds = None
-    else:
-        st.session_state.creds = None
+# Load credentials
+try:
+    creds = standardize_columns(pd.read_excel(PERSIST_FILE))
+except Exception as e:
+    st.error(f"Could not read credentials file at {PERSIST_FILE}: {e}")
+    st.stop()
 
-if st.session_state.creds is None:
-    st.subheader("Upload your Excel/CSV (only once)")
-    up = st.file_uploader("Choose file", type=["xlsx", "xls", "csv"], accept_multiple_files=False)
-    if not up:
-        st.stop()
-    try:
-        df = standardize_columns(read_any(up))
-        st.session_state.creds = df
+# Optional: allow replacing saved file after auth
+with st.expander("Replace saved credentials (optional)"):
+    new_up = st.file_uploader("Upload new Excel/CSV to replace saved file", type=["xlsx", "xls", "csv"], key="replacer")
+    if new_up is not None:
         try:
-            PERSIST_FILE.parent.mkdir(parents=True, exist_ok=True)
-            df.to_excel(PERSIST_FILE, index=False)
-            st.success(f"Saved to {PERSIST_FILE}. You won't need to upload again.")
+            df_new = standardize_columns(read_any(new_up))
+            df_new.to_excel(PERSIST_FILE, index=False)
+            st.success(f"Replaced saved file at {PERSIST_FILE}. Reload to use the new data.")
         except Exception as e:
-            st.warning(f"Loaded for this session, but could not save for persistence: {e}")
-    except Exception as e:
-        st.error(f"Failed to read file: {e}")
-        st.stop()
-else:
-    with st.expander("Replace saved file (optional)"):
-        new_up = st.file_uploader("Upload new Excel/CSV to replace saved file", type=["xlsx", "xls", "csv"], key="replacer")
-        if new_up is not None:
-            try:
-                df_new = standardize_columns(read_any(new_up))
-                st.session_state.creds = df_new
-                try:
-                    PERSIST_FILE.parent.mkdir(parents=True, exist_ok=True)
-                    df_new.to_excel(PERSIST_FILE, index=False)
-                    st.success(f"Replaced saved file at {PERSIST_FILE}.")
-                except Exception as e:
-                    st.warning(f"Replaced in memory only, could not save: {e}")
-            except Exception as e:
-                st.error(f"Failed to read new file: {e}")
-
-# -------------- search UI --------------
+            st.error(f"Failed to save new file: {e}")
+
+# Search UI (same as before)
 st.subheader("Find your password")
 q = st.text_input("Search by platform/site/username")
 if q.strip():
     Q = q.lower().strip()
-    df = st.session_state.creds
     mask = (
-        df["name"].str.lower().str.contains(Q, na=False)
-        | df["username"].str.lower().str.contains(Q, na=False)
-        | df["url"].str.lower().str.contains(Q, na=False)
+        creds["name"].str.lower().str.contains(Q, na=False)
+        | creds["username"].str.lower().str.contains(Q, na=False)
+        | creds["url"].str.lower().str.contains(Q, na=False)
     )
-    results = df[mask]
+    results = creds[mask]
     if results.empty:
         st.warning("No matches found.")
     else:
@@ -218,13 +223,3 @@ if q.strip():
                     st.caption("Note: " + str(row["note"]))
 else:
     st.caption("Type a keyword above to search.")
-
-with st.expander("Security notes"):
-    st.markdown(
-        """
-- This gate protects the UI with **password + TOTP** (Google/Microsoft Authenticator).
-- Your credentials file is still a plain Excel you provided; this app **does not re-encrypt** it.
-- On Hugging Face Spaces, data persists under `/data`. Locally it’s `./`.
-- For stronger protection, store credentials in an encrypted vault (I can upgrade this to AES-GCM with a master password if you want).
-        """
-    )