Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,8 @@ ignore = [
"G004",
"ANN001",
"ANN204",
"ANN206",
"CPY001",
"BLE001",
]
select = ["ALL"]
extend-safe-fixes = ["D415"]
Expand Down
174 changes: 104 additions & 70 deletions src/endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -109,70 +150,15 @@ 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(
dep: BrowserDep,
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
Expand All @@ -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]
Expand All @@ -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
2 changes: 1 addition & 1 deletion src/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
2 changes: 1 addition & 1 deletion src/owui.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}))
Expand Down
2 changes: 1 addition & 1 deletion tests/main_test.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import base64
from http import HTTPStatus
from json import JSONDecodeError
from unittest.mock import AsyncMock, MagicMock
Expand Down Expand Up @@ -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-"
Expand Down
Loading