diff --git a/main.py b/main.py index 5071ae7..1ceb5cd 100644 --- a/main.py +++ b/main.py @@ -6,7 +6,7 @@ import time import uvicorn import uvicorn.config -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.responses import RedirectResponse from src.models.requests import LinkRequest, LinkResponse @@ -36,12 +36,15 @@ async def read_item(request: LinkRequest): timeout = request.maxTimeout if timeout == 0: timeout = None - - challenged = await asyncio.wait_for(bypass_cloudflare(page), timeout=timeout) + try: + challenged = await asyncio.wait_for(bypass_cloudflare(page), timeout=timeout) + except asyncio.TimeoutError as e: + await page.save_screenshot("errors/error.png") + raise HTTPException(detail="Timeout", status_code=408) from e logger.info(f"Got webpage: {request.url}") - response = await LinkResponse.create( + response = LinkResponse.create( page=page, start_timestamp=start_time, challenged=challenged,