mirror of
https://github.com/guillaumemeyer/watermarks-remover.git
synced 2026-08-22 13:11:57 +02:00
254 lines
8.2 KiB
Python
254 lines
8.2 KiB
Python
|
"""Tests for the binary-input guard on the text-only tools."""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import io
|
|||
|
|
import os
|
|||
|
|
import subprocess
|
|||
|
|
import sys
|
|||
|
|
import zipfile
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
|
|||
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|||
|
|
SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts"
|
|||
|
|
sys.path.insert(0, str(SCRIPTS))
|
|||
|
|
|
|||
|
|
from common import guard_binary, looks_binary # noqa: E402
|
|||
|
|
|
|||
|
|
DOCX_XML = (
|
|||
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
|||
|
|
'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
|||
|
|
"<w:body><w:p><w:r><w:t>Table 1 holds the results.</w:t></w:r></w:p></w:body>"
|
|||
|
|
"</w:document>"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def make_docx(path: Path) -> Path:
|
|||
|
|
with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as zf:
|
|||
|
|
zf.writestr("[Content_Types].xml", "<Types/>")
|
|||
|
|
zf.writestr("word/document.xml", DOCX_XML)
|
|||
|
|
return path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run(script: str, *args: str) -> subprocess.CompletedProcess:
|
|||
|
|
return subprocess.run(
|
|||
|
|
[sys.executable, str(SCRIPTS / script), *args],
|
|||
|
|
capture_output=True,
|
|||
|
|
text=True,
|
|||
|
|
timeout=60,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# --- looks_binary ----------------------------------------------------------
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize(
|
|||
|
|
"data,expected_fragment",
|
|||
|
|
[
|
|||
|
|
(b"PK\x03\x04rest", "ZIP"),
|
|||
|
|
(b"%PDF-1.7\n", "PDF"),
|
|||
|
|
(b"\x89PNG\r\n\x1a\nrest", "PNG"),
|
|||
|
|
(b"\xff\xd8\xff\xe0rest", "JPEG"),
|
|||
|
|
(b"\x7fELF\x02\x01", "ELF"),
|
|||
|
|
(b"SQLite format 3\x00", "SQLite"),
|
|||
|
|
(b"plain text\x00with a nul", "NUL"),
|
|||
|
|
],
|
|||
|
|
)
|
|||
|
|
def test_flags_binary(data, expected_fragment):
|
|||
|
|
kind = looks_binary(data)
|
|||
|
|
assert kind is not None
|
|||
|
|
assert expected_fragment.lower() in kind.lower()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize(
|
|||
|
|
"data",
|
|||
|
|
[
|
|||
|
|
b"",
|
|||
|
|
b"Just some prose.\n",
|
|||
|
|
b"# Markdown\n\n- bullet\n",
|
|||
|
|
"Accented prose: naïve café résumé\n".encode("utf-8"),
|
|||
|
|
"Zero width and nbsp here\n".encode("utf-8"),
|
|||
|
|
b"Latin-1 bytes: caf\xe9 na\xefve\n", # not UTF-8, still text
|
|||
|
|
b"a\tb\r\nc\x0cd\x1b[0m\n", # tabs, CRLF, form feed, ANSI escape
|
|||
|
|
],
|
|||
|
|
)
|
|||
|
|
def test_allows_text(data):
|
|||
|
|
assert looks_binary(data) is None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_compressed_bytes_are_flagged(tmp_path):
|
|||
|
|
data = make_docx(tmp_path / "x.docx").read_bytes()
|
|||
|
|
assert looks_binary(data) is not None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_guard_binary_can_be_overridden():
|
|||
|
|
guard_binary(b"PK\x03\x04", "x.docx", allow_binary=True) # must not raise
|
|||
|
|
with pytest.raises(SystemExit) as exc:
|
|||
|
|
guard_binary(b"PK\x03\x04", "x.docx")
|
|||
|
|
assert exc.value.code == 2
|
|||
|
|
|
|||
|
|
|
|||
|
|
# --- CLI behaviour ---------------------------------------------------------
|
|||
|
|
|
|||
|
|
def test_inspect_text_refuses_docx(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
r = run("inspect_text.py", str(docx))
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert "looks like" in r.stderr
|
|||
|
|
assert "inspect_file.py" in r.stderr
|
|||
|
|
assert "Suspicious:" not in r.stdout
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_inspect_text_force_text_still_works(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
r = run("inspect_text.py", str(docx), "--force-text")
|
|||
|
|
assert r.returncode in (0, 1)
|
|||
|
|
assert "Length:" in r.stdout
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_text_refuses_docx_and_writes_nothing(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
before = docx.read_bytes()
|
|||
|
|
out = tmp_path / "doc.cleaned.docx"
|
|||
|
|
r = run("clean_text.py", str(docx), "-o", str(out))
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert not out.exists()
|
|||
|
|
assert docx.read_bytes() == before
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_text_in_place_leaves_docx_intact(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
before = docx.read_bytes()
|
|||
|
|
r = run("clean_text.py", str(docx), "--in-place")
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert docx.read_bytes() == before
|
|||
|
|
assert not (tmp_path / "doc.docx.bak").exists()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_file_still_routes_docx_to_container(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
out = tmp_path / "out.docx"
|
|||
|
|
r = run("clean_file.py", str(docx), "-o", str(out), "--json")
|
|||
|
|
assert r.returncode == 0, r.stderr
|
|||
|
|
assert out.exists()
|
|||
|
|
with zipfile.ZipFile(out) as zf:
|
|||
|
|
assert zf.testzip() is None
|
|||
|
|
assert "word/document.xml" in zf.namelist()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_file_refuses_unknown_binary(tmp_path):
|
|||
|
|
blob = tmp_path / "mystery.bin"
|
|||
|
|
blob.write_bytes(b"\x00\x01\x02\x03" * 64)
|
|||
|
|
out = tmp_path / "out.bin"
|
|||
|
|
r = run("clean_file.py", str(blob), "-o", str(out))
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert not out.exists()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_file_in_place_refuses_before_taking_a_backup(tmp_path):
|
|||
|
|
"""The refusal must land before backup_path(), not after.
|
|||
|
|
|
|||
|
|
Sniffing after the backup left a .bak sidecar for a file the run never
|
|||
|
|
touches, which is exactly what clean_text.py avoids.
|
|||
|
|
"""
|
|||
|
|
blob = tmp_path / "mystery.bin"
|
|||
|
|
blob.write_bytes(b"\x00\x01\x02\x03" * 64)
|
|||
|
|
before = blob.read_bytes()
|
|||
|
|
r = run("clean_file.py", str(blob), "--in-place")
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert blob.read_bytes() == before
|
|||
|
|
assert not (tmp_path / "mystery.bin.bak").exists()
|
|||
|
|
assert list(tmp_path.iterdir()) == [blob]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_clean_file_in_place_as_text_on_docx_leaves_no_backup(tmp_path):
|
|||
|
|
"""--as text bypasses classify(), so the guard is the only thing left."""
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
before = docx.read_bytes()
|
|||
|
|
r = run("clean_file.py", str(docx), "--in-place", "--as", "text")
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert docx.read_bytes() == before
|
|||
|
|
assert not (tmp_path / "doc.docx.bak").exists()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_router_advice_is_not_circular(tmp_path):
|
|||
|
|
"""clean_file.py / inspect_file.py must not point back at themselves."""
|
|||
|
|
blob = tmp_path / "mystery.bin"
|
|||
|
|
blob.write_bytes(b"\x00\x01\x02\x03" * 64)
|
|||
|
|
for script in ("clean_file.py", "inspect_file.py"):
|
|||
|
|
r = run(script, str(blob))
|
|||
|
|
assert r.returncode == 2, script
|
|||
|
|
assert "no supported text, image or container format" in r.stderr, script
|
|||
|
|
assert "Use inspect_file.py / clean_file.py" not in r.stderr, script
|
|||
|
|
assert "--force-text" in r.stderr, script
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_text_only_scripts_keep_the_pointer_to_the_routers(tmp_path):
|
|||
|
|
docx = make_docx(tmp_path / "doc.docx")
|
|||
|
|
r = run("clean_text.py", str(docx))
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert "Use inspect_file.py / clean_file.py" in r.stderr
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_text_files_are_unaffected(tmp_path):
|
|||
|
|
src = tmp_path / "note.txt"
|
|||
|
|
src.write_text("Hiddenmark here.\n", encoding="utf-8")
|
|||
|
|
out = tmp_path / "note.cleaned.txt"
|
|||
|
|
r = run("clean_text.py", str(src), "-o", str(out))
|
|||
|
|
assert r.returncode == 0, r.stderr
|
|||
|
|
assert out.read_text(encoding="utf-8") == "Hiddenmark here.\n"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_stdin_binary_is_refused():
|
|||
|
|
docx = io.BytesIO()
|
|||
|
|
with zipfile.ZipFile(docx, "w") as zf:
|
|||
|
|
zf.writestr("word/document.xml", DOCX_XML)
|
|||
|
|
r = subprocess.run(
|
|||
|
|
[sys.executable, str(SCRIPTS / "inspect_text.py")],
|
|||
|
|
input=docx.getvalue(),
|
|||
|
|
capture_output=True,
|
|||
|
|
timeout=60,
|
|||
|
|
)
|
|||
|
|
assert r.returncode == 2
|
|||
|
|
assert b"looks like" in r.stderr
|
|||
|
|
|
|||
|
|
|
|||
|
|
PNG_HEADER = b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR" + b"\x00" * 32
|
|||
|
|
JPEG_HEADER = b"\xff\xd8\xff\xe0\x00\x10JFIF\x00" + b"A" * 64
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.mark.parametrize("data,label", [(PNG_HEADER, "PNG"), (JPEG_HEADER, "JPEG")])
|
|||
|
|
@pytest.mark.parametrize("io_encoding", [None, "cp1252", "latin-1"])
|
|||
|
|
def test_stdin_non_ascii_magic_is_refused_whatever_the_codec(data, label, io_encoding):
|
|||
|
|
"""The ZIP test alone could not catch this: 'PK' is ASCII.
|
|||
|
|
|
|||
|
|
PNG's 0x89 and JPEG's 0xff only survive to the sniff if stdin is read as
|
|||
|
|
bytes. Decoding first makes detection depend on the console codec.
|
|||
|
|
"""
|
|||
|
|
env = dict(os.environ)
|
|||
|
|
if io_encoding is None:
|
|||
|
|
env.pop("PYTHONIOENCODING", None)
|
|||
|
|
else:
|
|||
|
|
env["PYTHONIOENCODING"] = io_encoding
|
|||
|
|
r = subprocess.run(
|
|||
|
|
[sys.executable, str(SCRIPTS / "inspect_text.py")],
|
|||
|
|
input=data,
|
|||
|
|
capture_output=True,
|
|||
|
|
timeout=60,
|
|||
|
|
env=env,
|
|||
|
|
)
|
|||
|
|
assert r.returncode == 2, (label, io_encoding, r.stderr)
|
|||
|
|
assert label.encode() in r.stderr, (label, io_encoding, r.stderr)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_stdin_text_still_flows_through():
|
|||
|
|
r = subprocess.run(
|
|||
|
|
[sys.executable, str(SCRIPTS / "clean_text.py")],
|
|||
|
|
input="plaintext\n".encode("utf-8"),
|
|||
|
|
capture_output=True,
|
|||
|
|
timeout=60,
|
|||
|
|
)
|
|||
|
|
assert r.returncode == 0, r.stderr
|
|||
|
|
assert r.stdout.replace(b"\r\n", b"\n") == b"plaintext\n"
|