rbox/mywebdav/main.py

import argparse
import uvicorn
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, HTTPException, status
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse, Response, RedirectResponse
from fastapi.templating import Jinja2Templates
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,
    manage,
)
from . import webdav
from .schemas import ErrorResponse
from .middleware import UsageTrackingMiddleware, RateLimitMiddleware, SecurityHeadersMiddleware
from .monitoring import health_router

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

enterprise_components = {
    "db_manager": None,
    "cache": None,
    "lock_manager": None,
    "token_manager": None,
    "rate_limiter": None,
    "webdav_locks": None,
    "task_queue": None,
    "dal": None,
    "write_buffer": None,
}


async def init_enterprise_components():
    from .database.manager import init_db_manager
    from .cache.layer import init_cache
    from .concurrency.locks import init_lock_manager
    from .concurrency.webdav_locks import init_webdav_locks
    from .concurrency.atomic import init_atomic_ops
    from .auth_tokens import init_token_manager
    from .middleware.rate_limit import init_rate_limiter
    from .workers.queue import init_task_queue
    from .enterprise.dal import init_dal
    from .enterprise.write_buffer import init_write_buffer, WriteType

    enterprise_components["db_manager"] = await init_db_manager("data")
    logger.info("Enterprise: Per-user database manager initialized")

    enterprise_components["cache"] = await init_cache(maxsize=10000, flush_interval=30)
    logger.info("Enterprise: Cache layer initialized")

    enterprise_components["dal"] = await init_dal(cache_size=50000)
    logger.info("Enterprise: Data Access Layer initialized")

    enterprise_components["write_buffer"] = await init_write_buffer(
        flush_interval=60.0,
        max_buffer_size=1000,
    )
    from .activity import _flush_activities
    from .billing.usage_tracker import _flush_usage_records
    enterprise_components["write_buffer"].register_handler(WriteType.ACTIVITY, _flush_activities)
    enterprise_components["write_buffer"].register_handler(WriteType.USAGE_RECORD, _flush_usage_records)
    logger.info("Enterprise: Write buffer initialized (60s flush interval)")

    enterprise_components["lock_manager"] = await init_lock_manager(default_timeout=30.0)
    logger.info("Enterprise: Lock manager initialized")

    init_atomic_ops()
    logger.info("Enterprise: Atomic operations initialized")

    enterprise_components["webdav_locks"] = await init_webdav_locks(
        enterprise_components["db_manager"]
    )
    logger.info("Enterprise: Persistent WebDAV locks initialized")

    enterprise_components["token_manager"] = await init_token_manager(
        enterprise_components["db_manager"]
    )
    logger.info("Enterprise: Token manager with revocation initialized")

    enterprise_components["rate_limiter"] = await init_rate_limiter()
    logger.info("Enterprise: Rate limiter initialized")

    enterprise_components["task_queue"] = await init_task_queue(max_workers=4)
    logger.info("Enterprise: Background task queue initialized")


async def shutdown_enterprise_components():
    from .database.manager import shutdown_db_manager
    from .cache.layer import shutdown_cache
    from .concurrency.locks import shutdown_lock_manager
    from .concurrency.webdav_locks import shutdown_webdav_locks
    from .auth_tokens import shutdown_token_manager
    from .middleware.rate_limit import shutdown_rate_limiter
    from .workers.queue import shutdown_task_queue
    from .enterprise.dal import shutdown_dal
    from .enterprise.write_buffer import shutdown_write_buffer

    await shutdown_task_queue()
    logger.info("Enterprise: Task queue stopped")

    await shutdown_rate_limiter()
    logger.info("Enterprise: Rate limiter stopped")

    await shutdown_token_manager()
    logger.info("Enterprise: Token manager stopped")

    await shutdown_webdav_locks()
    logger.info("Enterprise: WebDAV locks stopped")

    await shutdown_lock_manager()
    logger.info("Enterprise: Lock manager stopped")

    await shutdown_write_buffer()
    logger.info("Enterprise: Write buffer stopped (flushed pending writes)")

    await shutdown_dal()
    logger.info("Enterprise: Data Access Layer stopped")

    await shutdown_cache()
    logger.info("Enterprise: Cache layer stopped")

    await shutdown_db_manager()
    logger.info("Enterprise: Database manager stopped")


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up...")

    await init_enterprise_components()
    logger.info("All enterprise components initialized")

    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")

    await shutdown_enterprise_components()
    logger.info("All enterprise components shut down")
    print("Shutting down...")


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

templates = Jinja2Templates(directory="mywebdav/templates")

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(manage.router)
app.include_router(webdav.router)
app.include_router(health_router)

app.add_middleware(SecurityHeadersMiddleware, enable_hsts=False, enable_csp=True)
app.add_middleware(RateLimitMiddleware)
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}"
    )

    headers = exc.headers

    # For WebDAV authentication challenges, we must return the headers
    # from the exception and an empty body. A JSON body will confuse WebDAV clients.
    if request.url.path.startswith("/webdav") and exc.status_code == status.HTTP_401_UNAUTHORIZED:
        return Response(status_code=exc.status_code, headers=headers)

    # For other WebDAV errors, it's better to return a text body than JSON
    if request.url.path.startswith("/webdav"):
        return Response(content=exc.detail, status_code=exc.status_code, headers=headers)

    return JSONResponse(
        status_code=exc.status_code,
        content=ErrorResponse(code=exc.status_code, message=exc.detail).model_dump(),
        headers=headers,
    )


@app.get("/", response_class=HTMLResponse)
async def splash_page(request: Request):
    return templates.TemplateResponse("splash.html", {"request": request})


@app.get("/features", response_class=HTMLResponse)
async def features_page(request: Request):
    return templates.TemplateResponse("features.html", {"request": request})


@app.get("/pricing", response_class=HTMLResponse)
async def pricing_page(request: Request):
    return templates.TemplateResponse("pricing.html", {"request": request})


@app.get("/support", response_class=HTMLResponse)
async def support_page(request: Request):
    return templates.TemplateResponse("support.html", {"request": request})


@app.get("/legal/{document_name}", response_class=HTMLResponse)
async def legal_document(request: Request, document_name: str):
    try:
        return templates.TemplateResponse(f"legal/{document_name}.html", {"request": request})
    except Exception:
        raise HTTPException(status_code=404, detail="Legal document not found")


@app.get("/login")
async def login_redirect():
    return RedirectResponse(url="/app", status_code=302)


@app.get("/app", response_class=HTMLResponse)
async def web_app():
    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=9004, help="Port to listen on")
    args = parser.parse_args()

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


if __name__ == "__main__":
    main()