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 21347d5..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,10 +19,6 @@ from create_book import (
wp_get_cookies,
fetch_story_id,
)
-import tempfile
-from io import BytesIO
-from fastapi.staticfiles import StaticFiles
-from aiohttp import ClientResponseError
app = FastAPI()
BUILD_PATH = Path(__file__).parent / "build"
@@ -37,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,
@@ -45,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,
@@ -64,69 +87,42 @@ async def handle_download(
else:
cookies = None
- try:
- match mode:
- case DownloadMode.story:
- story_id = download_id
- case DownloadMode.part:
- story_id = await fetch_story_id(download_id, cookies)
+ match mode:
+ case DownloadMode.story:
+ story_id = download_id
+ case DownloadMode.part:
+ story_id = await fetch_story_id(download_id, cookies)
- metadata = await retrieve_story(story_id, cookies)
- book = epub.EpubBook()
+ 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',
- )
+ metadata = await retrieve_story(story_id, cookies)
+ set_metadata(book, metadata)
- set_metadata(book, metadata)
- await set_cover(book, metadata, cookies=cookies)
+ await set_cover(book, metadata, cookies=cookies)
- async for title in add_chapters(
- book, metadata, download_images=download_images, cookies=cookies
- ):
- ...
+ async for title in add_chapters(
+ book, metadata, download_images=download_images, cookies=cookies
+ ):
+ ...
- # Book is compiled
- temp_file = tempfile.NamedTemporaryFile(
- suffix=".epub", delete=True
- ) # Thanks https://stackoverflow.com/a/75398222
+ # Book is compiled
+ temp_file = tempfile.NamedTemporaryFile(
+ suffix=".epub", delete=True
+ ) # Thanks https://stackoverflow.com/a/75398222
- # create epub file
- epub.write_epub(temp_file, book, {})
+ # create epub file
+ epub.write_epub(temp_file, book, {})
- temp_file.file.seek(0)
- book_data = temp_file.file.read()
+ temp_file.file.seek(0)
+ book_data = temp_file.file.read()
- return StreamingResponse(
- BytesIO(book_data),
- media_type="application/epub+zip",
- headers={
- "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}_{"images" if download_images else ""}.epub"' # Thanks https://stackoverflow.com/a/72729058
- },
- )
-
- 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',
- )
+ return StreamingResponse(
+ BytesIO(book_data),
+ media_type="application/epub+zip",
+ headers={
+ "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}_{"images" if download_images else ""}.epub"' # Thanks https://stackoverflow.com/a/72729058
+ },
+ )
app.mount("/", StaticFiles(directory=BUILD_PATH), "static")