staging

zocrates.discovery.staging

Temporary download staging, SSRF validation, byte capping, and cleanup for candidate PDF papers.

DEFAULT_DOWNLOAD_TIMEOUT module-attribute

DEFAULT_DOWNLOAD_TIMEOUT = httpx.Timeout(connect=15.0, read=30.0, write=15.0, pool=30.0)

MAX_PDF_BYTES module-attribute

MAX_PDF_BYTES = 100 * 1024 * 1024

PDF_MAGIC_BYTES module-attribute

PDF_MAGIC_BYTES = b'%PDF-'

logger module-attribute

logger = logging.getLogger(__name__)

CandidatePaper dataclass

CandidatePaper(candidate_id, title, authors=list(), year=None, venue=None, doi=None, arxiv_id=None, openalex_id='', abstract=None, tldr=None, is_oa=False, oa_pdf_url=None, oa_pdf_urls=list(), citation_count=0, primary_topic=None, keywords=list(), in_library=False)

Represents a discovered scholarly paper candidate.

from_dict classmethod

from_dict(data)

Deserialize CandidatePaper from a dictionary.

Source code in zocrates/discovery/models.py

@classmethod
def from_dict(cls, data: dict[str, Any]) -> "CandidatePaper":
    """Deserialize CandidatePaper from a dictionary."""
    citation_count_val = data.get("citation_count", 0)
    try:
        citation_count = int(citation_count_val) if citation_count_val is not None else 0
    except (ValueError, TypeError):
        citation_count = 0

    year_val = data.get("year")
    year: int | None = None
    if year_val is not None:
        try:
            year = int(year_val)
        except (ValueError, TypeError):
            year = None

    raw_authors = data.get("authors")
    if isinstance(raw_authors, list):
        authors = [str(a).strip() for a in raw_authors if a and str(a).strip()]
    elif isinstance(raw_authors, str) and raw_authors.strip():
        authors = [raw_authors.strip()]
    else:
        authors = []

    raw_keywords = data.get("keywords")
    if isinstance(raw_keywords, list):
        keywords = [str(k).strip() for k in raw_keywords if k and str(k).strip()]
    elif isinstance(raw_keywords, str) and raw_keywords.strip():
        keywords = [raw_keywords.strip()]
    else:
        keywords = []

    def _clean_opt_str(val: Any) -> str | None:
        """Clean and trim an optional string or scalar value.

        Args:
            val: Input value to convert to a stripped string.

        Returns:
            Stripped string, or None if empty or invalid.
        """
        if val is not None and isinstance(val, (str, int, float)):
            s = str(val).strip()
            return s if s else None
        return None

    raw_urls = data.get("oa_pdf_urls")
    oa_pdf_urls: list[str] = []
    if isinstance(raw_urls, list):
        oa_pdf_urls = [str(u).strip() for u in raw_urls if u and str(u).strip()]
    primary_pdf = _clean_opt_str(data.get("oa_pdf_url"))
    if primary_pdf and primary_pdf not in oa_pdf_urls:
        oa_pdf_urls.insert(0, primary_pdf)

    return cls(
        candidate_id=str(data.get("candidate_id") or ""),
        title=str(data.get("title") or "").strip(),
        authors=authors,
        year=year,
        venue=_clean_opt_str(data.get("venue")),
        doi=_clean_opt_str(data.get("doi")),
        arxiv_id=_clean_opt_str(data.get("arxiv_id")),
        openalex_id=str(data.get("openalex_id") or ""),
        abstract=_clean_opt_str(data.get("abstract")),
        tldr=_clean_opt_str(data.get("tldr")),
        is_oa=bool(data.get("is_oa", False)),
        oa_pdf_url=primary_pdf,
        oa_pdf_urls=oa_pdf_urls,
        citation_count=citation_count,
        primary_topic=_clean_opt_str(data.get("primary_topic")),
        keywords=keywords,
        in_library=bool(data.get("in_library", False)),
    )

to_dict

to_dict()

Serialize CandidatePaper to a plain dictionary.

Source code in zocrates/discovery/models.py

def to_dict(self) -> dict[str, Any]:
    """Serialize CandidatePaper to a plain dictionary."""
    return {
        "candidate_id": self.candidate_id,
        "title": self.title,
        "authors": list(self.authors),
        "year": self.year,
        "venue": self.venue,
        "doi": self.doi,
        "arxiv_id": self.arxiv_id,
        "openalex_id": self.openalex_id,
        "abstract": self.abstract,
        "tldr": self.tldr,
        "is_oa": self.is_oa,
        "oa_pdf_url": self.oa_pdf_url,
        "oa_pdf_urls": list(self.oa_pdf_urls),
        "citation_count": self.citation_count,
        "primary_topic": self.primary_topic,
        "keywords": list(self.keywords),
        "in_library": self.in_library,
    }

_stream_to_file async

_stream_to_file(client, initial_url, dest_path, allow_html_landing=True)

Stream remote URL to local destination path with hop-by-hop HTTPS validation, byte capping, and magic byte checks.

Source code in zocrates/discovery/staging.py

async def _stream_to_file(
    client: httpx.AsyncClient,
    initial_url: str,
    dest_path: Path,
    allow_html_landing: bool = True,
) -> str | None:
    """Stream remote URL to local destination path with hop-by-hop HTTPS validation, byte capping, and magic byte checks."""
    curr_url = initial_url
    max_redirects = 5

    try:
        for _ in range(max_redirects + 1):
            parsed = urllib.parse.urlparse(curr_url)
            if parsed.scheme.lower() != "https":
                if parsed.scheme.lower() == "http":
                    curr_url = "https://" + curr_url[7:]
                    parsed = urllib.parse.urlparse(curr_url)
                else:
                    logger.warning("SSRF Protection: Blocked non-https URL or redirect: %s", curr_url)
                    return None

            hostname = (parsed.hostname or "").lower()
            if hostname in ("localhost", "127.0.0.1", "::1", "169.254.169.254") or hostname.startswith("192.168.") or hostname.startswith("10."):
                logger.warning("SSRF Protection: Blocked request to private/local host %s", hostname)
                return None

            headers = {
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
                ),
                "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,application/pdf,*/*;q=0.8",
            }
            async with client.stream("GET", curr_url, headers=headers, follow_redirects=False) as response:
                if response.status_code in (301, 302, 303, 307, 308):
                    loc = response.headers.get("Location")
                    if not loc:
                        logger.warning("Redirect status %d with no Location header for %s", response.status_code, curr_url)
                        return None
                    next_url = urllib.parse.urljoin(curr_url, loc)
                    curr_url = next_url
                    continue

                if response.status_code != 200:
                    logger.debug("PDF download failed with status %d for %s", response.status_code, curr_url)
                    return None

                link_header = response.headers.get("Link")
                candidate_pdf_from_header = (
                    extract_pdf_from_link_header(link_header, curr_url)
                    if allow_html_landing
                    else None
                )

                total_bytes = 0
                header_checked = False
                header_buffer = bytearray()
                is_html_candidate = False
                html_buf = bytearray()

                with open(dest_path, "wb") as f:
                    async for chunk in response.aiter_bytes():
                        if not chunk:
                            continue
                        total_bytes += len(chunk)
                        if total_bytes > MAX_PDF_BYTES:
                            logger.warning("PDF download exceeded %d bytes limit for %s", MAX_PDF_BYTES, curr_url)
                            raise ValueError(f"Byte cap exceeded ({total_bytes} > {MAX_PDF_BYTES})")

                        if not header_checked:
                            header_buffer.extend(chunk)
                            if len(header_buffer) >= 5:
                                if header_buffer.startswith(PDF_MAGIC_BYTES):
                                    header_checked = True
                                    f.write(header_buffer)
                                elif allow_html_landing:
                                    content_type = response.headers.get("content-type", "").lower()
                                    if (
                                        candidate_pdf_from_header
                                        or ("html" in content_type)
                                        or header_buffer.lstrip().startswith(
                                            (b"<!doc", b"<!DOC", b"<html", b"<HTML", b"<?xml", b"<head", b"<body")
                                        )
                                    ):
                                        is_html_candidate = True
                                        html_buf.extend(header_buffer)
                                        header_checked = True
                                    else:
                                        logger.warning(
                                            "File magic bytes mismatch for %s: expected %r, got %r",
                                            curr_url,
                                            PDF_MAGIC_BYTES,
                                            bytes(header_buffer[:5]),
                                        )
                                        raise ValueError("Invalid PDF magic bytes")
                                else:
                                    logger.warning(
                                        "File magic bytes mismatch for %s: expected %r, got %r",
                                        curr_url,
                                        PDF_MAGIC_BYTES,
                                        bytes(header_buffer[:5]),
                                    )
                                    raise ValueError("Invalid PDF magic bytes")
                        elif is_html_candidate:
                            html_buf.extend(chunk)
                            if len(html_buf) >= 524288:
                                break
                        else:
                            f.write(chunk)

                if is_html_candidate:
                    if dest_path.exists():
                        try:
                            dest_path.unlink()
                        except OSError:
                            pass
                    pdf_target_url = candidate_pdf_from_header
                    html_text = ""
                    if not pdf_target_url:
                        html_text = html_buf.decode("utf-8", errors="replace")
                        pdf_target_url = extract_pdf_url_from_html(html_text, curr_url)
                    if not pdf_target_url:
                        pdf_target_url = await resolve_dspace7_pdf(
                            client, curr_url, html_text, initial_url=initial_url
                        )
                    if pdf_target_url:
                        logger.info(
                            "Extracted direct PDF URL %s from HTML landing page %s",
                            pdf_target_url,
                            curr_url,
                        )
                        return await _stream_to_file(
                            client=client,
                            initial_url=pdf_target_url,
                            dest_path=dest_path,
                            allow_html_landing=False,
                        )
                    logger.warning("No PDF URL found in HTML landing page %s", curr_url)
                    raise ValueError("Invalid PDF magic bytes")

                if not header_checked:
                    logger.warning("File too short (< 5 bytes) for %s", curr_url)
                    raise ValueError("File too short")

                return str(dest_path.resolve())

        logger.warning("Too many redirects (> %d) for %s", max_redirects, initial_url)
        return None
    except Exception as exc:
        logger.debug("Download and stage failed for %s: %s", curr_url, exc)
        if dest_path.exists():
            try:
                dest_path.unlink()
            except OSError:
                pass
        return None

cleanup_stale_staging_directories

cleanup_stale_staging_directories(max_age_seconds=7200)

Remove staging subdirectories older than max_age_seconds (default: 2 hours).

Parameters:

Name	Type	Description	Default
max_age_seconds	int	Age threshold in seconds.	7200

Returns:

Type	Description
int	Number of directories removed.

Source code in zocrates/discovery/staging.py

def cleanup_stale_staging_directories(max_age_seconds: int = 7200) -> int:
    """
    Remove staging subdirectories older than `max_age_seconds` (default: 2 hours).

    Args:
        max_age_seconds: Age threshold in seconds.

    Returns:
        Number of directories removed.
    """
    root = get_staging_root()
    if not root.exists():
        return 0

    removed = 0
    now = time.time()
    for item in root.iterdir():
        if item.is_dir():
            try:
                mtime = item.stat().st_mtime
                if (now - mtime) > max_age_seconds:
                    shutil.rmtree(item, ignore_errors=True)
                    removed += 1
            except Exception as e:
                logger.debug("Failed removing stale staging directory %s: %s", item, e)
    return removed

download_and_stage_pdf async

download_and_stage_pdf(candidate, staging_dir, client=None)

Download and stage an Open Access PDF for a candidate paper.

Enforces

• SSRF Protection: strictly requires https:// URL scheme.
• Safe Filename: sanitizes candidate_id against directory traversal.
• Directory Containment: ensures target path is strictly within staging_dir.
• Magic Byte Verification: first 5 bytes must equal b"%PDF-".
• Streaming Byte Cap: aborts if cumulative bytes exceed MAX_PDF_BYTES (100 MB).

Parameters:

Name	Type	Description	Default
candidate	CandidatePaper	CandidatePaper instance.	required
staging_dir	Path	Target batch directory Path.	required
client	AsyncClient | None	Optional existing httpx.AsyncClient for connection reuse.	None

Returns:

Type	Description
str | None	Absolute file path string on success, or None on failure/non-PDF.

Source code in zocrates/discovery/staging.py

async def download_and_stage_pdf(
    candidate: CandidatePaper,
    staging_dir: Path,
    client: httpx.AsyncClient | None = None,
) -> str | None:
    """
    Download and stage an Open Access PDF for a candidate paper.

    Enforces:
      - SSRF Protection: strictly requires https:// URL scheme.
      - Safe Filename: sanitizes candidate_id against directory traversal.
      - Directory Containment: ensures target path is strictly within staging_dir.
      - Magic Byte Verification: first 5 bytes must equal b"%PDF-".
      - Streaming Byte Cap: aborts if cumulative bytes exceed MAX_PDF_BYTES (100 MB).

    Args:
        candidate: CandidatePaper instance.
        staging_dir: Target batch directory Path.
        client: Optional existing httpx.AsyncClient for connection reuse.

    Returns:
        Absolute file path string on success, or None on failure/non-PDF.
    """
    urls_to_try: list[str] = []
    if candidate.oa_pdf_url and isinstance(candidate.oa_pdf_url, str):
        urls_to_try.append(candidate.oa_pdf_url.strip())
    for alt in getattr(candidate, "oa_pdf_urls", []):
        if alt and isinstance(alt, str):
            clean_alt = alt.strip()
            if clean_alt and clean_alt not in urls_to_try:
                urls_to_try.append(clean_alt)

    if not urls_to_try:
        return None

    raw_id = candidate.candidate_id.strip() if candidate.candidate_id else "candidate"
    safe_id = re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_id)
    if not safe_id:
        safe_id = "candidate"
    id_hash = hashlib.sha1(raw_id.encode("utf-8")).hexdigest()[:8]
    filename = f"{safe_id}_{id_hash}.pdf"

    resolved_staging_dir = Path(staging_dir).resolve()
    dest_path = (resolved_staging_dir / filename).resolve()

    try:
        if not dest_path.is_relative_to(resolved_staging_dir) or dest_path == resolved_staging_dir:
            logger.warning("Path traversal attempt in filename %r", filename)
            return None
    except AttributeError:
        if resolved_staging_dir not in dest_path.parents:
            return None

    resolved_staging_dir.mkdir(parents=True, exist_ok=True)

    async def _try_download(c: httpx.AsyncClient) -> str | None:
        """Attempt downloading candidate PDF across candidate URLs sequentially.

        Args:
            c: Asynchronous HTTP client instance to use for downloads.

        Returns:
            Absolute file path string if downloaded successfully, or None.
        """
        for raw_url in urls_to_try:
            url = raw_url
            if url.lower().startswith("http://"):
                url = "https://" + url[7:]
            elif not url.lower().startswith("https://"):
                logger.warning(
                    "SSRF Protection: Rejected non-https URL %r for candidate %s",
                    url,
                    candidate.candidate_id,
                )
                continue

            res = await _stream_to_file(c, url, dest_path)
            if res and Path(res).exists() and Path(res).stat().st_size > 0:
                return res
        return None

    if client is not None:
        return await _try_download(client)
    else:
        async with httpx.AsyncClient(timeout=DEFAULT_DOWNLOAD_TIMEOUT, follow_redirects=True) as own_client:
            return await _try_download(own_client)

extract_pdf_from_link_header

extract_pdf_from_link_header(link_header, base_url)

Extract a direct PDF URL from HTTP Link header (Signposting RFC 8288).

Source code in zocrates/discovery/staging.py

def extract_pdf_from_link_header(link_header: str, base_url: str) -> str | None:
    """Extract a direct PDF URL from HTTP Link header (Signposting RFC 8288)."""
    if not link_header:
        return None
    for part in link_header.split(","):
        if "application/pdf" in part and ("rel=\"item\"" in part or "rel=\"alternate\"" in part or "rel=item" in part):
            m = re.search(r"""<([^>]+)>""", part)
            if m:
                raw_url = m.group(1).strip()
                if "/bitstreams/" in raw_url and raw_url.endswith("/download"):
                    raw_url = re.sub(
                        r"/bitstreams/([0-9a-fA-F\-]{36})/download",
                        r"/server/api/core/bitstreams/\1/content",
                        raw_url,
                    )
                return urllib.parse.urljoin(base_url, raw_url)
    return None

extract_pdf_url_from_html

extract_pdf_url_from_html(html, base_url)

Extract a direct PDF URL from an HTML landing page using standard scholarly metadata tags (Highwire/Google Scholar, Open Access Signposting, EPrints, and DSpace).

Source code in zocrates/discovery/staging.py

def extract_pdf_url_from_html(html: str, base_url: str) -> str | None:
    """
    Extract a direct PDF URL from an HTML landing page using standard scholarly metadata tags
    (Highwire/Google Scholar, Open Access Signposting, EPrints, and DSpace).
    """
    if not html:
        return None

    # 1. Google Scholar / Highwire meta tag: <meta name="citation_pdf_url" content="...">
    m = re.search(
        r"""<meta\s+[^>]*name=["\']citation_pdf_url["\'][^>]*content=["\']([^"\']+)["\']""",
        html,
        re.IGNORECASE,
    )
    if not m:
        m = re.search(
            r"""<meta\s+[^>]*content=["\']([^"\']+)["\'][^>]*name=["\']citation_pdf_url["\']""",
            html,
            re.IGNORECASE,
        )

    # 2. Open Access Signposting / Link relations: <link rel="item" type="application/pdf" href="...">
    if not m:
        m = re.search(
            r"""<link\s+[^>]*type=["\']application/pdf["\'][^>]*href=["\']([^"\']+)["\']""",
            html,
            re.IGNORECASE,
        )
    if not m:
        m = re.search(
            r"""<link\s+[^>]*href=["\']([^"\']+)["\'][^>]*type=["\']application/pdf["\']""",
            html,
            re.IGNORECASE,
        )

    # 3. EPrints document URL: <meta name="eprints.document_url" content="...">
    if not m:
        m = re.search(
            r"""<meta\s+[^>]*name=["\']eprints\.document_url["\'][^>]*content=["\']([^"\']+)["\']""",
            html,
            re.IGNORECASE,
        )

    # 4. Embedded DSpace 7 REST bitstream API link: /server/api/core/bitstreams/{uuid}/content
    if not m:
        m_dspace = re.search(r"""/server/api/core/bitstreams/([0-9a-fA-F\-]{36})/content""", html)
        if m_dspace:
            return urllib.parse.urljoin(base_url, m_dspace.group(0))

    if m:
        raw_url = m.group(1).strip()
        # DSpace 7 SPA rewrite: /bitstreams/{uuid}/download -> /server/api/core/bitstreams/{uuid}/content
        if "/bitstreams/" in raw_url and raw_url.endswith("/download"):
            raw_url = re.sub(
                r"/bitstreams/([0-9a-fA-F\-]{36})/download",
                r"/server/api/core/bitstreams/\1/content",
                raw_url,
            )
        return urllib.parse.urljoin(base_url, raw_url)

    return None

get_staging_root

get_staging_root()

Return the root directory for temporary discovery batch staging.

Prefers /staging/ (which is within user home, accessible to sandboxed Flatpak Zotero instances having filesystem=home permissions), falling back to ~/.zocrates/staging, and finally system temp.

Source code in zocrates/discovery/staging.py

def get_staging_root() -> Path:
    """Return the root directory for temporary discovery batch staging.

    Prefers <zocrates-data>/staging/ (which is within user home, accessible to sandboxed
    Flatpak Zotero instances having filesystem=home permissions), falling back to
    ~/.zocrates/staging, and finally system temp.
    """
    try:
        from zocrates.config import get_zocrates_data_dir

        root = get_zocrates_data_dir() / "staging"
        root.mkdir(parents=True, exist_ok=True)
        return root
    except Exception:
        pass

    try:
        root = Path.home() / ".zocrates" / "staging"
        root.mkdir(parents=True, exist_ok=True)
        return root
    except Exception:
        pass

    root = Path(tempfile.gettempdir()) / "zocrates_staging"
    root.mkdir(parents=True, exist_ok=True)
    return root

resolve_dspace7_pdf async

resolve_dspace7_pdf(client, page_url, html, initial_url='')

Resolve a direct bitstream PDF URL from a modern DSpace 7 Angular SPA landing page.

DSpace 7 returns a lightweight JavaScript SPA shell () whose primary PDF bitstreams are not present in raw HTML meta tags, but accessible via the native DSpace 7 REST API (/server/api/core/items/{uuid}/bundles).

Source code in zocrates/discovery/staging.py

async def resolve_dspace7_pdf(
    client: httpx.AsyncClient,
    page_url: str,
    html: str,
    initial_url: str = "",
) -> str | None:
    """Resolve a direct bitstream PDF URL from a modern DSpace 7 Angular SPA landing page.

    DSpace 7 returns a lightweight JavaScript SPA shell (<ds-app>) whose primary
    PDF bitstreams are not present in raw HTML meta tags, but accessible via the
    native DSpace 7 REST API (/server/api/core/items/{uuid}/bundles).
    """
    if "dspace" not in html.lower() and "<ds-app" not in html.lower():
        return None

    parsed = urllib.parse.urlparse(page_url)
    origin = f"{parsed.scheme}://{parsed.netloc}"

    # 1. Check if the redirected or current URL contains a DSpace item/publication UUID
    bundles_href: str | None = None
    m_uuid = re.search(r"/(?:entities/[a-zA-Z0-9_\-]+|items|item)/([0-9a-fA-F\-]{36})", page_url)
    if m_uuid:
        item_uuid = m_uuid.group(1)
        bundles_href = f"{origin}/server/api/core/items/{item_uuid}/bundles"
    else:
        # 2. Check if initial or final URL contains a handle path, e.g. handle/10344/10509
        for u in (initial_url, page_url):
            m_handle = re.search(r"handle/([0-9a-zA-Z\.\_\-]+/[0-9a-zA-Z\.\_\-]+)", u)
            if m_handle:
                pid_url = f"{origin}/server/api/pid/find?id={m_handle.group(1)}"
                try:
                    r = await client.get(pid_url, timeout=10)
                    if r.status_code == 200:
                        bundles_href = r.json().get("_links", {}).get("bundles", {}).get("href")
                        if bundles_href:
                            break
                except Exception:
                    pass

    if not bundles_href:
        return None

    try:
        b_res = await client.get(bundles_href, timeout=10)
        if b_res.status_code != 200:
            return None
        b_data = b_res.json()
        bundles = b_data.get("_embedded", {}).get("bundles", [])
        original_bundles = [b for b in bundles if b.get("name") == "ORIGINAL"]
        target_bundles = original_bundles if original_bundles else bundles

        for bundle in target_bundles:
            bs_href = bundle.get("_links", {}).get("bitstreams", {}).get("href")
            if not bs_href:
                continue
            bs_res = await client.get(bs_href, timeout=10)
            if bs_res.status_code != 200:
                continue
            bs_data = bs_res.json()
            for bs in bs_data.get("_embedded", {}).get("bitstreams", []):
                name = bs.get("name", "").lower()
                mime = (bs.get("format") or {}).get("mimetype", "").lower() if isinstance(bs.get("format"), dict) else ""
                if name.endswith(".pdf") or "pdf" in mime:
                    content_href = bs.get("_links", {}).get("content", {}).get("href")
                    if content_href:
                        return urllib.parse.urljoin(origin, content_href)
    except Exception as exc:
        logger.debug("DSpace 7 REST bitstream resolution failed for %s: %s", page_url, exc)

    return None