Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,15 +1,17 @@
 import json
-import os
 import time
 import base64
 import secrets
+import os
+from pathlib import Path
 from dataclasses import dataclass, asdict
 from typing import List, Optional, Dict, Any
 
 import streamlit as st
 import pandas as pd  # for Excel/CSV import
 
-# ---- Crypto deps (install: pip install cryptography argon2-cffi==23.1.0) ----
+# ---- Crypto deps ----
+# pip install: streamlit cryptography argon2-cffi pandas openpyxl
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
@@ -17,9 +19,15 @@ from cryptography.hazmat.primitives.hmac import HMAC
 from cryptography.hazmat.backends import default_backend
 from argon2.low_level import hash_secret_raw, Type
 
-APP_TITLE = "SmartPass — Local, Zero‑Knowledge Password Manager"
+APP_TITLE = "SmartPass — Local, Zero‑Knowledge Password Manager (uploads-enabled)"
 VERSION = 1
 
+# -------------------------- Uploads / temp dirs --------------------------
+UPLOAD_DIR = Path(os.environ.get("UPLOAD_DIR", "./uploads")).resolve()
+TMP_DIR = Path(os.environ.get("TMPDIR", "./tmp")).resolve()
+for p in (UPLOAD_DIR, TMP_DIR):
+    p.mkdir(parents=True, exist_ok=True)
+
 # -------------------------- Utility helpers --------------------------
 
 def b64e(b: bytes) -> str:
@@ -52,8 +60,7 @@ class KDFParams:
     hash_len: int = 32
 
     def to_dict(self):
-        d = asdict(self)
-        return d
+        return asdict(self)
 
 
 @dataclass
@@ -131,10 +138,8 @@ def aesgcm_decrypt(key: bytes, nonce_b64: str, ct_b64: str, aad: Optional[bytes]
 # -------------------------- Integrity (HMAC) --------------------------
 
 def canonical_vault_for_hmac(v: Dict[str, Any]) -> bytes:
-    # Exclude integrity field itself to avoid recursion
     tmp = dict(v)
     tmp.pop("integrity_hmac_b64", None)
-    # Stable ordering
     return json.dumps(tmp, sort_keys=True, separators=(",", ":")).encode("utf-8")
 
 
@@ -162,7 +167,6 @@ def new_vault(password: str) -> Vault:
     kdf.salt_b64 = b64e(secrets.token_bytes(16))
     master_key = derive_master_key(password, kdf)
 
-    # Derive subkeys
     wrap_key = hkdf_expand(master_key, b"wrap-key")  # for wrapping data key
     mac_key = hkdf_expand(master_key, b"vault-mac")  # for HMAC integrity
 
@@ -175,23 +179,18 @@ def new_vault(password: str) -> Vault:
         items=[],
         integrity_hmac_b64="",
     )
-    # Compute initial HMAC
-    vdict = vault.to_dict()
-    vdict["integrity_hmac_b64"] = ""
-    hmac_b64 = compute_hmac(mac_key, vdict)
-    vault.integrity_hmac_b64 = hmac_b64
-    # Clear secrets from locals (best-effort)
+    vdict = vault.to_dict(); vdict["integrity_hmac_b64"] = ""
+    vault.integrity_hmac_b64 = compute_hmac(mac_key, vdict)
     return vault
 
 
 def unlock_vault(vault_dict: Dict[str, Any], password: str) -> Dict[str, Any]:
-    # Returns {"data_key": bytes, "mac_key": bytes}
-    kdf = KDFParams(**vault_dict["kdf"]) if isinstance(vault_dict["kdf"], dict) else vault_dict["kdf"]
+    kdfd = vault_dict["kdf"]
+    kdf = KDFParams(**kdfd) if isinstance(kdfd, dict) else kdfd
     master_key = derive_master_key(password, kdf)
     wrap_key = hkdf_expand(master_key, b"wrap-key")
     mac_key = hkdf_expand(master_key, b"vault-mac")
 
-    # Verify integrity first
     if not verify_hmac(mac_key, vault_dict):
         raise ValueError("Integrity check failed. The vault may be corrupted or the password is incorrect.")
 
@@ -233,6 +232,50 @@ def update_item(data_key: bytes, it: VaultItem, payload: Dict[str, Any]) -> Vaul
     return it
 
 
+# -------------------------- Import helpers (Excel/CSV) --------------------------
+
+def _read_tabular(file_or_path) -> pd.DataFrame:
+    # Accept a file-like object or a path
+    if hasattr(file_or_path, "read"):
+        # Try both excel/csv from buffer name fallback
+        try:
+            return pd.read_excel(file_or_path)
+        except Exception:
+            file_or_path.seek(0)
+            return pd.read_csv(file_or_path)
+    name = str(file_or_path).lower()
+    if name.endswith(".csv"):
+        return pd.read_csv(file_or_path)
+    return pd.read_excel(file_or_path)
+
+
+def _standardize_columns(df: pd.DataFrame) -> pd.DataFrame:
+    aliases = {
+        "name": ["name", "title", "site", "account"],
+        "username": ["username", "user", "login", "email"],
+        "url": ["url", "link", "website", "domain"],
+        "password": ["password", "pass", "pwd"],
+        "note": ["note", "notes", "remark", "remarks"],
+        "tags": ["tags", "label", "labels", "category", "categories"],
+        "type": ["type", "item_type"],
+    }
+    colmap: Dict[str, str] = {}
+    lower_cols = {c.lower().strip(): c for c in df.columns}
+    for target, alias_list in aliases.items():
+        for a in alias_list:
+            if a in lower_cols:
+                colmap[target] = lower_cols[a]
+                break
+    for col in ["name", "username", "url", "password", "note", "tags", "type"]:
+        if col not in colmap:
+            df[col] = ""
+            colmap[col] = col
+    return (
+        df[[colmap[c] for c in ["name", "username", "url", "password", "note", "tags", "type"]]]
+        .rename(columns={colmap[c]: c for c in colmap})
+    )
+
+
 # -------------------------- Streamlit UI --------------------------
 st.set_page_config(page_title=APP_TITLE, page_icon="🔐", layout="wide")
 
@@ -263,47 +306,10 @@ def do_lock():
     touch()
 
 
-# Sidebar: Create / Open / Lock / Export
-# -------------------------- Import helpers --------------------------
-
-def _read_tabular(file) -> pd.DataFrame:
-    name = (file.name or "").lower()
-    if name.endswith(".csv"):
-        return pd.read_csv(file)
-    # Excel requires openpyxl (installed implicitly by pandas if not, pip install openpyxl)
-    return pd.read_excel(file)
-
-
-def _standardize_columns(df: pd.DataFrame) -> pd.DataFrame:
-    # Map common column names to our schema
-    aliases = {
-        "name": ["name", "title", "site", "account"],
-        "username": ["username", "user", "login", "email"],
-        "url": ["url", "link", "website"],
-        "password": ["password", "pass", "pwd"],
-        "note": ["note", "notes", "remark", "remarks"],
-        "tags": ["tags", "label", "labels", "category", "categories"],
-        "type": ["type", "item_type"],
-    }
-    colmap: Dict[str, str] = {}
-    lower_cols = {c.lower().strip(): c for c in df.columns}
-    for target, alias_list in aliases.items():
-        for a in alias_list:
-            if a in lower_cols:
-                colmap[target] = lower_cols[a]
-                break
-    # Ensure all expected columns exist; create if missing
-    for col in ["name", "username", "url", "password", "note", "tags", "type"]:
-        if col not in colmap:
-            df[col] = ""
-            colmap[col] = col
-    return df[[colmap[c] for c in ["name", "username", "url", "password", "note", "tags", "type"]]]\
-             .rename(columns={colmap[c]: c for c in colmap})
-
-
+# Sidebar: Create / Open / Lock / Export / Import
 with st.sidebar:
     st.title("🔐 SmartPass")
-    st.caption("Local, file-based, zero-knowledge vault. No servers.")
+    st.caption(f"Uploads dir: {UPLOAD_DIR}")
 
     # Auto-lock
     lock_minutes = st.number_input("Auto-lock (minutes)", min_value=1, max_value=120, value=5)
@@ -324,13 +330,17 @@ with st.sidebar:
             st.error("Please enter your master password.")
         else:
             try:
-                vault_dict = json.loads(up.getvalue().decode("utf-8"))
+                # Persist a copy of uploaded vault and load from disk
+                saved_vault_path = UPLOAD_DIR / (up.name or "uploaded_vault.json")
+                with open(saved_vault_path, "wb") as f:
+                    f.write(up.getbuffer())
+                vault_dict = json.loads(saved_vault_path.read_text(encoding="utf-8"))
                 keys = unlock_vault(vault_dict, password_open)
                 st.session_state.vault = vault_dict
                 st.session_state.keys = keys
                 st.session_state.unlocked = True
                 touch()
-                st.success("Vault unlocked.")
+                st.success(f"Vault unlocked. Saved copy: {saved_vault_path}")
             except Exception as e:
                 st.session_state.unlocked = False
                 st.error(f"Failed to unlock: {e}")
@@ -341,7 +351,7 @@ with st.sidebar:
     new_pw2 = st.text_input("Confirm password", type="password")
     with st.expander("Advanced KDF (Argon2id)"):
         m_mib = st.slider("Memory (MiB)", 64, 512, 128, step=32)
-        t_cost = st.slider("Iterations", 1, 5, 3)
+        t_cost = st.slider("Iterations", 1, 6, 3)
         par = st.slider("Parallelism", 1, 4, 1)
     if st.button("Create Vault", use_container_width=True):
         if not new_pw:
@@ -350,66 +360,74 @@ with st.sidebar:
             st.error("Passwords do not match.")
         else:
             v = new_vault(new_pw)
-            # override KDF knobs from UI
             v.kdf.m_cost_kib = m_mib * 1024
             v.kdf.t_cost = t_cost
             v.kdf.parallelism = par
-            # Recompute HMAC with same master pw but updated kdf params
             mk = derive_master_key(new_pw, v.kdf)
             mac_key = hkdf_expand(mk, b"vault-mac")
-            vdict = v.to_dict()
-            vdict["integrity_hmac_b64"] = ""
-            vdict["integrity_hmac_b64"] = compute_hmac(mac_key, vdict)
+            vdict = v.to_dict(); vdict["integrity_hmac_b64"] = compute_hmac(mac_key, vdict)
             st.session_state.vault = vdict
             st.session_state.keys = unlock_vault(vdict, new_pw)
             st.session_state.unlocked = True
+            # Immediately save a new vault file under uploads
+            new_path = UPLOAD_DIR / "vault.smartpass.json"
+            new_path.write_text(json.dumps(vdict, separators=(",", ":")), encoding="utf-8")
             touch()
-            st.success("New vault created and unlocked.")
+            st.success(f"New vault created, unlocked, and saved at {new_path}.")
 
     st.divider()
-    st.subheader("Export")
-if st.session_state.vault is not None:
-    export_json = json.dumps(st.session_state.vault, separators=(",", ":"))
-    st.download_button(
-        label="Download vault JSON",
-        data=export_json,
-        file_name="vault.smartpass.json",
-        mime="application/json",
-        use_container_width=True,
-    )
+    st.subheader("Export / Backup")
+    if st.session_state.vault is not None:
+        export_json = json.dumps(st.session_state.vault, separators=(",", ":"))
+        st.download_button(
+            label="Download vault JSON",
+            data=export_json,
+            file_name="vault.smartpass.json",
+            mime="application/json",
+            use_container_width=True,
+        )
+        if st.button("Save vault to uploads/", use_container_width=True):
+            path = UPLOAD_DIR / "vault.smartpass.json"
+            path.write_text(export_json, encoding="utf-8")
+            st.success(f"Saved: {path}")
+
+    st.divider()
+    st.subheader("Import from Excel/CSV")
+    st.caption("Headers: name, username, url, password, note, tags, optional type (login/note). Aliases auto-mapped.")
+    imp = st.file_uploader("Upload creds file (.xlsx/.xls/.csv)", type=["xlsx","xls","csv"], accept_multiple_files=False, key="importer")
+    if imp is not None and st.session_state.unlocked:
+        try:
+            saved_creds_path = UPLOAD_DIR / (imp.name or "creds_upload.xlsx")
+            with open(saved_creds_path, "wb") as f:
+                f.write(imp.getbuffer())
+            df = _read_tabular(saved_creds_path)
+            df = _standardize_columns(df)
+            df['type'] = df['type'].fillna('').astype(str).str.lower().replace({'': 'login'})
+            added = 0
+            items_local: List[VaultItem] = [VaultItem(**it) for it in st.session_state.vault.get("items", [])]
+            for _, row in df.iterrows():
+                payload = {
+                    "type": (row.get('type') or 'login') if (row.get('type') in ['login','note']) else 'login',
+                    "name": str(row.get('name') or '').strip(),
+                    "username": str(row.get('username') or '').strip(),
+                    "url": str(row.get('url') or '').strip(),
+                    "password": str(row.get('password') or '').strip(),
+                    "note": str(row.get('note') or '').strip(),
+                    "tags": [t.strip() for t in str(row.get('tags') or '').split(',') if t.strip()],
+                }
+                if not any([payload['name'], payload['username'], payload['url'], payload['password'], payload['note'], payload['tags']]):
+                    continue
+                items_local.append(encrypt_item(st.session_state.keys["data_key"], payload))
+                added += 1
+            st.session_state.vault["items"] = [it.to_dict() for it in items_local]
+            re_hmac(st.session_state.vault, st.session_state.keys["mac_key"])
+            # Save updated vault immediately
+            (UPLOAD_DIR / "vault.smartpass.json").write_text(json.dumps(st.session_state.vault, separators=(",", ":")), encoding="utf-8")
+            st.success(f"Imported {added} item(s). Saved updated vault to uploads/.")
+            st.experimental_rerun()
+        except Exception as e:
+            st.error(f"Import failed: {e}")
 
-st.divider()
-st.subheader("Import from Excel/CSV")
-imp = st.file_uploader("Upload creds file (.xlsx/.xls/.csv)", type=["xlsx","xls","csv"], accept_multiple_files=False, key="importer")
-if imp is not None and st.session_state.unlocked:
-    try:
-        df = _read_tabular(imp)
-        df = _standardize_columns(df)
-        # Default type to 'login' if empty
-        df['type'] = df['type'].fillna('').astype(str).str.lower().replace({'': 'login'})
-        added = 0
-        items_local: List[VaultItem] = [VaultItem(**it) for it in st.session_state.vault.get("items", [])]
-        for _, row in df.iterrows():
-            payload = {
-                "type": (row.get('type') or 'login') if (row.get('type') in ['login','note']) else 'login',
-                "name": str(row.get('name') or '').strip(),
-                "username": str(row.get('username') or '').strip(),
-                "url": str(row.get('url') or '').strip(),
-                "password": str(row.get('password') or '').strip(),
-                "note": str(row.get('note') or '').strip(),
-                "tags": [t.strip() for t in str(row.get('tags') or '').split(',') if t.strip()],
-            }
-            # Skip empty rows
-            if not any(payload.values()):
-                continue
-            items_local.append(encrypt_item(st.session_state.keys["data_key"], payload))
-            added += 1
-        st.session_state.vault["items"] = [it.to_dict() for it in items_local]
-        re_hmac(st.session_state.vault, st.session_state.keys["mac_key"])
-        st.success(f"Imported {added} item(s).")
-        st.experimental_rerun()
-    except Exception as e:
-        st.error(f"Import failed: {e}")
 
 # -------------------------- Main Area --------------------------
 st.title(APP_TITLE)
@@ -418,7 +436,7 @@ if not st.session_state.unlocked:
     st.info("Unlock or create a vault from the left sidebar to begin.")
     st.stop()
 
-# Search and actions
+# Top actions
 col1, col2, col3, col4 = st.columns([3, 1, 1, 1])
 with col1:
     q = st.text_input("Search (name / username / url / tags)", value="")
@@ -433,13 +451,59 @@ with col4:
         do_lock()
         st.experimental_rerun()
 
+# Quick Password Lookup
+st.markdown("### 🔎 Quick Password Lookup")
+lookup = st.text_input("Which site/app? (e.g., gmail, netflix.com, bank)", key="quick_lookup")
+if lookup.strip():
+    LQ = lookup.strip().lower()
+    _matches = []
+    _items_all: List[VaultItem] = [VaultItem(**it) for it in st.session_state.vault.get("items", [])]
+    for _it in _items_all:
+        try:
+            _pl = decrypt_item(st.session_state.keys["data_key"], _it)
+        except Exception:
+            continue
+        _name = str(_pl.get("name", ""))
+        _user = str(_pl.get("username", ""))
+        _url  = str(_pl.get("url", ""))
+        _tags = ",".join(_pl.get("tags", []))
+        hay = " ".join([_name, _user, _url, _tags]).lower()
+        if LQ in hay:
+            _matches.append((_it, _pl))
+    if not _matches:
+        st.info("No matches found in your vault. If this account isn't stored here yet, import or add it first.")
+    else:
+        for _it, _pl in _matches[:10]:
+            with st.expander(f"{_pl.get('name') or '(unnamed)'} — {_pl.get('username','')}  [{_it.type}]"):
+                c1, c2 = st.columns([2, 1])
+                with c1:
+                    st.write(f"**URL:** {_pl.get('url','')}")
+                    show_pw = st.checkbox("Show password", key=f"show_{_it.id}")
+                    pw_val = _pl.get('password','') if _it.type == 'login' else ''
+                    st.text_input("Password", value=pw_val,
+                                  type=("default" if show_pw else "password"),
+                                  key=f"quick_pw_{_it.id}")
+                with c2:
+                    if st.button("Edit", key=f"quick_edit_{_it.id}"):
+                        st.session_state["edit_id"] = _it.id
+                        st.session_state["edit_payload"] = {
+                            **_pl,
+                            "type": _it.type,
+                            "name": _pl.get("name",""),
+                            "username": _pl.get("username",""),
+                            "url": _pl.get("url",""),
+                            "password": _pl.get("password",""),
+                            "note": _pl.get("note",""),
+                            "tags": ", ".join(_pl.get("tags", [])),
+                        }
+                        st.experimental_rerun()
+
 touch()
 
-# Items table
+# Decrypt all for listing
 vault_dict = st.session_state.vault
 keys = st.session_state.keys
 
-# Decrypt all (in-memory only)
 items: List[VaultItem] = [VaultItem(**it) for it in vault_dict.get("items", [])]
 rows = []
 for it in items:
@@ -459,7 +523,7 @@ for it in items:
     except Exception:
         rows.append({"_id": it.id, "type": it.type, "name": "<decrypt error>", "username": "", "url": "", "password": "", "note": "", "tags": "", "updated": ""})
 
-# Filter
+# Filter main list
 if q.strip():
     Q = q.lower()
     rows = [r for r in rows if any(Q in (str(r[k]) or "").lower() for k in ["name", "username", "url", "tags", "note"]) ]
@@ -485,11 +549,13 @@ for r in rows:
                 items = [it for it in items if it.id != r['_id']]
                 vault_dict["items"] = [it.to_dict() for it in items]
                 re_hmac(vault_dict, keys["mac_key"])
+                # Save updated vault immediately
+                (UPLOAD_DIR / "vault.smartpass.json").write_text(json.dumps(vault_dict, separators=(",", ":")), encoding="utf-8")
                 st.session_state.vault = vault_dict
-                st.success("Deleted.")
+                st.success("Deleted and saved.")
                 st.experimental_rerun()
 
-# Add / Edit forms
+# Add / Edit modals
 if st.session_state.get("add_type"):
     with st.modal("Add Item"):
         add_type = st.session_state.get("add_type")
@@ -513,9 +579,10 @@ if st.session_state.get("add_type"):
             items.append(new_item)
             vault_dict["items"] = [it.to_dict() for it in items]
             re_hmac(vault_dict, keys["mac_key"])
+            (UPLOAD_DIR / "vault.smartpass.json").write_text(json.dumps(vault_dict, separators=(",", ":")), encoding="utf-8")
             st.session_state.vault = vault_dict
             st.session_state["add_type"] = None
-            st.success("Item added.")
+            st.success("Item added and saved.")
             st.experimental_rerun()
         if st.button("Cancel"):
             st.session_state["add_type"] = None
@@ -544,25 +611,26 @@ if st.session_state.get("edit_id"):
                 "tags": [t.strip() for t in tags.split(",") if t.strip()],
             }
             updated = update_item(keys["data_key"], original, new_payload)
-            # Replace in list
             items = [updated if it.id == eid else it for it in items]
             vault_dict["items"] = [it.to_dict() for it in items]
             re_hmac(vault_dict, keys["mac_key"])
+            (UPLOAD_DIR / "vault.smartpass.json").write_text(json.dumps(vault_dict, separators=(",", ":")), encoding="utf-8")
             st.session_state.vault = vault_dict
             st.session_state["edit_id"] = None
-            st.success("Updated.")
+            st.success("Updated and saved.")
             st.experimental_rerun()
         if st.button("Cancel"):
             st.session_state["edit_id"] = None
 
-# Security footnotes
+# Security notes
 with st.expander("Security Notes & Tips"):
     st.markdown(
-        """
-- **All encryption/decryption happens locally.** Your vault is a JSON file that only contains ciphertext. Keep backups.
-- **Zero-knowledge:** The app never stores your master password. Only a key derived via **Argon2id** is used transiently in memory.
-- **Integrity:** The vault includes an HMAC to detect tampering/corruption. Wrong password also fails integrity.
-- **Clipboard caution:** This demo does not auto-copy passwords to avoid OS clipboard risks.
-- **KDF tuning:** If your device is slow, reduce memory/iterations in the sidebar's Advanced KDF.
+        f"""
+- **All crypto is local.** The vault JSON contains only ciphertext + metadata.
+- **Uploads dir:** `{UPLOAD_DIR}` (configurable via `UPLOAD_DIR`). Vault and imported files are saved here.
+- **Zero-knowledge:** The app never stores your master password; derivation happens in memory.
+- **Integrity:** Vault includes an HMAC to detect tampering/corruption.
+- **KDF tuning:** If your device is slow, reduce Argon2 memory/iterations in the sidebar.
+- **Backups:** Keep copies of your exported vault JSON in safe places.
         """
     )