From 9e1faedb1a96ab996616464de4d1251dce0e0775 Mon Sep 17 00:00:00 2001 From: Thephaseless Date: Fri, 18 Oct 2024 13:55:32 +0000 Subject: [PATCH] add exception handling --- main.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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,