186 lines
7.1 KiB
Python
186 lines
7.1 KiB
Python
"""
|
||
enhanced_sanitize_pdf.py
|
||
========================
|
||
|
||
This utility removes SANS‑style watermark overlays (names, e‑mails, hashes,
|
||
dates and licence lines) from a PDF without leaving selection boxes or
|
||
visible artefacts. It can open password‑protected files by prompting the
|
||
user for a password and will automatically write the cleaned document
|
||
alongside the original with ``_clean`` appended to its file name.
|
||
|
||
In addition to the basic cleaning implemented in ``sanitize_pdf_watermark.py``,
|
||
this version includes two important improvements:
|
||
|
||
* **Password handling:** If the input PDF is encrypted, the script will ask
|
||
for a password using ``input()`` and attempt to authenticate before
|
||
continuing. Should the password prove wrong, the script aborts with a
|
||
helpful message. When processing unencrypted files, no prompt will be
|
||
shown.
|
||
|
||
* **Robust saving:** Certain PDFs have corrupted object streams or other
|
||
structural issues that prevent PyMuPDF from saving them directly. If a
|
||
``ValueError`` or ``RuntimeError`` is raised during ``doc.save()``, the
|
||
script falls back to creating a repaired copy in memory using
|
||
``doc.tobytes(garbage=4, deflate=True)``. It then opens this repaired
|
||
content as a new document and saves it to disk. This technique is
|
||
recommended by the PyMuPDF documentation for recovering from broken
|
||
object streams【206307526196463†L135-L154】.
|
||
|
||
Usage::
|
||
|
||
python enhanced_sanitize_pdf.py path/to/input.pdf
|
||
|
||
The cleaned PDF will be created in the same directory as ``input.pdf`` with
|
||
``_clean`` appended before the ``.pdf`` extension. The script prints the
|
||
number of watermark segments removed.
|
||
|
||
This script requires the PyMuPDF (``fitz``) package.
|
||
"""
|
||
|
||
import argparse
|
||
import sys
|
||
import os
|
||
import re
|
||
import getpass
|
||
|
||
import fitz # type: ignore[import]
|
||
|
||
|
||
def open_with_password(path: str, max_attempts: int = 3) -> fitz.Document:
|
||
"""Open ``path`` and prompt for a password when the PDF is encrypted.
|
||
|
||
PyMuPDF happily opens encrypted PDFs without immediately failing, but it
|
||
marks them as needing authentication via ``doc.needs_pass``. The original
|
||
implementation relied on ``fitz.open`` raising a ``RuntimeError`` to detect
|
||
encryption, which meant passwords were never requested and the caller
|
||
would hit authorization errors later. We now inspect ``doc.needs_pass``
|
||
directly and allow up to ``max_attempts`` attempts at entering the password
|
||
before aborting.
|
||
"""
|
||
|
||
try:
|
||
doc = fitz.open(path)
|
||
except RuntimeError as exc:
|
||
raise RuntimeError(f"Unable to open '{path}': {exc}") from exc
|
||
|
||
if not doc.needs_pass:
|
||
return doc
|
||
|
||
print(f"The document '{path}' is encrypted.")
|
||
for attempt in range(1, max_attempts + 1):
|
||
password = getpass.getpass("Please enter the password: ")
|
||
if doc.authenticate(password):
|
||
return doc
|
||
print("Incorrect password. Try again." if attempt < max_attempts else "Incorrect password.")
|
||
|
||
doc.close()
|
||
raise ValueError("Failed to authenticate after multiple attempts. Unable to open the PDF.")
|
||
|
||
|
||
def remove_watermark_segments(doc: fitz.Document) -> int:
|
||
"""Remove known watermark text blocks from every page of the document.
|
||
|
||
Watermark segments are identified by unusual font sizes and text
|
||
rendering modes used exclusively by the SANS licence overlays:
|
||
|
||
* Invisible 36‑point text with rendering mode 1 or 3 (hidden names,
|
||
e‑mails, numeric codes and hashes).
|
||
* Rotated overlays drawn with a –25 point font.
|
||
* Footer lines drawn with 10‑ or 18‑point fonts containing the
|
||
licence owner’s information.
|
||
|
||
The patterns below use regular expressions to locate complete
|
||
``BT … ET`` text objects matching these criteria and delete them.
|
||
|
||
Returns
|
||
-------
|
||
int
|
||
The total number of removed segments across all pages.
|
||
"""
|
||
# Compile regular expressions once for efficiency
|
||
patterns = [
|
||
# Invisible 36‑point text with rendering mode 1 or 3
|
||
re.compile(r"BT[^B]*?/F\d+\s+36\s+Tf[^E]*?\s+(?:1|3)\s+Tr[^E]*?ET", re.S),
|
||
# Diagonal overlay drawn with a –25 point font
|
||
re.compile(r"BT[^B]*?/F\d+\s+-25\s+Tf[^E]*?ET", re.S),
|
||
# Horizontal footers drawn with 10‑point fonts
|
||
re.compile(r"BT[^B]*?/F\d+\s+10\s+Tf[^E]*?ET", re.S),
|
||
# Horizontal footers drawn with 18‑point fonts
|
||
re.compile(r"BT[^B]*?/F\d+\s+18\s+Tf[^E]*?ET", re.S),
|
||
]
|
||
removed_count = 0
|
||
for page in doc:
|
||
# Iterate over all content streams referenced by this page
|
||
for xref in page.get_contents():
|
||
stream = doc.xref_stream(xref).decode("latin1")
|
||
modified = stream
|
||
for pat in patterns:
|
||
modified, count = pat.subn("", modified)
|
||
removed_count += count
|
||
if modified != stream:
|
||
# Update the stream in place
|
||
doc.update_stream(xref, modified.encode("latin1"))
|
||
return removed_count
|
||
|
||
|
||
def build_output_path(input_path: str) -> str:
|
||
"""Derive the output filename by inserting ``_clean`` before the extension."""
|
||
directory, filename = os.path.split(input_path)
|
||
base, ext = os.path.splitext(filename)
|
||
return os.path.join(directory, f"{base}_clean{ext}")
|
||
|
||
|
||
def save_document(doc: fitz.Document, out_path: str) -> None:
|
||
"""Save the document to ``out_path`` with a robust fallback.
|
||
|
||
PyMuPDF may raise errors when saving PDFs with damaged object streams or
|
||
encryption. In such cases, attempt to repair the document by
|
||
exporting its bytes and reopening it as a new document, then save
|
||
again. If the fallback also fails, the exception is propagated.
|
||
"""
|
||
try:
|
||
doc.save(out_path, garbage=4, deflate=True)
|
||
except (ValueError, RuntimeError) as exc:
|
||
# Attempt to repair the document by exporting and reopening
|
||
print(f"Encountered an error while saving: {exc}. Attempting repair …")
|
||
# Export the document to bytes, performing garbage collection and
|
||
# decompression. This will also decrypt the file if a password
|
||
# was used【206307526196463†L135-L154】.
|
||
buffer = doc.tobytes(garbage=4, deflate=True)
|
||
repaired = fitz.open("pdf", buffer)
|
||
repaired.save(out_path, garbage=4, deflate=True)
|
||
repaired.close()
|
||
|
||
|
||
def sanitize_pdf(path: str) -> None:
|
||
"""Perform watermark removal on ``path`` and write a cleaned copy."""
|
||
doc = open_with_password(path)
|
||
removed = remove_watermark_segments(doc)
|
||
out_path = build_output_path(path)
|
||
save_document(doc, out_path)
|
||
doc.close()
|
||
print(f"Removed {removed} watermark segment(s). Cleaned PDF saved as '{out_path}'.")
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
description="Remove SANS-style watermark overlays from the provided PDF. "
|
||
"Encrypted PDFs are supported and will trigger a password prompt."
|
||
)
|
||
parser.add_argument(
|
||
"input_pdf",
|
||
metavar="INPUT.pdf",
|
||
help="Path to the PDF you want sanitized."
|
||
)
|
||
return parser
|
||
|
||
|
||
def main(argv: list[str]) -> None:
|
||
parser = build_parser()
|
||
args = parser.parse_args(argv[1:])
|
||
sanitize_pdf(args.input_pdf)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main(sys.argv)
|