Update app.py

app.py

@@ -1,22 +1,45 @@
 import os
 import json
 import secrets
-from datetime import datetime
+import smtplib
+import ssl
+from email.message import EmailMessage
+from datetime import datetime, timedelta
 
 import streamlit as st
 from sqlalchemy import (
-    create_engine, Column, Integer, String, Float, DateTime, ForeignKey, Text
+    create_engine, Column, Integer, String, Float, DateTime, ForeignKey, Text, Boolean
 )
 from sqlalchemy.orm import declarative_base, sessionmaker, relationship
 from passlib.context import CryptContext
 
+
+# ----------------------------
+# App Config
+# ----------------------------
+APP_NAME = "Bank AHR"
+
+# Email / SMTP settings (set these in your hosting Secrets/Env)
+SMTP_HOST = os.getenv("SMTP_HOST", "")
+SMTP_PORT = int(os.getenv("SMTP_PORT", "587"))
+SMTP_USER = os.getenv("SMTP_USER", "")
+SMTP_PASS = os.getenv("SMTP_PASS", "")
+SMTP_FROM = os.getenv("SMTP_FROM", SMTP_USER)
+
+# Optional: for development only (shows codes in UI when email isn't configured)
+DEV_MODE = os.getenv("DEV_MODE", "false").lower() in ("1", "true", "yes")
+
+# Verification settings
+CODE_TTL_MINUTES = int(os.getenv("CODE_TTL_MINUTES", "15"))          # code expires after 15 min
+RESEND_COOLDOWN_SECONDS = int(os.getenv("RESEND_COOLDOWN_SECONDS", "60"))  # resend cooldown
+
+
 # ----------------------------
-# Config: DB path ("memory")
+# Config: DB path (persistent)
 # ----------------------------
 # Hugging Face Spaces persistent storage (if enabled) is typically mounted at /data
-# If /data is not available, it will use local file (may reset on rebuild).
 DB_DIR = "/data" if os.path.isdir("/data") else "."
-DB_PATH = os.path.join(DB_DIR, "ibank.db")
+DB_PATH = os.path.join(DB_DIR, "bank_ahr.db")
 
 engine = create_engine(f"sqlite:///{DB_PATH}", connect_args={"check_same_thread": False})
 SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
@@ -36,6 +59,15 @@ class User(Base):
     password_hash = Column(String, nullable=False)
     created_at = Column(DateTime, default=datetime.utcnow)
 
+    # Email verification
+    email_verified = Column(Boolean, default=False, nullable=False)
+    email_verification_code_hash = Column(String, nullable=True)
+    email_verification_sent_at = Column(DateTime, nullable=True)
+
+    # Password reset
+    password_reset_code_hash = Column(String, nullable=True)
+    password_reset_sent_at = Column(DateTime, nullable=True)
+
     accounts = relationship("BankAccount", back_populates="owner", cascade="all, delete-orphan")
     payees = relationship("Payee", back_populates="owner", cascade="all, delete-orphan")
 
@@ -82,25 +114,64 @@ class Transaction(Base):
 Base.metadata.create_all(bind=engine)
 
 
+# ----------------------------
+# Lightweight schema upgrade (SQLite ALTER TABLE for new columns)
+# ----------------------------
+def ensure_schema():
+    """
+    If an older DB exists, add missing columns without needing migrations.
+    """
+    session = SessionLocal()
+    try:
+        # Check columns in users table
+        cols = session.execute("PRAGMA table_info(users);").fetchall()
+        col_names = {c[1] for c in cols}  # c[1] is column name
+
+        def add_col(sql):
+            session.execute(sql)
+
+        if "email_verified" not in col_names:
+            add_col("ALTER TABLE users ADD COLUMN email_verified BOOLEAN NOT NULL DEFAULT 0;")
+        if "email_verification_code_hash" not in col_names:
+            add_col("ALTER TABLE users ADD COLUMN email_verification_code_hash VARCHAR;")
+        if "email_verification_sent_at" not in col_names:
+            add_col("ALTER TABLE users ADD COLUMN email_verification_sent_at DATETIME;")
+        if "password_reset_code_hash" not in col_names:
+            add_col("ALTER TABLE users ADD COLUMN password_reset_code_hash VARCHAR;")
+        if "password_reset_sent_at" not in col_names:
+            add_col("ALTER TABLE users ADD COLUMN password_reset_sent_at DATETIME;")
+
+        session.commit()
+    except Exception:
+        session.rollback()
+        # If PRAGMA/ALTER fails for any reason, app still runs with fresh DB
+    finally:
+        session.close()
+
+ensure_schema()
+
+
 # ----------------------------
 # Helpers
 # ----------------------------
+def db_session():
+    return SessionLocal()
+
 def hash_password(pw: str) -> str:
     return pwd_context.hash(pw)
 
 def verify_password(pw: str, hashed: str) -> bool:
     return pwd_context.verify(pw, hashed)
 
-def db_session():
-    return SessionLocal()
+def hash_code(code: str) -> str:
+    return pwd_context.hash(code)
 
-def generate_account_number(session) -> str:
-    # 12-digit random number, ensure uniqueness
-    while True:
-        acct = "".join(str(secrets.randbelow(10)) for _ in range(12))
-        exists = session.query(BankAccount).filter_by(account_number=acct).first()
-        if not exists:
-            return acct
+def verify_code(code: str, hashed: str) -> bool:
+    return pwd_context.verify(code, hashed)
+
+def generate_otp_code() -> str:
+    # 6-digit numeric code
+    return f"{secrets.randbelow(1_000_000):06d}"
 
 def money(x: float) -> str:
     return f"${x:,.2f}"
@@ -111,25 +182,234 @@ def require_login():
         st.stop()
 
 def get_user(session):
-    return session.query(User).get(st.session_state["user_id"])
+    return session.get(User, st.session_state["user_id"])
 
 def get_user_accounts(session, user_id: int):
     return session.query(BankAccount).filter_by(owner_id=user_id).order_by(BankAccount.created_at.asc()).all()
 
+def can_resend(last_sent_at: datetime | None) -> tuple[bool, str]:
+    if not last_sent_at:
+        return True, ""
+    delta = datetime.utcnow() - last_sent_at
+    if delta.total_seconds() >= RESEND_COOLDOWN_SECONDS:
+        return True, ""
+    wait = int(RESEND_COOLDOWN_SECONDS - delta.total_seconds())
+    return False, f"Please wait {wait}s before requesting another code."
+
+def is_code_expired(sent_at: datetime | None) -> bool:
+    if not sent_at:
+        return True
+    return datetime.utcnow() > (sent_at + timedelta(minutes=CODE_TTL_MINUTES))
+
+
+# ----------------------------
+# Email sending
+# ----------------------------
+def smtp_configured() -> bool:
+    return bool(SMTP_HOST and SMTP_PORT and SMTP_USER and SMTP_PASS and SMTP_FROM)
+
+def send_email(to_email: str, subject: str, body: str) -> tuple[bool, str]:
+    """
+    Sends email via SMTP. Requires SMTP_* env vars.
+    """
+    if not smtp_configured():
+        if DEV_MODE:
+            return False, "SMTP not configured (DEV_MODE enabled)."
+        return False, "Email service is not configured. Please set SMTP secrets."
+
+    try:
+        msg = EmailMessage()
+        msg["From"] = SMTP_FROM
+        msg["To"] = to_email
+        msg["Subject"] = subject
+        msg.set_content(body)
+
+        context = ssl.create_default_context()
+        with smtplib.SMTP(SMTP_HOST, SMTP_PORT) as server:
+            server.starttls(context=context)
+            server.login(SMTP_USER, SMTP_PASS)
+            server.send_message(msg)
+        return True, "Email sent."
+    except Exception as e:
+        return False, f"Failed to send email: {e}"
+
+
+# ----------------------------
+# Verification & Reset flows
+# ----------------------------
+def request_email_verification(email: str) -> tuple[bool, str, str | None]:
+    """
+    Generates and emails a verification code, storing hash + timestamp.
+    Returns (ok, msg, raw_code_if_dev)
+    """
+    session = db_session()
+    raw_code = None
+    try:
+        user = session.query(User).filter_by(email=email.lower().strip()).first()
+        if not user:
+            return False, "No account found for that email.", None
+        if user.email_verified:
+            return True, "Email already verified.", None
+
+        ok_resend, msg_resend = can_resend(user.email_verification_sent_at)
+        if not ok_resend:
+            return False, msg_resend, None
+
+        raw_code = generate_otp_code()
+        user.email_verification_code_hash = hash_code(raw_code)
+        user.email_verification_sent_at = datetime.utcnow()
+        session.commit()
+
+        subject = f"{APP_NAME} - Email Verification Code"
+        body = (
+            f"Your {APP_NAME} verification code is: {raw_code}\n\n"
+            f"This code expires in {CODE_TTL_MINUTES} minutes.\n"
+            "If you did not request this, you can ignore this email."
+        )
+        sent, send_msg = send_email(user.email, subject, body)
+
+        if sent:
+            return True, "Verification code sent to your email.", (raw_code if DEV_MODE else None)
+
+        # If not sent and DEV_MODE is on, we still allow showing the code
+        if DEV_MODE:
+            return True, f"{send_msg} (Showing code in DEV_MODE)", raw_code
+
+        return False, send_msg, None
+    except Exception as e:
+        session.rollback()
+        return False, f"Error: {e}", None
+    finally:
+        session.close()
+
+def verify_email_with_code(email: str, code: str) -> tuple[bool, str]:
+    session = db_session()
+    try:
+        user = session.query(User).filter_by(email=email.lower().strip()).first()
+        if not user:
+            return False, "No account found for that email."
+        if user.email_verified:
+            return True, "Email is already verified."
+
+        if not user.email_verification_code_hash or not user.email_verification_sent_at:
+            return False, "No verification code requested yet. Please request a new code."
+
+        if is_code_expired(user.email_verification_sent_at):
+            return False, "Code expired. Please request a new verification code."
+
+        if not verify_code(code.strip(), user.email_verification_code_hash):
+            return False, "Invalid code."
+
+        user.email_verified = True
+        user.email_verification_code_hash = None
+        user.email_verification_sent_at = None
+        session.commit()
+        return True, "Email verified successfully! You can now log in."
+    except Exception as e:
+        session.rollback()
+        return False, f"Error: {e}"
+    finally:
+        session.close()
+
+def request_password_reset(email: str) -> tuple[bool, str, str | None]:
+    """
+    Sends reset code and stores hash.
+    """
+    session = db_session()
+    raw_code = None
+    try:
+        user = session.query(User).filter_by(email=email.lower().strip()).first()
+        # Do not reveal whether account exists in real apps; but per your request we keep it straightforward.
+        if not user:
+            return False, "No account found for that email.", None
+
+        ok_resend, msg_resend = can_resend(user.password_reset_sent_at)
+        if not ok_resend:
+            return False, msg_resend, None
+
+        raw_code = generate_otp_code()
+        user.password_reset_code_hash = hash_code(raw_code)
+        user.password_reset_sent_at = datetime.utcnow()
+        session.commit()
+
+        subject = f"{APP_NAME} - Password Reset Code"
+        body = (
+            f"Your {APP_NAME} password reset code is: {raw_code}\n\n"
+            f"This code expires in {CODE_TTL_MINUTES} minutes.\n"
+            "If you did not request this, you can ignore this email."
+        )
+        sent, send_msg = send_email(user.email, subject, body)
+
+        if sent:
+            return True, "Password reset code sent to your email.", (raw_code if DEV_MODE else None)
+
+        if DEV_MODE:
+            return True, f"{send_msg} (Showing code in DEV_MODE)", raw_code
+
+        return False, send_msg, None
+    except Exception as e:
+        session.rollback()
+        return False, f"Error: {e}", None
+    finally:
+        session.close()
+
+def reset_password_with_code(email: str, code: str, new_password: str) -> tuple[bool, str]:
+    session = db_session()
+    try:
+        user = session.query(User).filter_by(email=email.lower().strip()).first()
+        if not user:
+            return False, "No account found for that email."
+
+        if not user.password_reset_code_hash or not user.password_reset_sent_at:
+            return False, "No reset code requested yet."
+
+        if is_code_expired(user.password_reset_sent_at):
+            return False, "Reset code expired. Please request a new one."
+
+        if not verify_code(code.strip(), user.password_reset_code_hash):
+            return False, "Invalid reset code."
+
+        if len(new_password) < 6:
+            return False, "Password should be at least 6 characters."
+
+        user.password_hash = hash_password(new_password)
+        user.password_reset_code_hash = None
+        user.password_reset_sent_at = None
+        session.commit()
+        return True, "Password updated successfully! You can now log in."
+    except Exception as e:
+        session.rollback()
+        return False, f"Error: {e}"
+    finally:
+        session.close()
+
 
 # ----------------------------
 # Core banking operations
 # ----------------------------
+def generate_account_number(session) -> str:
+    # 12-digit random number, ensure uniqueness
+    while True:
+        acct = "".join(str(secrets.randbelow(10)) for _ in range(12))
+        exists = session.query(BankAccount).filter_by(account_number=acct).first()
+        if not exists:
+            return acct
+
 def open_account(full_name: str, email: str, password: str, nickname: str = "Checking"):
     session = db_session()
     try:
-        if session.query(User).filter_by(email=email.lower().strip()).first():
+        email_clean = email.lower().strip()
+        if session.query(User).filter_by(email=email_clean).first():
             return False, "Email already registered."
 
+        if len(password) < 6:
+            return False, "Password should be at least 6 characters."
+
         user = User(
             full_name=full_name.strip(),
-            email=email.lower().strip(),
+            email=email_clean,
             password_hash=hash_password(password),
+            email_verified=False,
         )
         session.add(user)
         session.flush()
@@ -142,8 +422,42 @@ def open_account(full_name: str, email: str, password: str, nickname: str = "Che
             balance=0.0,
         )
         session.add(acct)
+
+        # Create verification code on register
+        raw_code = generate_otp_code()
+        user.email_verification_code_hash = hash_code(raw_code)
+        user.email_verification_sent_at = datetime.utcnow()
+
         session.commit()
-        return True, f"Account opened! Your new account number: {acct_num}"
+
+        # Send verification email
+        subject = f"{APP_NAME} - Email Verification Code"
+        body = (
+            f"Welcome to {APP_NAME}!\n\n"
+            f"Your verification code is: {raw_code}\n\n"
+            f"This code expires in {CODE_TTL_MINUTES} minutes."
+        )
+        sent, send_msg = send_email(user.email, subject, body)
+
+        if sent:
+            return True, (
+                f"Account created! Your new account number: {acct_num}\n"
+                f"A verification code has been sent to your email."
+            )
+
+        if DEV_MODE:
+            return True, (
+                f"Account created! Your new account number: {acct_num}\n"
+                f"{send_msg}\n"
+                f"DEV_MODE code: {raw_code}"
+            )
+
+        return True, (
+            f"Account created! Your new account number: {acct_num}\n"
+            f"NOTE: {send_msg}\n"
+            f"Ask admin to configure SMTP so you can verify email."
+        )
+
     except Exception as e:
         session.rollback()
         return False, f"Error: {e}"
@@ -157,6 +471,9 @@ def login(email: str, password: str):
         if not user or not verify_password(password, user.password_hash):
             return False, "Invalid email or password."
 
+        if not user.email_verified:
+            return False, "Email not verified. Please verify your email before logging in."
+
         st.session_state["user_id"] = user.id
         st.session_state["user_name"] = user.full_name
         return True, f"Welcome, {user.full_name}!"
@@ -174,7 +491,9 @@ def deposit(user_id: int, account_number: str, amount: float, note: str = ""):
 
     session = db_session()
     try:
-        acct = session.query(BankAccount).filter_by(owner_id=user_id, account_number=account_number).with_for_update().first()
+        acct = session.query(BankAccount).filter_by(
+            owner_id=user_id, account_number=account_number
+        ).with_for_update().first()
         if not acct:
             return False, "Account not found."
 
@@ -204,11 +523,15 @@ def transfer(user_id: int, from_acct_num: str, to_acct_num: str, amount: float,
 
     session = db_session()
     try:
-        from_acct = session.query(BankAccount).filter_by(owner_id=user_id, account_number=from_acct_num).with_for_update().first()
+        from_acct = session.query(BankAccount).filter_by(
+            owner_id=user_id, account_number=from_acct_num
+        ).with_for_update().first()
         if not from_acct:
             return False, "Source account not found."
 
-        to_acct = session.query(BankAccount).filter_by(account_number=to_acct_num).with_for_update().first()
+        to_acct = session.query(BankAccount).filter_by(
+            account_number=to_acct_num
+        ).with_for_update().first()
         if not to_acct:
             return False, "Destination account number not found."
 
@@ -255,7 +578,9 @@ def pay_bill(user_id: int, from_acct_num: str, payee_id: int, amount: float, not
 
     session = db_session()
     try:
-        acct = session.query(BankAccount).filter_by(owner_id=user_id, account_number=from_acct_num).with_for_update().first()
+        acct = session.query(BankAccount).filter_by(
+            owner_id=user_id, account_number=from_acct_num
+        ).with_for_update().first()
         if not acct:
             return False, "Account not found."
 
@@ -288,7 +613,7 @@ def pay_bill(user_id: int, from_acct_num: str, payee_id: int, amount: float, not
 # ----------------------------
 # UI
 # ----------------------------
-st.set_page_config(page_title="iBank", page_icon="🏦", layout="wide")
+st.set_page_config(page_title=APP_NAME, page_icon="🏦", layout="wide")
 
 st.markdown(
     """
@@ -297,6 +622,7 @@ st.markdown(
       .subtle { color: #6b7280; margin-top: 0; }
       .card { padding: 1rem; border-radius: 14px; border: 1px solid #e5e7eb; background: #fff; }
       .small { font-size: 0.9rem; color: #6b7280; }
+      .codehint { font-size: 0.9rem; padding: 0.75rem; border-radius: 10px; border: 1px dashed #e5e7eb; }
     </style>
     """,
     unsafe_allow_html=True,
@@ -305,7 +631,8 @@ st.markdown(
 left, right = st.columns([3, 2], gap="large")
 
 with left:
-    st.markdown('<div class="big-title">🏦 iBank</div>', unsafe_allow_html=True)
+    st.markdown(f'<div class="big-title">🏦 {APP_NAME}</div>', unsafe_allow_html=True)
+    st.markdown('<p class="subtle">Secure banking demo with email verification</p>', unsafe_allow_html=True)
 
 with right:
     if st.session_state.get("user_id"):
@@ -318,6 +645,9 @@ with right:
 
 st.divider()
 
+# ----------------------------
+# Auth Screens (not logged in)
+# ----------------------------
 if not st.session_state.get("user_id"):
     colA, colB = st.columns(2, gap="large")
 
@@ -335,22 +665,75 @@ if not st.session_state.get("user_id"):
                 else:
                     st.error(msg)
 
+        st.markdown("#### Verify email")
+        with st.form("verify_email_form", clear_on_submit=False):
+            v_email = st.text_input("Email to verify", placeholder="you@example.com", key="verify_email_input")
+            v_code = st.text_input("Verification code (6 digits)", placeholder="123456", key="verify_code_input")
+            c1, c2 = st.columns(2)
+            with c1:
+                verify_btn = st.form_submit_button("Verify")
+            with c2:
+                resend_btn = st.form_submit_button("Resend code")
+
+            if resend_btn:
+                ok, msg, raw = request_email_verification(v_email)
+                (st.success if ok else st.error)(msg)
+                if raw and DEV_MODE:
+                    st.markdown(f'<div class="codehint">DEV_MODE code: <b>{raw}</b></div>', unsafe_allow_html=True)
+
+            if verify_btn:
+                ok, msg = verify_email_with_code(v_email, v_code)
+                (st.success if ok else st.error)(msg)
+
+        st.markdown("#### Forgot password")
+        with st.form("forgot_password_form", clear_on_submit=False):
+            fp_email = st.text_input("Account email", placeholder="you@example.com", key="fp_email")
+            fp_code = st.text_input("Reset code (6 digits)", placeholder="123456", key="fp_code")
+            fp_new_pw = st.text_input("New password", type="password", key="fp_new_pw")
+
+            c1, c2 = st.columns(2)
+            with c1:
+                send_reset = st.form_submit_button("Send reset code")
+            with c2:
+                do_reset = st.form_submit_button("Reset password")
+
+            if send_reset:
+                ok, msg, raw = request_password_reset(fp_email)
+                (st.success if ok else st.error)(msg)
+                if raw and DEV_MODE:
+                    st.markdown(f'<div class="codehint">DEV_MODE reset code: <b>{raw}</b></div>', unsafe_allow_html=True)
+
+            if do_reset:
+                ok, msg = reset_password_with_code(fp_email, fp_code, fp_new_pw)
+                (st.success if ok else st.error)(msg)
+
     with colB:
         st.markdown("### Open account (Register)")
         with st.form("register_form", clear_on_submit=True):
             full_name = st.text_input("Full name", placeholder="Md Ahsan Habib")
             email = st.text_input("Email", placeholder="you@example.com")
-            password = st.text_input("Password", type="password")
+            password = st.text_input("Password (min 6 chars)", type="password")
             nickname = st.text_input("First account nickname", value="Checking")
-            submit = st.form_submit_button("Open account")
+            submit = st.form_submit_button("Create account")
             if submit:
                 ok, msg = open_account(full_name, email, password, nickname)
                 if ok:
                     st.success(msg)
-                    st.info("Now log in with your email & password.")
+                    st.info("Now verify your email using the code, then log in.")
                 else:
                     st.error(msg)
 
+        if not smtp_configured():
+            st.warning(
+                "SMTP is not configured, so email codes can't be sent.\n\n"
+                "Set these environment variables/secrets:\n"
+                "- SMTP_HOST\n- SMTP_PORT (usually 587)\n- SMTP_USER\n- SMTP_PASS\n- SMTP_FROM\n\n"
+                "For local testing, you can set DEV_MODE=true to show codes in the UI."
+            )
+
+# ----------------------------
+# Main App (logged in)
+# ----------------------------
 else:
     require_login()
     session = db_session()
@@ -365,7 +748,7 @@ else:
         st.markdown('<div class="card">', unsafe_allow_html=True)
         st.markdown("**Total balance**")
         st.markdown(f"### {money(total_balance)}")
-        st.markdown('<div class="small">Across all your iBank accounts</div>', unsafe_allow_html=True)
+        st.markdown(f'<div class="small">Across all your {APP_NAME} accounts</div>', unsafe_allow_html=True)
         st.markdown("</div>", unsafe_allow_html=True)
 
     with c2:
@@ -404,10 +787,15 @@ else:
             if submit:
                 try:
                     acct_num = generate_account_number(session)
-                    acct = BankAccount(owner_id=user.id, account_number=acct_num, nickname=nickname.strip() or "Savings", balance=0.0)
+                    acct = BankAccount(
+                        owner_id=user.id,
+                        account_number=acct_num,
+                        nickname=nickname.strip() or "Savings",
+                        balance=0.0
+                    )
                     session.add(acct)
                     session.commit()
-                    st.success(f"Created new account: {nickname} — {acct_num}")
+                    st.success(f"Created new account: {acct.nickname} — {acct_num}")
                     st.rerun()
                 except Exception as e:
                     session.rollback()
@@ -440,7 +828,7 @@ else:
         else:
             with st.form("transfer_form"):
                 from_label = st.selectbox("From account", list(from_choices.keys()))
-                to_acct = st.text_input("To account number", placeholder="12-digit iBank account number")
+                to_acct = st.text_input("To account number", placeholder="12-digit account number")
                 amount = st.number_input("Amount", min_value=0.0, value=25.0, step=5.0)
                 note = st.text_input("Note (optional)", placeholder="Rent, gift, etc.")
                 submit = st.form_submit_button("Transfer")
@@ -498,7 +886,7 @@ else:
                 meta = {}
                 try:
                     meta = json.loads(t.meta or "{}")
-                except:
+                except Exception:
                     meta = {}
                 rows.append({
                     "Time (UTC)": t.created_at.strftime("%Y-%m-%d %H:%M:%S"),