File size: 2,857 Bytes
6165ba9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import logging
import os
import sys
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from .config import OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES, CLEANUP_INTERVAL
from .controllers.web_controller import router as web_router
from .utils import RateLimitMiddleware, ConcurrencyLimitMiddleware, RequestSizeLimitMiddleware, perform_cleanup
# Ensure registry is initialized
from .models import get_field_registry_manager
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("aibom_generator")
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
logger.info("Starting AI SBOM Generator WebApp")
try:
get_field_registry_manager() # Ensure registry is loaded
logger.info("Registry loaded successfully")
except Exception as e:
logger.error(f"Failed to load registry: {e}")
# Initial cleanup
try:
perform_cleanup(OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES)
except Exception as e:
logger.warning(f"Initial cleanup failed: {e}")
yield
# Shutdown (if needed)
app = FastAPI(title="AI SBOM Generator", lifespan=lifespan)
# --- Middleware ---
app.add_middleware(
RateLimitMiddleware,
rate_limit_per_minute=10,
rate_limit_window=60,
protected_routes=["/generate"]
)
app.add_middleware(
ConcurrencyLimitMiddleware,
max_concurrent_requests=5,
timeout=5.0,
protected_routes=["/generate"]
)
app.add_middleware(
RequestSizeLimitMiddleware,
max_content_length=1024*1024 # 1MB
)
# --- Cleanup Middleware ---
request_counter = 0
@app.middleware("http")
async def cleanup_middleware(request: Request, call_next):
global request_counter
request_counter += 1
if request_counter % CLEANUP_INTERVAL == 0:
try:
removed = perform_cleanup(OUTPUT_DIR, MAX_AGE_DAYS, MAX_FILES)
logger.info(f"Scheduled cleanup removed {removed} files")
except Exception as e:
logger.error(f"Error during scheduled cleanup: {e}")
response = await call_next(request)
return response
# --- Static Files ---
os.makedirs(OUTPUT_DIR, exist_ok=True)
app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
# Mount static files (CSS/JS)
os.makedirs("src/static", exist_ok=True)
app.mount("/static", StaticFiles(directory="src/static"), name="static")
# --- Routes ---
app.include_router(web_router)
if __name__ == "__main__":
import uvicorn
# Print clear access URL to avoid 0.0.0.0 confusion
print("🚀 Application ready! Access it at: http://localhost:8000")
uvicorn.run("src.main:app", host="0.0.0.0", port=8000, reload=True)
|