diff --git a/.gitignore b/.gitignore index 1c01200..20ecf1f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ venv data *ipynb build +.vscode diff --git a/src/api/src/create_book.py b/src/api/src/create_book.py index 1d8d75d..fb6c6c5 100644 --- a/src/api/src/create_book.py +++ b/src/api/src/create_book.py @@ -111,9 +111,6 @@ async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> dict: async with session.get( f"https://www.wattpad.com/api/v3/stories/{story_id}?fields=tags,id,title,createDate,modifyDate,language(name),description,completed,mature,url,isPaywalled,user(username),parts(id,title),cover" ) as response: - if not response.ok: - if response.status in [404, 400]: - return {} response.raise_for_status() body = await response.json() @@ -132,9 +129,6 @@ async def fetch_part_content(part_id: int, cookies: Optional[dict] = None) -> st async with session.get( f"https://www.wattpad.com/apiv2/?m=storytext&id={part_id}" ) as response: - if not response.ok: - if response.status in [404, 400]: - return "" response.raise_for_status() body = await response.text() @@ -151,9 +145,6 @@ async def fetch_cover(url: str, cookies: Optional[dict] = None) -> bytes: else ClientSession(headers=headers, cookies=cookies) ) as session: # Don't cache requests with Cookies. async with session.get(url) as response: - if not response.ok: - if response.status in [404, 400]: - return bytes() response.raise_for_status() body = await response.read() diff --git a/src/api/src/main.py b/src/api/src/main.py index d9242f5..f1129d6 100644 --- a/src/api/src/main.py +++ b/src/api/src/main.py @@ -1,8 +1,14 @@ +"""WattpadDownloader API Server.""" + from typing import Optional +import tempfile from pathlib import Path +from io import BytesIO from enum import Enum -from fastapi import FastAPI +from aiohttp import ClientResponseError +from fastapi import FastAPI, Request from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse +from fastapi.staticfiles import StaticFiles from ebooklib import epub from create_book import ( retrieve_story, @@ -13,9 +19,6 @@ from create_book import ( wp_get_cookies, fetch_story_id, ) -import tempfile -from io import BytesIO -from fastapi.staticfiles import StaticFiles app = FastAPI() BUILD_PATH = Path(__file__).parent / "build" @@ -36,6 +39,28 @@ def home(): return FileResponse(BUILD_PATH / "index.html") +@app.exception_handler(ClientResponseError) +def download_error_handler(request: Request, exception: ClientResponseError): + match exception.status: + case 400 | 404: + return HTMLResponse( + status_code=404, + content='This story does not exist, or has been deleted. Support is available on the Discord', + ) + case 429: + # Rate-limit by Wattpad + return HTMLResponse( + status_code=429, + content='The website is overloaded. Please try again in a few minutes. Support is available on the Discord', + ) + case _: + # Unhandled error + return HTMLResponse( + status_code=500, + content='Something went wrong. Yell at me on the Discord', + ) + + @app.get("/download/{download_id}") async def handle_download( download_id: int, @@ -44,7 +69,6 @@ async def handle_download( username: Optional[str] = None, password: Optional[str] = None, ): - if username and not password or password and not username: return HTMLResponse( status_code=422, @@ -69,10 +93,11 @@ async def handle_download( case DownloadMode.part: story_id = await fetch_story_id(download_id, cookies) - metadata = await retrieve_story(story_id, cookies) book = epub.EpubBook() + metadata = await retrieve_story(story_id, cookies) set_metadata(book, metadata) + await set_cover(book, metadata, cookies=cookies) async for title in add_chapters(