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")
...