diff --git a/src/api/src/main.py b/src/api/src/main.py index 0342517..9850c08 100644 --- a/src/api/src/main.py +++ b/src/api/src/main.py @@ -1,8 +1,16 @@ +from typing import Optional from pathlib import Path from fastapi import FastAPI, HTTPException from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse from ebooklib import epub -from create_book import retrieve_story, set_cover, set_metadata, add_chapters, slugify +from create_book import ( + retrieve_story, + set_cover, + set_metadata, + add_chapters, + slugify, + wp_get_cookies, +) import tempfile from io import BytesIO from fastapi.staticfiles import StaticFiles @@ -17,29 +25,47 @@ def home(): @app.get("/download/{story_id}") -async def download_book(story_id: int, download_images: bool = False): - data = await retrieve_story(story_id) +async def download_book( + story_id: int, + download_images: bool = False, + username: Optional[str] = None, + password: Optional[str] = None, +): + if username and not password or password and not username: + return HTMLResponse( + status_code=422, + content='Include both the username _and_ password, or neither. Support is available on the Discord', + ) + + if username and password: + try: + cookies = await wp_get_cookies(username=username, password=password) + except ValueError: + return HTMLResponse( + status_code=403, + content='Incorrect Username and/or Password. Support is available on the Discord', + ) + else: + cookies = None + + data = await retrieve_story(story_id, cookies=cookies) book = epub.EpubBook() - # Metadata and Cover are updated try: set_metadata(book, data) except KeyError: - # raise HTTPException( - # status_code=404, - # detail='Story not found. Check the ID - Support is available on the Discord', - # ) - # return FileResponse(BUILD_PATH / "index.html", status_code=404) return HTMLResponse( status_code=404, content='Story not found. Check the ID - Support is available on the Discord', ) - await set_cover(book, data) + await set_cover(book, data, cookies=cookies) # print("Metadata Downloaded") # Chapters are downloaded - async for title in add_chapters(book, data, download_images=download_images): + async for title in add_chapters( + book, data, download_images=download_images, cookies=cookies + ): # print(f"Part ({title}) downloaded") ...