From 8bed900297798063d4fd7650fb56f569ceb20d29 Mon Sep 17 00:00:00 2001
From: NguyenND <ngodanguyen@gmail.com>
Date: Thu, 4 Jun 2026 10:48:31 -0400
Subject: [PATCH] 06/04 Optimize app

---
 app/config.py                  |  4 ++++
 app/extensions.py              |  2 +-
 app/routes/auth.py             | 12 +++++++++++-
 app/routes/schwab.py           |  2 +-
 app/services/schwab_service.py |  7 ++++---
 app/utils/crypto.py            |  6 +-----
 6 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/app/config.py b/app/config.py
index d325e95..caa3dbc 100644
--- a/app/config.py
+++ b/app/config.py
@@ -39,6 +39,10 @@ class Config:
     # Sentry error monitoring (optional — leave blank to disable)
     SENTRY_DSN = os.environ.get('SENTRY_DSN', '')
 
+    # Rate limiter shared storage (use Redis in production to share limits across Gunicorn workers)
+    # Example: RATELIMIT_STORAGE_URI=redis://localhost:6379
+    RATELIMIT_STORAGE_URI = os.environ.get('RATELIMIT_STORAGE_URI', 'memory://')
+
     # Session idle timeout in minutes (default 60)
     SESSION_IDLE_MINUTES = int(os.environ.get('SESSION_IDLE_MINUTES', 60))
 
diff --git a/app/extensions.py b/app/extensions.py
index 68e4b11..61b75ab 100644
--- a/app/extensions.py
+++ b/app/extensions.py
@@ -9,7 +9,7 @@ db = SQLAlchemy()
 login_manager = LoginManager()
 migrate = Migrate()
 csrf = CSRFProtect()
-limiter = Limiter(key_func=get_remote_address, storage_uri='memory://')
+limiter = Limiter(key_func=get_remote_address)   # storage backend set via RATELIMIT_STORAGE_URI in app config
 
 login_manager.login_view = 'auth.login'
 login_manager.login_message = 'Please log in to access this page.'
diff --git a/app/routes/auth.py b/app/routes/auth.py
index cb39e63..eb98af3 100644
--- a/app/routes/auth.py
+++ b/app/routes/auth.py
@@ -75,7 +75,7 @@ def logout():
 # ── TOTP: second-factor verification ─────────────────────────────────────────
 
 @auth_bp.route('/totp/verify', methods=['GET', 'POST'])
-@limiter.limit('10 per minute')
+@limiter.limit('10 per minute; 30 per hour')
 def totp_verify():
     pending_id = session.get('_totp_pending_id')
     if not pending_id:
@@ -102,8 +102,17 @@ def totp_verify():
             audit('login_success_2fa', f'user={user.username}')
             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)
+        attempts = session.get('_totp_attempts', 0) + 1
+        session['_totp_attempts'] = attempts
+        if attempts >= 5:
+            session.pop('_totp_pending_id', None)
+            session.pop('_totp_attempts', None)
+            audit('login_failed_2fa', f'user={user.username} — max attempts reached')
+            flash('Too many failed attempts. Please log in again.', 'danger')
+            return redirect(url_for('auth.login'))
         error = 'Invalid code — please try again.'
 
     return render_template('auth/totp_verify.html', error=error)
@@ -113,6 +122,7 @@ def totp_verify():
 
 @auth_bp.route('/totp/setup', methods=['GET', 'POST'])
 @login_required
+@limiter.limit('10 per minute')
 def totp_setup():
     user = current_user
 
diff --git a/app/routes/schwab.py b/app/routes/schwab.py
index 16bf9e4..9f23ad3 100644
--- a/app/routes/schwab.py
+++ b/app/routes/schwab.py
@@ -149,7 +149,7 @@ def map_accounts():
         for sa in schwab_accounts:
             val = request.form.get(f'pfm_account_{sa.id}', '')
             if val == 'new':
-                pfm_type = ACCOUNT_TYPE_MAP.get(sa.account_type, 'investment')
+                pfm_type = ACCOUNT_TYPE_MAP.get(sa.account_type, 'other')
                 new_acct = Account(
                     name=sa.account_name,
                     account_type=pfm_type,
diff --git a/app/services/schwab_service.py b/app/services/schwab_service.py
index 1b27130..15276e2 100644
--- a/app/services/schwab_service.py
+++ b/app/services/schwab_service.py
@@ -153,9 +153,10 @@ def _apply_token_data(connection, data):
     """Write token fields from token response onto connection model."""
     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)
+        connection.refresh_token = data['refresh_token']
+    # Reset the 7-day expiry window on every successful token exchange so the
+    # dashboard warning doesn't trigger prematurely when the token isn't rotated.
+    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)
 
diff --git a/app/utils/crypto.py b/app/utils/crypto.py
index 78c334d..1d4c0b4 100644
--- a/app/utils/crypto.py
+++ b/app/utils/crypto.py
@@ -35,11 +35,7 @@ class EncryptedText(types.TypeDecorator):
     def process_bind_param(self, value, dialect):
         if value is None:
             return None
-        try:
-            return _fernet().encrypt(value.encode()).decode()
-        except Exception as exc:
-            log.warning('[crypto] encrypt failed: %s', exc)
-            return value   # store plaintext rather than lose data
+        return _fernet().encrypt(value.encode()).decode()
 
     def process_result_value(self, value, dialect):
         if value is None: