rbox/mywebdav/main.py

import argparse
import uvicorn
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from tortoise.contrib.fastapi import register_tortoise
from .settings import settings
from .routers import (
    auth,
    users,
    folders,
    files,
    shares,
    search,
    admin,
    starred,
    billing,
    admin_billing,
)
from . import webdav
from .schemas import ErrorResponse
from .middleware.usage_tracking import UsageTrackingMiddleware

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up...")
    logger.info("Database connected.")
    from .billing.scheduler import start_scheduler
    from .billing.models import PricingConfig
    from .mail import email_service

    start_scheduler()
    logger.info("Billing scheduler started")
    await email_service.start()
    logger.info("Email service started")
    pricing_count = await PricingConfig.all().count()
    if pricing_count == 0:
        from decimal import Decimal

        await PricingConfig.create(
            config_key="storage_per_gb_month",
            config_value=Decimal("0.0045"),
            description="Storage cost per GB per month",
            unit="per_gb_month",
        )
        await PricingConfig.create(
            config_key="bandwidth_egress_per_gb",
            config_value=Decimal("0.009"),
            description="Bandwidth egress cost per GB",
            unit="per_gb",
        )
        await PricingConfig.create(
            config_key="bandwidth_ingress_per_gb",
            config_value=Decimal("0.0"),
            description="Bandwidth ingress cost per GB (free)",
            unit="per_gb",
        )
        await PricingConfig.create(
            config_key="free_tier_storage_gb",
            config_value=Decimal("15"),
            description="Free tier storage in GB",
            unit="gb",
        )
        await PricingConfig.create(
            config_key="free_tier_bandwidth_gb",
            config_value=Decimal("15"),
            description="Free tier bandwidth in GB per month",
            unit="gb",
        )
        await PricingConfig.create(
            config_key="tax_rate_default",
            config_value=Decimal("0.0"),
            description="Default tax rate (0 = no tax)",
            unit="percentage",
        )
        logger.info("Default pricing configuration initialized")

    yield

    from .billing.scheduler import stop_scheduler

    stop_scheduler()
    logger.info("Billing scheduler stopped")
    await email_service.stop()
    logger.info("Email service stopped")
    print("Shutting down...")


app = FastAPI(
    title="MyWebdav Cloud Storage",
    description="A commercial cloud storage web application",
    version="0.1.0",
    lifespan=lifespan,
)

app.include_router(auth.router)
app.include_router(users.router)
app.include_router(folders.router)
app.include_router(files.router)
app.include_router(shares.router)
app.include_router(search.router)
app.include_router(admin.router)
app.include_router(starred.router)
app.include_router(billing.router)
app.include_router(admin_billing.router)
app.include_router(webdav.router)

app.add_middleware(UsageTrackingMiddleware)

app.mount("/static", StaticFiles(directory="static"), name="static")

register_tortoise(
    app,
    db_url=settings.DATABASE_URL,
    modules={"models": ["mywebdav.models"], "billing": ["mywebdav.billing.models"]},
    generate_schemas=True,
    add_exception_handlers=True,
)


@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    logger.error(
        f"HTTPException: {exc.status_code} - {exc.detail} for URL: {request.url}"
    )
    return JSONResponse(
        status_code=exc.status_code,
        content=ErrorResponse(code=exc.status_code, message=exc.detail).model_dump(),
    )


@app.get("/", response_class=HTMLResponse)  # Change response_class to HTMLResponse
async def read_root():
    with open("static/index.html", "r") as f:
        return f.read()


def main():
    parser = argparse.ArgumentParser(description="Run the MyWebdav application.")
    parser.add_argument(
        "--host", type=str, default="0.0.0.0", help="Host address to bind to"
    )
    parser.add_argument("--port", type=int, default=8000, help="Port to listen on")
    args = parser.parse_args()

    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()