diff --git a/pyproject.toml b/pyproject.toml index 46f50fe..0c200a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,8 @@ ignore = [ "G004", "ANN001", "ANN204", - "ANN206", + "CPY001", + "BLE001", ] select = ["ALL"] extend-safe-fixes = ["D415"] diff --git a/src/endpoints.py b/src/endpoints.py index 346b1b4..dd9785a 100644 --- a/src/endpoints.py +++ b/src/endpoints.py @@ -69,12 +69,53 @@ async def health_check(sb: BrowserDep): async def read_item(request: LinkRequest, dep: BrowserDep) -> LinkResponse: """Handle POST requests.""" start_time = int(time.time() * 1000) - timer = TimeoutTimer(duration=request.max_timeout) - request.url = request.url.replace('"', "").strip() + final_url = await setup_routes(request, dep) + + try: + challenge_detected, page_html, page_request, status = ( + await _navigate_and_solve(dep, request, timer) + ) + except (TimeoutError, PlaywrightTimeoutError) as e: + logger.error("Timed out while loading the page or solving the challenge") + raise HTTPException( + status_code=408, + detail="Timed out while loading the page or solving the challenge", + ) from e + + cookies = await dep.context.cookies() + content_type, response_content = await build_response_content( + dep, request, page_request, + challenge_detected=challenge_detected, + page_html=page_html, + ) + + return LinkResponse( + message="Success", + solution=Solution( + user_agent=await dep.page.evaluate("navigator.userAgent"), + url=final_url if final_url is not None else dep.page.url, + status=status, + cookies=cookies, + headers=page_request.headers if page_request else {}, + response=response_content, + content_type=content_type, + ), + start_timestamp=start_time, + ) + + +async def setup_routes(request: LinkRequest, dep: BrowserDep) -> str | None: + """ + Install request routes for media blocking and CSP stripping. + + Returns the final URL captured during navigation; callers read it after + the page settles. + """ if request.block_media: + async def block_media_route(route) -> None: if route.request.resource_type in ("image", "media", "font"): await route.abort() @@ -109,63 +150,7 @@ async def strip_csp_route(route) -> None: ) await dep.page.route("**/*", strip_csp_route) - - try: - challenge_detected, page_html, page_request, status = ( - await _navigate_and_solve(dep, request, timer) - ) - except (TimeoutError, PlaywrightTimeoutError) as e: - logger.error("Timed out while loading the page or solving the challenge") - raise HTTPException( - status_code=408, - detail="Timed out while loading the page or solving the challenge", - ) from e - - cookies = await dep.context.cookies() - - content_type = "text/html" - response_content = "" - - if request.return_only_cookies: - response_content = "" - elif ( - page_request - and page_request.headers.get("content-type", "").startswith( - "application/pdf" - ) - ): - content_type = "application/pdf" - try: - fetch_response = await dep.page.request.fetch(dep.page.url) - response_content = base64.b64encode( - await fetch_response.body() - ).decode("ascii") - except Exception: - logger.exception( - "Failed to fetch PDF bytes, falling back to viewer HTML" - ) - content_type = "text/html" - response_content = await dep.page.content() - else: - response_content = ( - page_html - if page_html is not None and not challenge_detected - else await dep.page.content() - ) - - return LinkResponse( - message="Success", - solution=Solution( - user_agent=await dep.page.evaluate("navigator.userAgent"), - url=final_url if final_url is not None else dep.page.url, - status=status, - cookies=cookies, - headers=page_request.headers if page_request else {}, - response=response_content, - content_type=content_type, - ), - start_timestamp=start_time, - ) + return final_url async def _navigate_and_solve( @@ -173,6 +158,7 @@ async def _navigate_and_solve( request: LinkRequest, timer: TimeoutTimer, ) -> tuple[bool, str | None, object, HTTPStatus]: + """Navigate to the URL, then solve a challenge or wait for network idle.""" page_html: str | None = None page_request = await dep.page.goto( request.url, timeout=timer.remaining() * 1000 @@ -188,17 +174,16 @@ async def _navigate_and_solve( ) if not challenge_active: page_html = await dep.page.content() - try: - await dep.page.wait_for_load_state( - "networkidle", timeout=timer.remaining() * 1000 - ) - except PlaywrightTimeoutError: - logger.info( - "networkidle timed out after domcontentloaded; " - "continuing with loaded page" - ) + await _wait_for_networkidle(dep, timer) return False, page_html, page_request, status + await _solve_challenge(dep, timer) + status = HTTPStatus.OK + return True, page_html, page_request, status + + +async def _solve_challenge(dep: BrowserDep, timer: TimeoutTimer) -> None: + """Attempt to solve a detected Cloudflare interstitial challenge.""" logger.info("Challenge detected, attempting to solve...") await wait_for( dep.solver.solve_captcha( # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] @@ -209,6 +194,55 @@ async def _navigate_and_solve( ), timeout=timer.remaining(), ) - status = HTTPStatus.OK logger.debug("Challenge solved successfully.") - return True, page_html, page_request, status + + +async def _wait_for_networkidle(dep: BrowserDep, timer: TimeoutTimer) -> None: + """Wait for network idle, tolerating post-DOM-load stalls.""" + try: + await dep.page.wait_for_load_state( + "networkidle", timeout=timer.remaining() * 1000 + ) + except PlaywrightTimeoutError: + logger.info( + "networkidle timed out after domcontentloaded; " + "continuing with loaded page" + ) + + +async def build_response_content( + dep: BrowserDep, + request: LinkRequest, + page_request: object, + *, + challenge_detected: bool, + page_html: str | None, +) -> tuple[str, str]: + """Build (content_type, response_content) from the settled page.""" + if request.return_only_cookies: + return "text/html", "" + + if page_request and page_request.headers.get("content-type", "").startswith( + "application/pdf" + ): + return await _fetch_pdf_content(dep) + + response_content = ( + page_html + if page_html is not None and not challenge_detected + else await dep.page.content() + ) + return "text/html", response_content + + +async def _fetch_pdf_content(dep: BrowserDep) -> tuple[str, str]: + """Fetch raw PDF bytes as base64, falling back to viewer HTML on failure.""" + try: + fetch_response = await dep.page.request.fetch(dep.page.url) + response_content = base64.b64encode( + await fetch_response.body() + ).decode("ascii") + except Exception: + logger.exception("Failed to fetch PDF bytes, falling back to viewer HTML") + return "text/html", await dep.page.content() + return "application/pdf", response_content diff --git a/src/models.py b/src/models.py index 61ed825..1c8fe30 100644 --- a/src/models.py +++ b/src/models.py @@ -78,7 +78,7 @@ class LinkResponse(BaseModel): version: str = consts.VERSION @classmethod - def invalid(cls, url: str): + def invalid(cls, url: str) -> LinkResponse: """ Return an invalid LinkResponse with default error values. diff --git a/src/owui.py b/src/owui.py index b924749..3efd10a 100644 --- a/src/owui.py +++ b/src/owui.py @@ -69,7 +69,7 @@ async def load_urls( except PlaywrightTimeoutError: logger.debug("networkidle timed out for %s; extracting anyway", url) content = await _extract_content(dep.page) - except Exception as exc: # noqa: BLE001 + except Exception as exc: logger.warning("Failed to load %s: %s", url, exc) content = "" results.append(LoadResult(page_content=content, metadata={"source": url})) diff --git a/tests/main_test.py b/tests/main_test.py index ced8a60..85a97d3 100644 --- a/tests/main_test.py +++ b/tests/main_test.py @@ -1,3 +1,4 @@ +import base64 from http import HTTPStatus from json import JSONDecodeError from unittest.mock import AsyncMock, MagicMock @@ -83,7 +84,6 @@ def test_pdf_handling(): if solution.get("contentType") != "application/pdf": pytest.skip("Skipping PDF test - PDF bytes could not be fetched (upstream issue)") assert solution["response"] # non-empty base64 - import base64 decoded = base64.b64decode(solution["response"]) assert decoded[:5] == b"%PDF-"