Initial MaterialSub release
Derived from pengchujin/jzsub at 222a90265d2a8797ca258eb1a980cee0863a8311; preserve the upstream MIT license and attribution.
This commit is contained in:
+196
@@ -0,0 +1,196 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user