197 lines
7.3 KiB
Python
197 lines
7.3 KiB
Python
import os
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.main import app
|
|
from app.services import engine
|
|
from app.services.store import document_store
|
|
|
|
client = TestClient(app)
|
|
|
|
CORPUS_PDF = os.path.join(os.path.dirname(__file__), "..", "..", "corpus", "basic", "hello_world.pdf")
|
|
|
|
|
|
@pytest.fixture
|
|
def protected_doc_id():
|
|
# 1. Upload unprotected PDF
|
|
with open(CORPUS_PDF, "rb") as f:
|
|
resp = client.post("/documents", files={"file": ("hello.pdf", f, "application/pdf")})
|
|
assert resp.status_code == 201
|
|
doc_id = resp.json()["id"]
|
|
|
|
# 2. Protect document
|
|
protect_payload = {
|
|
"userPassword": "UserPassword123!",
|
|
"ownerPassword": "OwnerPassword123!",
|
|
"confirmPassword": "UserPassword123!",
|
|
"permissions": {
|
|
"canPrint": True,
|
|
"canCopy": True,
|
|
"canModify": True,
|
|
},
|
|
}
|
|
prot_resp = client.post(f"/documents/{doc_id}/protect", json=protect_payload)
|
|
assert prot_resp.status_code == 200
|
|
assert prot_resp.json()["permissions"]["isEncrypted"] is True
|
|
return doc_id
|
|
|
|
|
|
def test_unlock_authenticated_protected_pdf(protected_doc_id):
|
|
"""TEST 1: Authenticated protected PDF -> unlock succeeds (isEncrypted == False)."""
|
|
unlock_resp = client.post(f"/documents/{protected_doc_id}/unlock")
|
|
assert unlock_resp.status_code == 200, unlock_resp.text
|
|
data = unlock_resp.json()
|
|
assert data["permissions"]["isEncrypted"] is False
|
|
assert data["permissions"]["encryption"] == "None"
|
|
|
|
|
|
def test_wrong_password_auth_fails():
|
|
"""TEST 2: Wrong password authentication fails at upload (401)."""
|
|
# Create protected bytes first
|
|
with open(CORPUS_PDF, "rb") as f:
|
|
resp = client.post("/documents", files={"file": ("hello.pdf", f, "application/pdf")})
|
|
doc_id = resp.json()["id"]
|
|
client.post(f"/documents/{doc_id}/protect", json={"userPassword": "pwd", "confirmPassword": "pwd"})
|
|
export_bytes = client.get(f"/documents/{doc_id}/export").content
|
|
|
|
# Attempt open with wrong password
|
|
upload_resp = client.post(
|
|
"/documents?password=WrongPassword",
|
|
files={"file": ("protected.pdf", export_bytes, "application/pdf")},
|
|
)
|
|
assert upload_resp.status_code == 401
|
|
|
|
|
|
def test_unlock_nonexistent_document():
|
|
"""TEST 3: Direct unlock API request for nonexistent document -> returns 404."""
|
|
resp = client.post("/documents/nonexistent-id-12345/unlock")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
def test_unlocked_pdf_opens_without_password(protected_doc_id):
|
|
"""TEST 4: Unlocked PDF -> open without password via engine."""
|
|
unlock_resp = client.post(f"/documents/{protected_doc_id}/unlock")
|
|
assert unlock_resp.status_code == 200
|
|
|
|
# Retrieve document from store and verify opening bytes without password
|
|
d = document_store.get_document(protected_doc_id)
|
|
pdfengine = engine.require()
|
|
unlocked_doc = pdfengine.PdfDocument.load_from_memory(d["bytes_data"], "")
|
|
assert unlocked_doc.page_count > 0
|
|
|
|
|
|
def test_unlocked_pdf_reports_not_encrypted(protected_doc_id):
|
|
"""TEST 5: Unlocked PDF reports isEncrypted == False."""
|
|
unlock_resp = client.post(f"/documents/{protected_doc_id}/unlock")
|
|
assert unlock_resp.status_code == 200
|
|
perms = unlock_resp.json()["permissions"]
|
|
assert perms["isEncrypted"] is False
|
|
assert perms["encryption"] == "None"
|
|
|
|
|
|
def test_unlocked_pdf_export(protected_doc_id):
|
|
"""TEST 6: Unlocked PDF exports successfully."""
|
|
client.post(f"/documents/{protected_doc_id}/unlock")
|
|
export_resp = client.get(f"/documents/{protected_doc_id}/export")
|
|
assert export_resp.status_code == 200
|
|
assert export_resp.content.startswith(b"%PDF-")
|
|
|
|
|
|
def test_exported_unlocked_pdf_reopens_without_password(protected_doc_id):
|
|
"""TEST 7: Exported unlocked PDF reopens without password."""
|
|
client.post(f"/documents/{protected_doc_id}/unlock")
|
|
export_bytes = client.get(f"/documents/{protected_doc_id}/export").content
|
|
|
|
# Re-open exported bytes without password
|
|
reopen_resp = client.post(
|
|
"/documents",
|
|
files={"file": ("unlocked.pdf", export_bytes, "application/pdf")},
|
|
)
|
|
assert reopen_resp.status_code == 201
|
|
assert reopen_resp.json()["permissions"]["isEncrypted"] is False
|
|
|
|
|
|
def test_unlock_edited_pdf(protected_doc_id):
|
|
"""TEST 8: Unlock edited PDF preserves edits & content."""
|
|
# Apply text edit
|
|
edit_payload = {
|
|
"version": "1.0",
|
|
"operations": [
|
|
{
|
|
"id": "text_edit_1",
|
|
"type": "free_text",
|
|
"pageIndex": 0,
|
|
"data": {
|
|
"x": 50,
|
|
"y": 50,
|
|
"width": 200,
|
|
"height": 40,
|
|
"text": "Preserved Edit",
|
|
"fontSize": 14,
|
|
"color": "#0000ff",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
edit_resp = client.post(f"/documents/{protected_doc_id}/edits", json=edit_payload)
|
|
assert edit_resp.status_code == 200, edit_resp.text
|
|
edited_doc_id = edit_resp.json()["newDocumentId"]
|
|
|
|
# Unlock edited PDF
|
|
unlock_resp = client.post(f"/documents/{edited_doc_id}/unlock")
|
|
assert unlock_resp.status_code == 200
|
|
assert unlock_resp.json()["permissions"]["isEncrypted"] is False
|
|
|
|
|
|
def test_unlock_failure_rollback(protected_doc_id):
|
|
"""TEST 9: Verification failure preserves original encrypted document in store."""
|
|
original_doc = document_store.get_document(protected_doc_id)
|
|
original_bytes = original_doc["bytes_data"]
|
|
|
|
# Try unlocking an invalid doc ID
|
|
fail_resp = client.post("/documents/invalid-id/unlock")
|
|
assert fail_resp.status_code == 404
|
|
|
|
# Document in store remains intact and encrypted
|
|
current_doc = document_store.get_document(protected_doc_id)
|
|
assert current_doc["bytes_data"] == original_bytes
|
|
assert current_doc["permissions"]["isEncrypted"] is True
|
|
|
|
|
|
def test_unprotected_pdf_unlock_rejected():
|
|
"""TEST 10: Unprotected PDF unlock request is rejected (400)."""
|
|
with open(CORPUS_PDF, "rb") as f:
|
|
resp = client.post("/documents", files={"file": ("hello.pdf", f, "application/pdf")})
|
|
doc_id = resp.json()["id"]
|
|
|
|
unlock_resp = client.post(f"/documents/{doc_id}/unlock")
|
|
assert unlock_resp.status_code == 400
|
|
assert "not password protected" in unlock_resp.json()["detail"]
|
|
|
|
|
|
def test_protect_pdf_again_after_unlock(protected_doc_id):
|
|
"""TEST 11: Protect PDF works again after unlocking."""
|
|
# 1. Unlock
|
|
client.post(f"/documents/{protected_doc_id}/unlock")
|
|
|
|
# 2. Protect again with new password
|
|
new_protect_payload = {
|
|
"userPassword": "NewUserPass456!",
|
|
"confirmPassword": "NewUserPass456!",
|
|
}
|
|
prot_resp = client.post(f"/documents/{protected_doc_id}/protect", json=new_protect_payload)
|
|
assert prot_resp.status_code == 200
|
|
assert prot_resp.json()["permissions"]["isEncrypted"] is True
|
|
assert prot_resp.json()["permissions"]["encryption"] == "AES-256"
|
|
|
|
|
|
def test_pdf_security_state_post_unlock(protected_doc_id):
|
|
"""TEST 12: PDF security state metadata is correct post-unlock."""
|
|
unlock_resp = client.post(f"/documents/{protected_doc_id}/unlock")
|
|
assert unlock_resp.status_code == 200
|
|
p = unlock_resp.json()["permissions"]
|
|
assert p["isEncrypted"] is False
|
|
assert p["encryption"] == "None"
|
|
assert p["securityRevision"] == -1
|
|
assert p["ownerUnlocked"] is False
|