Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,16 +1,26 @@
-# simple_lookup_persistent.py
+# simple_lookup_secure.py
+import json
+from pathlib import Path
+from io import BytesIO
+
 import streamlit as st
 import pandas as pd
-from pathlib import Path
+import pyotp
+import qrcode
+from PIL import Image
+from argon2 import PasswordHasher
+from argon2.exceptions import VerifyMismatchError
 
-st.set_page_config(page_title="Simple Password Lookup", page_icon="🔎", layout="centered")
-st.title("🔎 Simple Password Lookup")
-st.caption("Upload your Excel/CSV once → it will be saved for future sessions → search to reveal passwords.")
+st.set_page_config(page_title="Simple Password Lookup (2FA)", page_icon="🔐", layout="centered")
+st.title("🔐 Simple Password Lookup — protected with 2FA")
 
-# ---------------- Config ----------------
-# Hugging Face Spaces mounts /data as persistent storage; locally fall back to ./creds.xlsx
-PERSIST_FILE = Path("/data/creds.xlsx") if Path("/data").exists() else Path("creds.xlsx")
+# ---------- persistent locations ----------
+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
 
+# ---------- helpers ----------
+ph = PasswordHasher()
 ALIASES = {
     "name": ["name", "title", "site", "account", "platform", "service", "app"],
     "username": ["username", "user", "login", "email", "userid", "id"],
@@ -21,42 +31,124 @@ 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."""
-    # Normalize headers
     df.columns = [str(c).strip().lower() for c in df.columns]
     colmap = {}
-    # Map aliases
     for target, alias_list in ALIASES.items():
         for c in df.columns:
             if c in alias_list:
                 colmap[target] = c
                 break
-    # Ensure all expected columns exist
     for col in EXPECTED:
         if col not in colmap:
             df[col] = ""
             colmap[col] = col
-
     out = df[[colmap[c] for c in EXPECTED]].rename(columns={colmap[c]: c for c in colmap})
-    # Clean to strings
     for c in EXPECTED:
         out[c] = out[c].astype(str).fillna("").replace("nan", "")
-
-    # ✅ Keep rows that have at least one non-empty (after strip) among key fields
     key_cols = ["name", "username", "url", "password"]
     mask = out[key_cols].applymap(lambda x: str(x).strip() != "").any(axis=1)
-    out = out[mask].reset_index(drop=True)
-    return out
+    return out[mask].reset_index(drop=True)
 
 def read_any(file) -> pd.DataFrame:
-    """Read CSV or Excel from an uploaded file-like object."""
     name = (getattr(file, "name", "") or "").lower()
     if name.endswith(".csv"):
         return pd.read_csv(file)
-    # Default to Excel
     return pd.read_excel(file)
 
-# ---------------- Load or Upload once ----------------
+def load_config():
+    if CONFIG_FILE.exists():
+        return json.loads(CONFIG_FILE.read_text(encoding="utf-8"))
+    return None
+
+def save_config(cfg: dict):
+    CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
+    CONFIG_FILE.write_text(json.dumps(cfg, separators=(",", ":")), encoding="utf-8")
+
+def make_qr_png(data: str) -> bytes:
+    img = qrcode.make(data)
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+# ---------- auth gate (setup or login) ----------
+if "authed" not in st.session_state:
+    st.session_state.authed = False
+if "failures" not in st.session_state:
+    st.session_state.failures = 0
+
+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.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.")
+    st.stop()
+else:
+    # Login form
+    st.subheader("Login")
+    with st.form("login"):
+        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
+            if ph.check_needs_rehash(cfg["password_hash"]):
+                cfg["password_hash"] = ph.hash(pw or "")
+                save_config(cfg)
+            totp = pyotp.TOTP(cfg["totp_secret"])
+            if totp.verify(code, valid_window=1):
+                st.session_state.authed = True
+            else:
+                st.session_state.failures += 1
+                st.error("Invalid 2FA code.")
+        except VerifyMismatchError:
+            st.session_state.failures += 1
+            st.error("Invalid password.")
+        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.stop()
+
+# ---------- past this point: authenticated ----------
+st.success("Authenticated ✅")
+
+# -------------- load creds or upload once --------------
 if "creds" not in st.session_state:
     if PERSIST_FILE.exists():
         try:
@@ -69,14 +161,13 @@ if "creds" not in st.session_state:
         st.session_state.creds = None
 
 if st.session_state.creds is None:
-    st.subheader("1) Upload your Excel/CSV (only once)")
+    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
-        # Try to persist for future sessions
         try:
             PERSIST_FILE.parent.mkdir(parents=True, exist_ok=True)
             df.to_excel(PERSIST_FILE, index=False)
@@ -87,7 +178,6 @@ if st.session_state.creds is None:
         st.error(f"Failed to read file: {e}")
         st.stop()
 else:
-    # Optional: allow replacing the saved file
     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:
@@ -103,14 +193,12 @@ else:
             except Exception as e:
                 st.error(f"Failed to read new file: {e}")
 
-# ---------------- Search ----------------
-st.subheader("2) Find your password")
-q = st.text_input("Search by platform/site/username", placeholder="e.g., netflix, gmail.com, your@email.com")
-
+# -------------- search UI --------------
+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
-    # Robust contains with na=False to avoid NaN issues
     mask = (
         df["name"].str.lower().str.contains(Q, na=False)
         | df["username"].str.lower().str.contains(Q, na=False)
@@ -125,23 +213,18 @@ if q.strip():
             title = (row["name"] or row["username"] or row["url"]).strip()
             with st.expander(f"{title} — {row['username']} | {row['url']}"):
                 show = st.checkbox("Show password", key=f"show_{idx}")
-                st.text_input(
-                    "Password",
-                    value=row["password"],
-                    type=("default" if show else "password"),
-                    key=f"pw_{idx}",
-                )
+                st.text_input("Password", value=row["password"], type=("default" if show else "password"), key=f"pw_{idx}")
                 if str(row.get("note", "")).strip():
                     st.caption("Note: " + str(row["note"]))
 else:
     st.caption("Type a keyword above to search.")
 
-# ---------------- Footer ----------------
-with st.expander("Notes"):
+with st.expander("Security notes"):
     st.markdown(
         """
-- Accepted columns (case-insensitive; aliases auto-mapped): **name**, **username**, **url**, **password**, *(optional)* **note**.
-- Your first upload is saved to a persistent file so you **won’t have to re-upload again**.
-- Use the “Replace saved file” section above if you need to update your credentials later.
+- 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).
         """
     )