rbox/rbox/main.py

import argparse
import uvicorn
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, status, 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

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
    start_scheduler()
    logger.info("Billing scheduler 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")
    print("Shutting down...")

app = FastAPI(
    title="RBox Cloud Storage",
    description="A self-hosted 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)

from .middleware.usage_tracking import UsageTrackingMiddleware

app.add_middleware(UsageTrackingMiddleware)

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

register_tortoise(
    app,
    db_url=settings.DATABASE_URL,
    modules={
        "models": ["rbox.models"],
        "billing": ["rbox.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).dict(),
    )


@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 RBox 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()