From 308afde25fd684cc23df549661ff681ae6c9bac7 Mon Sep 17 00:00:00 2001 From: AaronBenDaniel <144371000+AaronBenDaniel@users.noreply.github.com> Date: Sun, 24 Nov 2024 21:42:07 -0500 Subject: [PATCH] fix(api): Handle invalid part IDs --- src/api/src/main.py | 39 ++++++++++++++++++++++++++------------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/src/api/src/main.py b/src/api/src/main.py index cca7941..21347d5 100644 --- a/src/api/src/main.py +++ b/src/api/src/main.py @@ -74,6 +74,13 @@ async def handle_download( metadata = await retrieve_story(story_id, cookies) book = epub.EpubBook() + if not metadata: + # Invalid ID + return HTMLResponse( + status_code=404, + content='The story you tried to download does not exist or has been deleted. Support is available on the Discord', + ) + set_metadata(book, metadata) await set_cover(book, metadata, cookies=cookies) @@ -101,19 +108,25 @@ async def handle_download( }, ) - except KeyError: - # Invalid ID - return HTMLResponse( - status_code=404, - content='The story you tried to download does not exist or has been deleted. Support is available on the Discord', - ) - - except ClientResponseError: - # Rate-limit by Wattpad - return HTMLResponse( - status_code=429, - content='Unfortunately, the downloader got rate-limited by Wattpad. Please try again later. Support is available on the Discord', - ) + except ClientResponseError as error: + if error.status == 429: + # Rate-limit by Wattpad + return HTMLResponse( + status_code=429, + content='Unfortunately, the downloader got rate-limited by Wattpad. Please try again later. Support is available on the Discord', + ) + elif error.status == 400: + # Invalid ID + return HTMLResponse( + status_code=404, + content='The story you tried to download does not exist or has been deleted. Support is available on the Discord', + ) + else: + # Unhandled error + return HTMLResponse( + status_code=500, + content='Something went wrong. Support is available on the Discord', + ) app.mount("/", StaticFiles(directory=BUILD_PATH), "static")