skill-materialsub/skills/materialsub/scripts/citation_watermark.py

#!/usr/bin/env python3
"""Attach a user-approved citation watermark specification to a MaterialSub job."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import sys
import tempfile
from typing import Any, Sequence
from urllib.parse import parse_qsl, urlsplit


MAX_CITATION_CHARACTERS = 1200
DEFAULT_INTERNAL_NOTICE = "内容引自网络，仅供内部交流"
SENSITIVE_QUERY_KEYS = re.compile(
    r"(?:^|_)(?:auth|credential|expires?|key|policy|signature|signed|token)(?:_|$)",
    re.IGNORECASE,
)


class CitationError(RuntimeError):
    """A malformed citation request or job manifest."""


def _read_json(path: Path) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise CitationError(f"could not read JSON manifest {path}: {exc}") from exc
    if not isinstance(value, dict):
        raise CitationError("download manifest root must be an object")
    return value


def _normalize_citation(value: str) -> str:
    lines = [" ".join(line.split()) for line in value.splitlines()]
    value = "\n".join(line for line in lines if line)
    if not value:
        raise CitationError("citation cannot be empty")
    if len(value) > MAX_CITATION_CHARACTERS:
        raise CitationError(
            f"citation is too long ({len(value)} characters; maximum {MAX_CITATION_CHARACTERS})"
        )
    if any(ord(character) < 32 and character != "\n" for character in value):
        raise CitationError("citation contains unsupported control characters")
    return value


def _normalize_notice(value: str) -> str:
    value = " ".join(value.split())
    if not value:
        raise CitationError("notice cannot be empty")
    if len(value) > 120 or any(ord(character) < 32 for character in value):
        raise CitationError("notice contains unsupported or excessive text")
    return value


def _public_source_url(manifest: dict[str, Any]) -> str:
    source = manifest.get("source")
    value = source.get("url") if isinstance(source, dict) else None
    if not isinstance(value, str) or not value.strip():
        raise CitationError("download manifest has no canonical source URL")
    parts = urlsplit(value.strip())
    if parts.scheme not in ("http", "https") or not parts.netloc or parts.username or parts.password:
        raise CitationError("manifest source URL is not a public HTTP(S) URL")
    for key, _ in parse_qsl(parts.query, keep_blank_values=True):
        if SENSITIVE_QUERY_KEYS.search(key):
            raise CitationError("refusing to put a signed or credential-like URL in a watermark")
    return value.strip()


def _atomic_write_text(path: Path, value: str) -> None:
    fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(value)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    except BaseException:
        try:
            os.close(fd)
        except OSError:
            pass
        Path(temporary).unlink(missing_ok=True)
        raise


def _atomic_write_json(path: Path, value: dict[str, Any]) -> None:
    _atomic_write_text(
        path,
        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
    )


def attach_citation(
    manifest_path: Path,
    citation: str,
    *,
    include_source_url: bool = False,
    notice: str | None = None,
) -> Path:
    manifest_path = manifest_path.expanduser().resolve()
    manifest = _read_json(manifest_path)
    citation = _normalize_citation(citation)
    notice = _normalize_notice(notice) if notice is not None else None
    source_url = _public_source_url(manifest) if include_source_url else None
    citation_body = citation + (f"\nSource: {source_url}" if source_url else "")
    rendered_text = (f"{notice}\n\n" if notice else "") + citation_body + "\n"
    citation_path = manifest_path.parent / "citation-watermark.txt"
    _atomic_write_text(citation_path, rendered_text)
    digest = hashlib.sha256(citation_path.read_bytes()).hexdigest()
    manifest["citation_watermark"] = {
        "enabled": True,
        "citation_file": citation_path.name,
        "citation_sha256": digest,
        "include_source_url": include_source_url,
        "layout": (
            "notice-plus-three-line-citation"
            if notice and len(citation.splitlines()) == 3 and not include_source_url
            else "notice-plus-citation" if notice else "citation-only"
        ),
        "notice": notice,
        "position": "top-left",
    }
    _atomic_write_json(manifest_path, manifest)
    return citation_path


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Bind a user-approved top-left citation watermark to a MaterialSub job."
    )
    parser.add_argument("manifest", type=Path, help="download-manifest.json")
    citation = parser.add_mutually_exclusive_group()
    citation.add_argument("--citation", help="complete formal citation")
    citation.add_argument("--citation-file", type=Path, help="UTF-8 file containing the citation")
    parser.add_argument("--authors", help="citation authors line")
    parser.add_argument("--title", help="citation title line")
    parser.add_argument("--publication", help="journal, identifiers, and year line")
    parser.add_argument(
        "--notice",
        nargs="?",
        const=DEFAULT_INTERNAL_NOTICE,
        help=(
            "add an approved notice above the citation; without TEXT uses "
            f"{DEFAULT_INTERNAL_NOTICE!r}"
        ),
    )
    parser.add_argument(
        "--include-source-url",
        action="store_true",
        help="also show the manifest's canonical public source URL",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    try:
        structured = (args.authors, args.title, args.publication)
        has_structured = any(value is not None for value in structured)
        if has_structured:
            if not all(value is not None for value in structured):
                raise CitationError("--authors, --title, and --publication must be used together")
            if args.citation is not None or args.citation_file is not None:
                raise CitationError("structured citation fields cannot be combined with --citation")
            citation = "\n".join(str(value) for value in structured)
        elif args.citation_file is not None:
            citation = args.citation_file.read_text(encoding="utf-8")
        elif args.citation is not None:
            citation = args.citation
        else:
            raise CitationError(
                "provide --citation, --citation-file, or all three structured citation fields"
            )
        path = attach_citation(
            args.manifest,
            citation,
            include_source_url=args.include_source_url,
            notice=args.notice,
        )
    except (CitationError, OSError, UnicodeDecodeError) as exc:
        print(f"citation watermark error: {exc}", file=sys.stderr)
        return 2
    print(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())