Update AUTH.py

AUTH.py

@@ -1,15 +1,14 @@
-# auth.py  (robust version for different Streamlit versions)
+# auth.py — robust, single-file authentication helper for Streamlit multipage
 import streamlit as st
 import hashlib
 from typing import Optional, Dict, Any
 
 # ---------------- CONFIG ----------------
 def _get_password_hash() -> Optional[str]:
-    raw = None
+    """Read APP_PASSWORD from st.secrets and return its SHA256 hex digest or None."""
     try:
         raw = st.secrets.get("APP_PASSWORD", None)
     except Exception:
-        # st.secrets might throw if not configured — treat as None
         raw = None
     if not raw:
         return None
@@ -29,71 +28,46 @@ def logout():
 
 # ---------------- Query param utilities (compat layer) ----------------
 def _get_query_params() -> Dict[str, Any]:
-    """
-    Return a dict-like of query params across Streamlit versions.
-    - Prefer st.query_params if present (newer API).
-    - Else try st.experimental_get_query_params.
-    - Else try st.get_query_params (older possibility).
-    - Else return {}.
-    """
+    """Return query params in a backward/forward compatible manner."""
     try:
-        # modern API returns a dict of lists
         if hasattr(st, "query_params"):
-            qp = dict(st.query_params or {})
-            return qp
+            return dict(st.query_params or {})
     except Exception:
         pass
-
     try:
         if hasattr(st, "experimental_get_query_params"):
-            qp = st.experimental_get_query_params() or {}
-            return dict(qp)
+            return dict(st.experimental_get_query_params() or {})
     except Exception:
         pass
-
     try:
-        # very old variants (some forks)
         if hasattr(st, "get_query_params"):
-            qp = st.get_query_params() or {}
-            return dict(qp)
+            return dict(st.get_query_params() or {})
     except Exception:
         pass
-
     return {}
 
 def _set_query_params(**kwargs):
-    """
-    Set query params using the available API.
-    kwargs should be key=value where value is string or list.
-    """
+    """Set query params using available API (best-effort)."""
     try:
-        # modern API
         if hasattr(st, "set_query_params"):
             st.set_query_params(**kwargs)
             return
     except Exception:
         pass
-
     try:
         if hasattr(st, "experimental_set_query_params"):
             st.experimental_set_query_params(**kwargs)
             return
     except Exception:
         pass
-
-    try:
-        if hasattr(st, "set_query_params"):  # redundant but safe
-            st.set_query_params(**kwargs)
-            return
-    except Exception:
-        pass
-
-    # if no API exists, do nothing (can't redirect)
     return
 
-# ---------------- Login widget (use in AUTH page) ----------------
-def login_ui():
-    """Return True only when login succeeded (use in AUTH page)."""
+# ---------------- Login UI and handler ----------------
+def login_ui_inline() -> bool:
+    """
+    Render login UI and return True if credentials were accepted during this call.
+    Does NOT change session state or redirect; caller should set session state.
+    """
     st.title("🔒 Authentication")
     st.markdown("Enter the app password to continue.")
     pwd = st.text_input("Password", type="password", key="_pwd_input")
@@ -103,34 +77,41 @@ def login_ui():
             return False
         entered_hash = hashlib.sha256(pwd.encode("utf-8")).hexdigest()
         if entered_hash == PASSWORD_HASH:
-            set_logged_in(True)
-            st.success("Login successful — redirecting...")
             return True
         else:
             st.error("Incorrect password")
             return False
     return False
 
+def auth_page_handler(redirect_page: str = "HOW TO USE"):
+    """
+    To be used in pages/AUTH.py.
+    Shows login UI; if successful it sets session logged flag and redirects to redirect_page.
+    """
+    succeeded = login_ui_inline()
+    if succeeded:
+        set_logged_in(True)
+        # set query param to requested page and rerun so multipage switches
+        _set_query_params(page=redirect_page)
+        st.experimental_rerun()
+
 # ---------------- Enforcement helper ----------------
 def require_login(auth_page_name: str = "AUTH", redirect_page: Optional[str] = None):
     """
-    Call at top of every page. If user is not logged in:
-      - redirect user to the auth page and stop rendering.
-    auth_page_name : exact page title of your AUTH page.
+    Call at the top of every non-AUTH page:
+        from auth import require_login
+        require_login(auth_page_name="AUTH", redirect_page="HOW TO USE")
+    If user is not logged in this will redirect to the AUTH page and stop rendering.
     """
-    # ensure session_state key exists
     if "_logged_in" not in st.session_state:
         st.session_state["_logged_in"] = False
 
-    # If user already logged in, do nothing
     if is_logged_in():
         return
 
-    # Get current page param in a robust way
     params = _get_query_params()
     cur_page = None
     try:
-        # params values often are lists (Streamlit), handle both cases
         if "page" in params:
             p = params["page"]
             if isinstance(p, (list, tuple)):
@@ -140,30 +121,10 @@ def require_login(auth_page_name: str = "AUTH", redirect_page: Optional[str] = N
     except Exception:
         cur_page = None
 
-    # Allow rendering if already on AUTH page (so user can login)
+    # allow rendering if already on AUTH page (user needs to see login UI)
     if cur_page == auth_page_name:
         return
 
-    # Otherwise set page to AUTH (if set_query API available) and stop
+    # otherwise redirect to AUTH page and stop further rendering
     _set_query_params(page=auth_page_name)
     st.stop()
-
-def auth_page_handler(redirect_page="HOME"):
-    """UI page where user enters password"""
-    st.title("🔐 Login Required")
-
-    # Input field
-    password = st.text_input("Enter password:", type="password")
-
-    # When user submits
-    if st.button("Login"):
-        if password == APP_PASSWORD:
-            st.session_state["authenticated"] = True
-            st.session_state["auth_time"] = now_ts()
-            st.success("Login successful! Redirecting...")
-
-            st.switch_page(f"{redirect_page}.py")
-            st.stop()
-        else:
-            st.error("Incorrect password")
-