src/streamlit_app.py

# simple_lookup_2fa.py
from __future__ import annotations
import json
from pathlib import Path
from io import BytesIO

import streamlit as st
import pandas as pd
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 (2FA)", page_icon="🔐", layout="centered")
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

# ---------------- Globals / helpers ----------------
ph = PasswordHasher()
ALIASES = {
    "name": ["name", "title", "site", "account", "platform", "service", "app"],
    "username": ["username", "user", "login", "email", "userid", "id"],
    "url": ["url", "link", "website", "domain"],
    "password": ["password", "pass", "pwd", "secret"],
    "note": ["note", "notes", "remark", "remarks"],
}
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():
        for c in df.columns:
            if c in alias_list:
                colmap[target] = c
                break
    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})
    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)

def read_any(file) -> pd.DataFrame:
    name = (getattr(file, "name", "") or "").lower()
    if name.endswith(".csv"):
        return pd.read_csv(file)
    return pd.read_excel(file)

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()

# ---------------- 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("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:
            # 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()

# ---------------- Step 3: Login (password + 2FA) ----------------
if not st.session_state.authed:
    st.subheader("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 "")
            # 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)
            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()

# ---------------- Authenticated area ----------------
st.success("Authenticated ✅")

# 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()

# 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:
            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.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()
    mask = (
        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 = creds[mask]
    if results.empty:
        st.warning("No matches found.")
    else:
        st.caption(f"Matches: {len(results)} (showing up to 50)")
        for idx, row in results.head(50).iterrows():
            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}")
                if str(row.get("note", "")).strip():
                    st.caption("Note: " + str(row["note"]))
else:
    st.caption("Type a keyword above to search.")