diff --git a/app.py b/app.py index 6f73f21..cfe316b 100644 --- a/app.py +++ b/app.py @@ -209,19 +209,36 @@ def _is_same_origin(): return referer.startswith(base_url) return False +def _ensure_rate_limits_table(conn): + """Create the rate_limits table if it doesn't exist (migration safety net).""" + conn.execute(''' + CREATE TABLE IF NOT EXISTS rate_limits ( + ip_address TEXT, + timestamp REAL, + PRIMARY KEY (ip_address, timestamp) + ) + ''') + conn.execute('CREATE INDEX IF NOT EXISTS idx_rate_limit_ts ON rate_limits(timestamp)') + conn.commit() + def _check_rate_limit(remote_ip, key_prefix='rl', window=600, limit=10): """Generic rate limiting via SQLite. Window is in seconds.""" now_ts = time.time() conn = get_db_connection() try: - count = conn.execute( - 'SELECT COUNT(*) FROM rate_limits WHERE ip_address = ? AND timestamp > ?', - (f"{key_prefix}:{remote_ip}", now_ts - window) - ).fetchone()[0] - + try: + count = conn.execute( + 'SELECT COUNT(*) FROM rate_limits WHERE ip_address = ? AND timestamp > ?', + (f"{key_prefix}:{remote_ip}", now_ts - window) + ).fetchone()[0] + except sqlite3.OperationalError: + # Table missing (pre-migration DB) — create it and proceed. + _ensure_rate_limits_table(conn) + count = 0 + if count >= limit: return False - + conn.execute('INSERT INTO rate_limits (ip_address, timestamp) VALUES (?, ?)', (f"{key_prefix}:{remote_ip}", now_ts)) conn.commit() return True