06/03 Optimize app

app/__init__.py

@@ -2,7 +2,7 @@ import logging
 import os
 from flask import Flask
 from app.config import config
-from app.extensions import db, login_manager, migrate, csrf
+from app.extensions import db, login_manager, migrate, csrf, limiter
 from app.utils.formatters import format_currency, format_percent, format_large_number
 
 
@@ -76,6 +76,7 @@ def create_app(config_name=None):
     login_manager.init_app(app)
     migrate.init_app(app, db)
     csrf.init_app(app)
+    limiter.init_app(app)
 
     from app.routes.auth import auth_bp
     from app.routes.dashboard import dashboard_bp

app/extensions.py

@@ -2,11 +2,14 @@ from flask_sqlalchemy import SQLAlchemy
 from flask_login import LoginManager
 from flask_migrate import Migrate
 from flask_wtf.csrf import CSRFProtect
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 
 db = SQLAlchemy()
 login_manager = LoginManager()
 migrate = Migrate()
 csrf = CSRFProtect()
+limiter = Limiter(key_func=get_remote_address, storage_uri='memory://')
 
 login_manager.login_view = 'auth.login'
 login_manager.login_message = 'Please log in to access this page.'

app/models/schwab_connection.py

@@ -7,10 +7,11 @@ class SchwabConnection(db.Model):
     __tablename__ = 'schwab_connections'
 
     id               = db.Column(db.Integer, primary_key=True)
-    access_token     = db.Column(db.Text, nullable=False)
-    refresh_token    = db.Column(db.Text, nullable=False)
-    token_expires_at = db.Column(db.DateTime, nullable=False)   # UTC
-    is_active        = db.Column(db.Boolean, default=True)
+    access_token              = db.Column(db.Text, nullable=False)
+    refresh_token             = db.Column(db.Text, nullable=False)
+    token_expires_at          = db.Column(db.DateTime, nullable=False)   # access token expiry
+    refresh_token_expires_at  = db.Column(db.DateTime, nullable=True)    # refresh token expiry (7 days)
+    is_active                 = db.Column(db.Boolean, default=True)
     created_at       = db.Column(db.DateTime, default=datetime.utcnow)
     last_synced_at   = db.Column(db.DateTime, nullable=True)
 

app/models/user.py

@@ -16,6 +16,8 @@ class User(UserMixin, db.Model):
     currency = db.Column(db.String(10), default='USD')
     currency_symbol = db.Column(db.String(5), default='$')
     groq_model = db.Column(db.String(64), default='llama-3.3-70b-versatile')
+    totp_secret  = db.Column(db.String(64), nullable=True)
+    totp_enabled = db.Column(db.Boolean, default=False, nullable=False)
     created_at = db.Column(db.DateTime, default=datetime.utcnow)
     last_login = db.Column(db.DateTime, nullable=True)
 

app/routes/auth.py

@@ -1,13 +1,21 @@
-from flask import Blueprint, render_template, redirect, url_for, flash, request
+import io
+import base64
+import logging
+
+import pyotp
+import qrcode
+from flask import (Blueprint, render_template, redirect, url_for,
+                   flash, request, session)
 from flask_login import login_user, logout_user, login_required, current_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, PasswordField, BooleanField, SubmitField
 from wtforms.validators import DataRequired, Length
 from app.models.user import User
-from app.extensions import db
+from app.extensions import db, limiter
 from datetime import datetime
 
 auth_bp = Blueprint('auth', __name__, url_prefix='/auth')
+log = logging.getLogger(__name__)
 
 
 class LoginForm(FlaskForm):
@@ -17,7 +25,10 @@ class LoginForm(FlaskForm):
     submit = SubmitField('Sign In')
 
 
+# ── Login / logout ────────────────────────────────────────────────────────────
+
 @auth_bp.route('/login', methods=['GET', 'POST'])
+@limiter.limit('10 per minute; 30 per hour')
 def login():
     if current_user.is_authenticated:
         return redirect(url_for('dashboard.index'))
@@ -26,11 +37,25 @@ def login():
     if form.validate_on_submit():
         user = User.query.filter_by(username=form.username.data.strip()).first()
         if user and user.check_password(form.password.data):
+            if user.totp_enabled:
+                # Park the user id in the session and redirect to TOTP verification
+                session['_totp_pending_id'] = user.id
+                session['_totp_remember']   = bool(form.remember_me.data)
+                session['_totp_next']       = request.args.get('next', '')
+                log.info('[auth] TOTP required for user %s', user.username)
+                return redirect(url_for('auth.totp_verify'))
+
             login_user(user, remember=form.remember_me.data)
             user.last_login = datetime.utcnow()
             db.session.commit()
-            next_page = request.args.get('next')
-            return redirect(next_page or url_for('dashboard.index'))
+            log.info('[auth] user %s logged in', user.username)
+            next_page = request.args.get('next', '')
+            if not next_page.startswith('/'):
+                next_page = url_for('dashboard.index')
+            return redirect(next_page)
+
+        log.warning('[auth] failed login attempt for username=%r ip=%s',
+                    form.username.data, request.remote_addr)
         flash('Invalid username or password.', 'danger')
 
     return render_template('auth/login.html', form=form)
@@ -42,3 +67,104 @@ def logout():
     logout_user()
     flash('You have been logged out.', 'info')
     return redirect(url_for('auth.login'))
+
+
+# ── TOTP: second-factor verification ─────────────────────────────────────────
+
+@auth_bp.route('/totp/verify', methods=['GET', 'POST'])
+@limiter.limit('10 per minute')
+def totp_verify():
+    pending_id = session.get('_totp_pending_id')
+    if not pending_id:
+        return redirect(url_for('auth.login'))
+
+    user = db.session.get(User, pending_id)
+    if not user or not user.totp_enabled:
+        session.pop('_totp_pending_id', None)
+        return redirect(url_for('auth.login'))
+
+    error = None
+    if request.method == 'POST':
+        code = request.form.get('code', '').strip().replace(' ', '')
+        totp = pyotp.TOTP(user.totp_secret)
+        if totp.verify(code, valid_window=1):
+            session.pop('_totp_pending_id', None)
+            remember = session.pop('_totp_remember', False)
+            next_url = session.pop('_totp_next', '') or url_for('dashboard.index')
+            if not next_url.startswith('/'):
+                next_url = url_for('dashboard.index')
+            login_user(user, remember=remember)
+            user.last_login = datetime.utcnow()
+            db.session.commit()
+            log.info('[auth] TOTP verified for user %s', user.username)
+            return redirect(next_url)
+        log.warning('[auth] invalid TOTP code for user %s ip=%s',
+                    user.username, request.remote_addr)
+        error = 'Invalid code — please try again.'
+
+    return render_template('auth/totp_verify.html', error=error)
+
+
+# ── TOTP: setup (enable) ──────────────────────────────────────────────────────
+
+@auth_bp.route('/totp/setup', methods=['GET', 'POST'])
+@login_required
+def totp_setup():
+    user = current_user
+
+    if request.method == 'GET':
+        # Generate a new secret on every GET so the user always sees a fresh QR code
+        secret = pyotp.random_base32()
+        session['_totp_setup_secret'] = secret
+        totp    = pyotp.TOTP(secret)
+        uri     = totp.provisioning_uri(
+            name=user.username,
+            issuer_name='PFM — Personal Finance',
+        )
+        img     = qrcode.make(uri)
+        buf     = io.BytesIO()
+        img.save(buf, format='PNG')
+        buf.seek(0)
+        qr_data = 'data:image/png;base64,' + base64.b64encode(buf.read()).decode()
+        return render_template('auth/totp_setup.html',
+                               qr_data=qr_data, secret=secret, uri=uri)
+
+    # POST — user confirmed with their first code
+    secret = session.get('_totp_setup_secret')
+    code   = request.form.get('code', '').strip().replace(' ', '')
+    if not secret:
+        flash('Setup session expired. Please try again.', 'warning')
+        return redirect(url_for('auth.totp_setup'))
+
+    totp = pyotp.TOTP(secret)
+    if totp.verify(code, valid_window=1):
+        user.totp_secret  = secret
+        user.totp_enabled = True
+        db.session.commit()
+        session.pop('_totp_setup_secret', None)
+        log.info('[auth] TOTP enabled for user %s', user.username)
+        flash('Two-factor authentication enabled successfully.', 'success')
+        return redirect(url_for('settings.index'))
+
+    flash('Invalid code — the QR code was not scanned correctly. Please try again.', 'danger')
+    return redirect(url_for('auth.totp_setup'))
+
+
+# ── TOTP: disable ─────────────────────────────────────────────────────────────
+
+@auth_bp.route('/totp/disable', methods=['POST'])
+@login_required
+def totp_disable():
+    user = current_user
+    # Require password confirmation before disabling
+    password = request.form.get('password', '')
+    if not user.check_password(password):
+        flash('Incorrect password — 2FA not disabled.', 'danger')
+        return redirect(url_for('settings.index'))
+
+    user.totp_enabled = False
+    user.totp_secret  = None
+    db.session.commit()
+    log.info('[auth] TOTP disabled for user %s', user.username)
+    flash('Two-factor authentication disabled.', 'info')
+    return redirect(url_for('settings.index'))

app/routes/dashboard.py

@@ -42,6 +42,19 @@ def index():
     period = request.args.get('period', 'this_month')
     date_from, date_to, period_label = _parse_date_range(period)
 
+    # ── Schwab token expiry check ─────────────────────
+    from app.models.schwab_connection import SchwabConnection
+    from datetime import timedelta
+    schwab_warning = None
+    schwab_conn = SchwabConnection.query.filter_by(is_active=True).first()
+    if schwab_conn and schwab_conn.refresh_token_expires_at:
+        days_left = (schwab_conn.refresh_token_expires_at - datetime.utcnow()).days
+        if days_left <= 2:
+            schwab_warning = (
+                f'Schwab connection expires in {max(days_left, 0)} day(s). '
+                f'Reconnect now to keep syncing.'
+            )
+
     # ── Summary cards ────────────────────────────────
     total_income = db.session.query(
         func.coalesce(func.sum(Transaction.amount), 0)
@@ -150,7 +163,8 @@ def index():
                            recent_txns=recent_txns,
                            fx=fx,
                            fx_history_data=fx_history_data,
-                           ai_insight=ai_insight)
+                           ai_insight=ai_insight,
+                           schwab_warning=schwab_warning)
 
 
 @dashboard_bp.route('/api/fx-history')

app/services/schwab_service.py

@@ -144,8 +144,11 @@ def refresh_tokens(connection):
 
 def _apply_token_data(connection, data):
     """Write token fields from token response onto connection model."""
-    connection.access_token     = data['access_token']
-    connection.refresh_token    = data.get('refresh_token', connection.refresh_token)
+    connection.access_token    = data['access_token']
+    if data.get('refresh_token'):
+        connection.refresh_token            = data['refresh_token']
+        # Schwab refresh tokens expire 7 days after issue
+        connection.refresh_token_expires_at = datetime.utcnow() + timedelta(days=7)
     expires_in                  = int(data.get('expires_in', 1800))
     connection.token_expires_at = datetime.utcnow() + timedelta(seconds=expires_in - 60)
 

app/templates/auth/totp_setup.html

@@ -0,0 +1,60 @@
+{% extends "base.html" %}
+{% block title %}Enable Two-Factor Auth{% endblock %}
+{% block page_title %}Enable Two-Factor Authentication{% endblock %}
+
+{% block content %}
+<div class="row justify-content-center">
+<div class="col-12 col-md-8 col-lg-6">
+
+<div class="pcard mb-3">
+    <div class="d-flex align-items-center gap-3 mb-4">
+        <div style="width:44px;height:44px;border-radius:11px;background:#ede9fe;color:#5b21b6;display:flex;align-items:center;justify-content:center;font-size:22px;flex-shrink:0;">
+            <i class="bi bi-shield-lock-fill"></i>
+        </div>
+        <div>
+            <div style="font-weight:700;font-size:15px;">Authenticator App Setup</div>
+            <div style="font-size:12px;color:var(--muted);">Scan the QR code with Google Authenticator, Authy, or any TOTP app.</div>
+        </div>
+    </div>
+
+    <!-- Step 1 -->
+    <div class="mb-4">
+        <div style="font-size:12px;font-weight:700;color:var(--muted);text-transform:uppercase;letter-spacing:.05em;margin-bottom:10px;">Step 1 — Scan QR code</div>
+        <div class="text-center">
+            <img src="{{ qr_data }}" alt="TOTP QR code" style="width:200px;height:200px;border:1px solid var(--border);border-radius:8px;padding:8px;background:#fff;">
+        </div>
+        <div class="mt-3 p-3" style="background:#f8fafc;border-radius:8px;font-size:12px;">
+            <div style="color:var(--muted);margin-bottom:4px;">Can't scan? Enter this key manually:</div>
+            <code class="mono" style="font-size:14px;font-weight:700;letter-spacing:.1em;word-break:break-all;">{{ secret }}</code>
+        </div>
+    </div>
+
+    <!-- Step 2 -->
+    <div>
+        <div style="font-size:12px;font-weight:700;color:var(--muted);text-transform:uppercase;letter-spacing:.05em;margin-bottom:10px;">Step 2 — Confirm with a code</div>
+        <form method="POST" autocomplete="off">
+            <input type="hidden" name="csrf_token" value="{{ csrf_token() }}">
+            <div class="d-flex gap-2">
+                <input type="text" name="code" inputmode="numeric" pattern="[0-9 ]*"
+                       maxlength="7" autofocus
+                       class="form-control mono text-center"
+                       style="font-size:22px;font-weight:700;letter-spacing:.15em;"
+                       placeholder="000000">
+                <button type="submit" class="btn btn-primary" style="white-space:nowrap;">
+                    Enable 2FA
+                </button>
+            </div>
+            <div style="font-size:11px;color:var(--muted);margin-top:6px;">
+                Enter the current 6-digit code shown in your authenticator app to confirm setup.
+            </div>
+        </form>
+    </div>
+</div>
+
+<div class="text-center">
+    <a href="{{ url_for('settings.index') }}" class="btn btn-outline-secondary btn-sm">Cancel</a>
+</div>
+
+</div>
+</div>
+{% endblock %}

app/templates/auth/totp_verify.html

@@ -0,0 +1,41 @@
+{% extends "base.html" %}
+{% block title %}Two-Factor Verification{% endblock %}
+{% block page_title %}Two-Factor Verification{% endblock %}
+
+{% block content %}
+<div class="row justify-content-center mt-5">
+<div class="col-12 col-sm-8 col-md-5 col-lg-4">
+<div class="pcard">
+    <div class="text-center mb-4">
+        <div style="width:56px;height:56px;border-radius:14px;background:#ede9fe;color:#5b21b6;display:flex;align-items:center;justify-content:center;font-size:26px;margin:0 auto 12px;">
+            <i class="bi bi-shield-lock-fill"></i>
+        </div>
+        <h5 style="font-weight:700;">Verify your identity</h5>
+        <p class="text-muted" style="font-size:13px;">Enter the 6-digit code from your authenticator app.</p>
+    </div>
+
+    {% if error %}
+    <div class="alert alert-danger" style="font-size:13px;">{{ error }}</div>
+    {% endif %}
+
+    <form method="POST" autocomplete="off">
+        <input type="hidden" name="csrf_token" value="{{ csrf_token() }}">
+        <div class="mb-3">
+            <input type="text" name="code" inputmode="numeric" pattern="[0-9 ]*"
+                   maxlength="7" autofocus
+                   class="form-control text-center mono"
+                   style="font-size:28px;font-weight:700;letter-spacing:.2em;"
+                   placeholder="000000">
+        </div>
+        <button type="submit" class="btn btn-primary w-100">Verify</button>
+    </form>
+
+    <div class="text-center mt-3">
+        <a href="{{ url_for('auth.login') }}" style="font-size:12px;color:var(--muted);">
+            ← Back to login
+        </a>
+    </div>
+</div>
+</div>
+</div>
+{% endblock %}

app/templates/dashboard/index.html

@@ -18,6 +18,15 @@
 {% endblock %}
 
 {% block content %}
+{% if schwab_warning %}
+<div class="alert alert-warning d-flex align-items-center gap-2 mb-3" style="font-size:13px;">
+    <i class="bi bi-exclamation-triangle-fill flex-shrink-0"></i>
+    <div>
+        {{ schwab_warning }}
+        <a href="{{ url_for('schwab.connect') }}" class="alert-link ms-1">Reconnect now →</a>
+    </div>
+</div>
+{% endif %}
 <div class="d-flex align-items-start align-items-sm-center justify-content-between flex-wrap gap-2 mb-4">
     <div>
         <h5 class="mb-0 fw-semibold">{{ period_label }}</h5>

app/templates/settings/index.html

@@ -61,6 +61,47 @@
     </div>
 </div>
 
+<!-- Two-Factor Authentication -->
+<div class="row g-3 mt-1">
+    <div class="col-12 col-md-6">
+        <div class="pcard h-100">
+            <div class="d-flex align-items-center gap-2 mb-3">
+                <i class="bi bi-shield-lock-fill" style="font-size:1.3rem;color:#5b21b6;"></i>
+                <span class="pcard-title mb-0">Two-Factor Authentication (2FA)</span>
+                {% if current_user.totp_enabled %}
+                <span class="badge ms-auto" style="background:#d1fae5;color:#065f46;font-size:11px;">Enabled</span>
+                {% else %}
+                <span class="badge ms-auto" style="background:#fee2e2;color:#991b1b;font-size:11px;">Disabled</span>
+                {% endif %}
+            </div>
+            {% if current_user.totp_enabled %}
+            <p style="font-size:13px;color:var(--muted);">
+                2FA is active. Each login requires a 6-digit code from your authenticator app.
+            </p>
+            <form method="POST" action="{{ url_for('auth.totp_disable') }}"
+                  onsubmit="return confirm('Disable two-factor authentication?\nYou will only need a password to log in.')">
+                <input type="hidden" name="csrf_token" value="{{ csrf_token() }}">
+                <div class="mb-2">
+                    <input type="password" name="password" class="form-control form-control-sm"
+                           placeholder="Confirm your password" required style="max-width:280px;">
+                </div>
+                <button type="submit" class="btn btn-sm btn-outline-danger">
+                    <i class="bi bi-shield-x me-1"></i>Disable 2FA
+                </button>
+            </form>
+            {% else %}
+            <p style="font-size:13px;color:var(--muted);">
+                Protect your account with a time-based one-time password (TOTP).
+                Use Google Authenticator, Authy, or any compatible app.
+            </p>
+            <a href="{{ url_for('auth.totp_setup') }}" class="btn btn-sm btn-outline-primary">
+                <i class="bi bi-shield-plus me-1"></i>Enable 2FA
+            </a>
+            {% endif %}
+        </div>
+    </div>
+</div>
+
 <!-- Maintenance + Email Alerts -->
 <div class="row g-3 mt-1">
     <div class="col-12 col-md-6">

requirements.txt

@@ -16,3 +16,7 @@ requests==2.32.3
 cryptography==44.0.2
 python-dateutil==2.9.0
 pdfplumber==0.11.4
+# Security
+flask-limiter==3.5.0
+pyotp==2.9.0
+qrcode==7.4.2

scripts/add_security_columns.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""
+Migration: add security columns.
+  users                   — totp_secret, totp_enabled
+  schwab_connections      — refresh_token_expires_at
+
+Run once:  python scripts/add_security_columns.py
+Safe to re-run — skips columns that already exist.
+"""
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from app import create_app
+from app.extensions import db
+
+app = create_app()
+
+COLUMNS = [
+    ("users", "totp_secret",
+     "VARCHAR(64) NULL DEFAULT NULL"),
+    ("users", "totp_enabled",
+     "TINYINT(1) NOT NULL DEFAULT 0"),
+    ("schwab_connections", "refresh_token_expires_at",
+     "DATETIME NULL DEFAULT NULL"),
+]
+
+with app.app_context():
+    with db.engine.connect() as conn:
+        for table, column, definition in COLUMNS:
+            exists = conn.execute(db.text(
+                "SELECT COUNT(*) FROM information_schema.columns "
+                "WHERE table_schema = DATABASE() "
+                "AND table_name = :t AND column_name = :c"
+            ), {"t": table, "c": column}).scalar()
+
+            if exists:
+                print(f"  {table}.{column} already exists — skipped.")
+            else:
+                conn.execute(db.text(
+                    f"ALTER TABLE {table} ADD COLUMN {column} {definition}"
+                ))
+                conn.commit()
+                print(f"  Added {table}.{column}.")
+
+print("Done.")

scripts/sync_schwab.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""
+Schwab daily auto-sync — balance, positions, and transactions.
+
+Add to crontab (crontab -e):
+    # Sync Schwab every day at 7 AM
+    0 7 * * * /home/pfm/venv/bin/python /home/pfm/web/scripts/sync_schwab.py >> /home/pfm/web/logs/sync_schwab.log 2>&1
+
+Or run manually:
+    python scripts/sync_schwab.py
+"""
+import sys
+import os
+import logging
+from datetime import datetime
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from app import create_app
+from app.extensions import db
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s|%(levelname)s|%(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S',
+)
+log = logging.getLogger(__name__)
+
+app = create_app()
+
+with app.app_context():
+    from app.models.schwab_connection import SchwabConnection, SchwabAccount
+    from app.services.schwab_service import (
+        sync_account_snapshot, sync_preview, import_transactions,
+    )
+
+    connection = SchwabConnection.query.filter_by(is_active=True).first()
+    if not connection:
+        log.warning('[sync_schwab] No active Schwab connection — skipping.')
+        sys.exit(0)
+
+    # Warn if refresh token is close to expiry
+    if connection.refresh_token_expires_at:
+        days_left = (connection.refresh_token_expires_at - datetime.utcnow()).days
+        if days_left <= 2:
+            log.warning('[sync_schwab] Schwab refresh token expires in %d day(s)! '
+                        'Log in and reconnect at /schwab/connect.', days_left)
+
+    accounts = SchwabAccount.query.filter(
+        SchwabAccount.pfm_account_id != None,
+        SchwabAccount.is_active == True,
+    ).all()
+
+    if not accounts:
+        log.warning('[sync_schwab] No mapped Schwab accounts found.')
+        sys.exit(0)
+
+    log.info('[sync_schwab] Syncing %d account(s)…', len(accounts))
+
+    for sa in accounts:
+        sa.connection = connection   # always use active connection
+
+        # ── 1. Balance + positions ────────────────────────────────────────────
+        try:
+            bal_updated, pos_synced = sync_account_snapshot(sa)
+            log.info('[sync_schwab] %s — balance updated=%s positions=%d',
+                     sa.account_name, bal_updated, pos_synced)
+        except Exception as e:
+            log.error('[sync_schwab] snapshot failed for %s: %s', sa.account_name, e,
+                      exc_info=True)
+
+        # ── 2. New transactions ───────────────────────────────────────────────
+        try:
+            preview = sync_preview(sa)
+            if preview:
+                imported, skipped = import_transactions(preview, sa)
+                log.info('[sync_schwab] %s — transactions: imported=%d skipped=%d',
+                         sa.account_name, imported, skipped)
+            else:
+                log.info('[sync_schwab] %s — no new transactions.', sa.account_name)
+        except Exception as e:
+            log.error('[sync_schwab] transaction sync failed for %s: %s',
+                      sa.account_name, e, exc_info=True)
+
+    log.info('[sync_schwab] Done.')