debt-manager/backend/app/main.py

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from contextlib import asynccontextmanager
from app.core.config import get_settings
from app.core.database import engine, async_session
from app.models.models import Base, User, Debt, RepaymentPlan, RepaymentRecord, Group
from app.core.security import hash_password
from app.api import auth, debts, repayments, admin, loans, notifications, twofa
from sqlalchemy import select, text

settings = get_settings()

METHOD_REV = {"equal_installment": "e", "interest_first": "p", "lump_sum": "i"}


async def seed_admin():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    async with async_session() as db:
        result = await db.execute(select(User))
        existing = result.scalars().all()
        if not existing:
            admin_user = User(
                email="admin",
                password_hash=hash_password(settings.ADMIN_PASSWORD),
                nickname="管理员",
                role="admin",
            )
            db.add(admin_user)
            test_user = User(
                email="test",
                password_hash=hash_password("test123"),
                nickname="测试用户",
                role="user",
            )
            db.add(test_user)
            await db.commit()


async def migrate_account_names():
    async with async_session() as db:
        result = await db.execute(select(User))
        users = result.scalars().all()
        for u in users:
            if "@" in u.email:
                new_email = u.email.split("@")[0]
                existing = await db.execute(select(User).where(User.email == new_email, User.id != u.id))
                if not existing.scalar_one_or_none():
                    u.email = new_email
        await db.commit()


async def migrate_add_real_email():
    async with engine.begin() as conn:
        try:
            await conn.execute(text(
                "ALTER TABLE users ADD COLUMN IF NOT EXISTS real_email VARCHAR(255) DEFAULT ''"
            ))
        except Exception:
            pass  # 表不存在时跳过，由 seed_admin 创建

async def migrate_add_2fa():
    async with engine.begin() as conn:
        await conn.execute(text(
            "ALTER TABLE users ADD COLUMN IF NOT EXISTS two_factor_secret VARCHAR(32)"
        ))
        await conn.execute(text(
            "ALTER TABLE users ADD COLUMN IF NOT EXISTS two_factor_enabled INTEGER DEFAULT 0"
        ))



async def migrate_debt_data():
    async with async_session() as db:
        result = await db.execute(select(Debt).where(Debt.schedule_data == None))
        debts = result.scalars().all()
        for debt in debts:
            plans_result = await db.execute(
                select(RepaymentPlan).where(RepaymentPlan.debt_id == debt.id).order_by(RepaymentPlan.due_date)
            )
            plans = plans_result.scalars().all()

            schedule = []
            paid_periods = []
            for p in plans:
                item = {
                    "p": p.principal_due / 100,
                    "date": p.due_date.strftime("%Y-%m-%d") if p.due_date else "",
                    "pay": p.total_due / 100,
                    "pr": p.principal_due / 100,
                    "it": p.interest_due / 100,
                    "rp": 0,
                }
                schedule.append(item)
                if p.status == "paid":
                    paid_periods.append(len(schedule))

            remaining = debt.amount / 100
            for item in schedule:
                item["rp"] = round(remaining, 2)
                remaining -= item["pr"]

            debt.schedule_data = schedule
            debt.paid_periods = paid_periods

        await db.commit()


@asynccontextmanager
async def lifespan(app: FastAPI):
    await seed_admin()
    await migrate_add_real_email()
    await migrate_add_2fa()
    await migrate_account_names()
    await migrate_debt_data()
    from app.services.scheduler import start_scheduler, check_upcoming_repayments
    start_scheduler()
    await check_upcoming_repayments()
    yield
    from app.services.scheduler import scheduler
    scheduler.shutdown()


app = FastAPI(title=settings.PROJECT_NAME, lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
    allow_headers=["Authorization", "Content-Type"],
)


@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    from app.core.rate_limit import is_ip_blocked, get_remaining_seconds
    ip = request.client.host if request.client else "unknown"
    if is_ip_blocked(ip):
        remaining = get_remaining_seconds(ip)
        return JSONResponse(status_code=429, content={"detail": f"登录失败次数过多，请 {remaining} 秒后再试", "retry_after": remaining})
    return await call_next(request)


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(status_code=500, content={"detail": "服务器内部错误"})

app.include_router(auth.router, prefix=settings.API_V1_PREFIX)
app.include_router(loans.router, prefix=settings.API_V1_PREFIX)
app.include_router(debts.router, prefix=settings.API_V1_PREFIX)
app.include_router(repayments.router, prefix=settings.API_V1_PREFIX)
app.include_router(admin.router, prefix=settings.API_V1_PREFIX)
app.include_router(notifications.router, prefix=settings.API_V1_PREFIX)
app.include_router(twofa.router, prefix=settings.API_V1_PREFIX)


@app.get("/health")
async def health():
    return {"status": "ok"}