Update src/streamlit_app.py

src/streamlit_app.py

@@ -17,10 +17,10 @@ st.title("🔐 Simple Password Lookup — 2FA protected")
 
 # ---------------- Persistence ----------------
 PERSIST_DIR  = Path("/data") if Path("/data").exists() else Path(".")
-PERSIST_FILE = PERSIST_DIR / "creds.xlsx"      # saved credentials file (uploaded once)
-CONFIG_FILE  = PERSIST_DIR / "auth.json"       # stores Argon2 password hash + TOTP secret + label
+PERSIST_FILE = PERSIST_DIR / "creds.xlsx"      # saved credentials (uploaded once)
+CONFIG_FILE  = PERSIST_DIR / "auth.json"       # Argon2 password hash + TOTP secret + label
 
-# ---------------- Globals / helpers ----------------
+# ---------------- Helpers ----------------
 ph = PasswordHasher()
 ALIASES = {
     "name": ["name", "title", "site", "account", "platform", "service", "app"],
@@ -53,6 +53,7 @@ def standardize_columns(df: pd.DataFrame) -> pd.DataFrame:
     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)
@@ -78,13 +79,10 @@ 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
+# setup state (only used during first-time 2FA setup)
+st.session_state.setdefault("setup_secret", None)
+st.session_state.setdefault("setup_label", None)
+st.session_state.setdefault("setup_pw", None)
 
 # ---------------- Step 1: Upload (only if not saved yet) ----------------
 if not PERSIST_FILE.exists():
@@ -110,24 +108,26 @@ if not cfg:
 
     # 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 "")
+        label = st.text_input(
+            "Account label (as shown in your authenticator)",
+            value=st.session_state.get("setup_label") or "SimpleLookup"
+        )
+        pw1 = st.text_input("New admin password", type="password", value=st.session_state.get("setup_pw") or "")
+        pw2 = st.text_input("Confirm password", type="password", value=st.session_state.get("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:
-            # generate secret only if not already generated
-            if not st.session_state.setup_secret:
+            if not st.session_state.get("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:
+    if st.session_state.get("setup_secret"):
         secret = st.session_state.setup_secret
-        label = st.session_state.setup_label or "SimpleLookup"
+        label = st.session_state.get("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")
@@ -135,16 +135,15 @@ if not cfg:
         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.")
+                # clear transient setup state
+                for k in ("setup_secret", "setup_label", "setup_pw"):
+                    st.session_state.pop(k, None)
+                st.success("2FA setup complete. Loading login…")
+                st.experimental_rerun()  # 🔁 rerun so cfg loads and login form appears
             else:
-                st.error("Invalid code. Open your authenticator and try a fresh code.")
+                st.error("Invalid code. Try a fresh code from your authenticator.")
     st.stop()
 
 # ---------------- Step 3: Login (password + 2FA) ----------------