onesvm-browser-server/bench/trafilatura-http/app.py

#!/usr/bin/env python3
"""Minimal trafilatura HTTP wrapper. POST /v1/read only (plus GET /health)."""
from __future__ import annotations

import ipaddress
import json
import socket
import threading
import traceback
from copy import deepcopy
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

import trafilatura
from trafilatura.metadata import extract_metadata
from trafilatura.settings import DEFAULT_CONFIG

SEM = threading.Semaphore(8)
MAX_DOWNLOAD_BYTES = 5_000_000
TIMEOUT_S = 15
LISTEN = ("0.0.0.0", 8080)

_CFG = deepcopy(DEFAULT_CONFIG)
_CFG["DEFAULT"]["DOWNLOAD_TIMEOUT"] = str(TIMEOUT_S)
_CFG["DEFAULT"]["MAX_FILE_SIZE"] = str(MAX_DOWNLOAD_BYTES)

_BLOCKED_NETS = (
    ipaddress.ip_network("0.0.0.0/8"),
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("::1/128"),
    ipaddress.ip_network("fc00::/7"),
    ipaddress.ip_network("fe80::/10"),
)
_BLOCKED_HOSTS = {
    "metadata.google.internal",
    "metadata.google.com",
    "kubernetes.default.svc",
}


def _forbidden_host(host: str) -> str | None:
    h = host.strip("[]").lower()
    if h in _BLOCKED_HOSTS or h.endswith(".internal"):
        return "blocked_metadata_host"
    try:
        infos = socket.getaddrinfo(h, None)
    except socket.gaierror as exc:
        return f"dns_fail:{exc}"
    for info in infos:
        ip = ipaddress.ip_address(info[4][0])
        if any(ip in net for net in _BLOCKED_NETS):
            return f"blocked_private_ip:{ip}"
    return None


def _guard_url(url: str) -> str | None:
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        return "scheme_not_http"
    if not parsed.hostname:
        return "no_host"
    return _forbidden_host(parsed.hostname)


def _extract(url: str, max_chars: int) -> dict:
    err = _guard_url(url)
    if err:
        return {"ok": False, "error": err, "fail_class": "fetch_fail"}
    downloaded = trafilatura.fetch_url(url, config=_CFG)
    if not downloaded:
        return {"ok": False, "error": "fetch_empty", "fail_class": "fetch_fail"}
    if len(downloaded.encode("utf-8", errors="replace")) > MAX_DOWNLOAD_BYTES:
        downloaded = downloaded.encode("utf-8", errors="replace")[:MAX_DOWNLOAD_BYTES].decode(
            "utf-8", errors="ignore"
        )
    markdown = trafilatura.extract(
        downloaded,
        output_format="markdown",
        include_comments=False,
        include_tables=True,
        favor_precision=True,
        config=_CFG,
    )
    if not markdown:
        return {"ok": False, "error": "empty_extract", "fail_class": "empty_extract", "title": ""}
    meta = extract_metadata(downloaded)
    title = (meta.title if meta and getattr(meta, "title", None) else "") or ""
    truncated = False
    if max_chars and len(markdown) > max_chars:
        markdown = markdown[:max_chars]
        truncated = True
    return {
        "ok": True,
        "title": title,
        "markdown": markdown,
        "char_count": len(markdown),
        "truncated": truncated,
        "url": url,
    }


class Handler(BaseHTTPRequestHandler):
    server_version = "trafilatura-http/s3b"

    def log_message(self, fmt: str, *args) -> None:
        print(f"{self.log_date_time_string()} {self.address_string()} {fmt % args}", flush=True)

    def _json(self, code: int, obj: dict) -> None:
        raw = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def do_GET(self) -> None:
        if self.path.split("?", 1)[0] in ("/health", "/"):
            self._json(200, {"ok": True, "service": "trafilatura-http"})
            return
        self._json(404, {"ok": False, "error": "not_found"})

    def do_POST(self) -> None:
        path = self.path.split("?", 1)[0]
        if path != "/v1/read":
            self._json(404, {"ok": False, "error": "not_found"})
            return
        length = int(self.headers.get("Content-Length") or 0)
        if length > 1_000_000:
            self._json(413, {"ok": False, "error": "body_too_large"})
            return
        try:
            payload = json.loads(self.rfile.read(length) or b"{}")
        except json.JSONDecodeError:
            self._json(400, {"ok": False, "error": "bad_json"})
            return
        url = (payload.get("url") or "").strip()
        max_chars = int(payload.get("max_chars") or 20000)
        if not url:
            self._json(400, {"ok": False, "error": "url_required"})
            return
        acquired = SEM.acquire(timeout=60)
        if not acquired:
            self._json(429, {"ok": False, "error": "queue_timeout", "fail_class": "timeout"})
            return
        try:
            result = _extract(url, max_chars)
            self._json(200 if result.get("ok") else 200, result)
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            self._json(500, {"ok": False, "error": str(exc), "fail_class": "fetch_fail"})
        finally:
            SEM.release()


def main() -> None:
    httpd = ThreadingHTTPServer(LISTEN, Handler)
    print(f"trafilatura-http listening on {LISTEN[0]}:{LISTEN[1]} semaphore=8", flush=True)
    httpd.serve_forever()


if __name__ == "__main__":
    main()