Files
SmartUp/backend/app/database.py
T
SmartUp Developer ad16618406 fix: address multiple code audit findings
- CORS: replace wildcard with explicit origin list from CORS_ORIGINS env
- Auth: enforce strong defaults, JWT blacklist (RevokedToken model), login rate limiting
- Auth: validate password length before bcrypt (72-byte limit)
- Scheduler: single-threaded worker to mitigate SQLite write contention
- Scheduler: graceful shutdown (wait=True)
- Snapshots: add prune_snapshots() with configurable retention count
- Storage: isolate localStorage keys via VITE_APP_KEY prefix
- Config: add cors_origins, login_rate_limit, snapshot_retention_count settings
2026-05-17 10:52:18 +08:00

79 lines
3.3 KiB
Python

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import sessionmaker, DeclarativeBase
from app.config import get_settings
settings = get_settings()
engine = create_engine(
settings.database_url,
connect_args={"check_same_thread": False},
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
class Base(DeclarativeBase):
pass
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
def init_db():
"""Create all tables."""
# import models so SQLAlchemy registers them
from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token # noqa: F401
Base.metadata.create_all(bind=engine)
_migrate_custom_pages()
def _migrate_custom_pages():
"""Apply small SQLite-safe migrations for deployments without Alembic."""
inspector = inspect(engine)
if "custom_pages" not in inspector.get_table_names():
return
columns = {col["name"] for col in inspector.get_columns("custom_pages")}
with engine.begin() as conn:
if "access_mode" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN access_mode VARCHAR(32) NOT NULL DEFAULT 'direct'"))
conn.execute(text("UPDATE custom_pages SET access_mode = CASE WHEN use_proxy = 1 THEN 'proxy' ELSE 'direct' END"))
if "login_username" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_username VARCHAR(255)"))
if "login_password" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_password TEXT"))
if "login_username_selector" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_username_selector VARCHAR(512)"))
if "login_password_selector" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_password_selector VARCHAR(512)"))
if "login_submit_selector" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_submit_selector VARCHAR(512)"))
if "login_autofill_enabled" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_autofill_enabled BOOLEAN NOT NULL DEFAULT 0"))
if "login_autofill_backfilled_at" not in columns:
conn.execute(text("ALTER TABLE custom_pages ADD COLUMN login_autofill_backfilled_at DATETIME"))
conn.execute(
text(
"UPDATE custom_pages "
"SET login_autofill_enabled = 1, login_autofill_backfilled_at = CURRENT_TIMESTAMP "
"WHERE login_autofill_enabled = 0 "
"AND NULLIF(TRIM(login_username), '') IS NOT NULL "
"AND NULLIF(TRIM(login_password), '') IS NOT NULL"
)
)
conn.execute(
text(
"UPDATE custom_pages "
"SET access_mode = 'remote_browser', use_proxy = 0 "
"WHERE url LIKE :host OR url LIKE :host_slash OR url LIKE :host_port"
),
{
"host": "%://170.106.100.210",
"host_slash": "%://170.106.100.210/%",
"host_port": "%://170.106.100.210:%",
},
)