#!/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())