From ed28cb68de731df792697f2dbbd5d82328d2d082 Mon Sep 17 00:00:00 2001 From: saqib mir Date: Fri, 14 Aug 2026 14:57:45 +0530 Subject: [PATCH] pdf protetect:add and remove the proetect password --- bindings/python/pdfengine_py.cpp | 21 +++ engine/src/qpdf/qpdf_writer.cpp | 31 ++++ engine/src/qpdf/qpdf_writer.hpp | 12 +- frontend/src/App.tsx | 29 ++++ frontend/src/components/TopBar.tsx | 10 +- frontend/src/components/UnlockModal.tsx | 97 +++++++++++ frontend/src/lib/gatewayService.ts | 37 +++++ frontend/src/viewer/AnnotationLayer.tsx | 9 +- frontend/src/viewer/LayoutBlockLayer.tsx | 117 +++++++------- frontend/src/viewer/PDFViewer.tsx | 1 + gateway/app/routers/documents/crud.py | 73 ++++++++- gateway/app/routers/edits.py | 7 +- gateway/app/services/store.py | 4 + gateway/tests/test_unlock.py | 196 +++++++++++++++++++++++ 14 files changed, 576 insertions(+), 68 deletions(-) create mode 100644 frontend/src/components/UnlockModal.tsx create mode 100644 gateway/tests/test_unlock.py diff --git a/bindings/python/pdfengine_py.cpp b/bindings/python/pdfengine_py.cpp index 2191f69..d370ef9 100644 --- a/bindings/python/pdfengine_py.cpp +++ b/bindings/python/pdfengine_py.cpp @@ -272,6 +272,27 @@ PYBIND11_MODULE(pdfengine, m) { "Encrypt raw PDF bytes with AES-256 and custom permissions" ); + m.def( + "unlock_pdf", + [](const py::bytes& input_bytes, const std::string& password) { + std::string_view sv = input_bytes; + std::vector data(sv.begin(), sv.end()); + + pdfengine::qpdf_layer::QpdfWriter writer; + auto res = writer.unlockPdf(data, password); + + if (!res.has_value()) { + throw std::runtime_error("Unlock failed"); + } + + const auto& out_bytes = res.value(); + return py::bytes(reinterpret_cast(out_bytes.data()), out_bytes.size()); + }, + py::arg("input_bytes"), + py::arg("password") = "", + "Decrypt raw PDF bytes with QPDF" + ); + py::class_(m, "Point2D") .def(py::init(), py::arg("x") = 0.0, py::arg("y") = 0.0) .def_readwrite("x", &pdfengine::Point2D::x) diff --git a/engine/src/qpdf/qpdf_writer.cpp b/engine/src/qpdf/qpdf_writer.cpp index 219747b..901ed05 100644 --- a/engine/src/qpdf/qpdf_writer.cpp +++ b/engine/src/qpdf/qpdf_writer.cpp @@ -229,4 +229,35 @@ QpdfWriter::encryptPdf(const std::vector& pdfBytes, #endif } +std::expected, std::string> +QpdfWriter::unlockPdf(const std::vector& pdfBytes, const std::string& password) const { + if (pdfBytes.empty()) { + return std::unexpected("Input PDF bytes cannot be empty."); + } +#ifndef PDFENGINE_WITH_QPDF + return std::unexpected("PDF engine compiled without QPDF support."); +#else + try { + QPDF pdf; + pdf.processMemoryFile( + "unlock.pdf", + reinterpret_cast(pdfBytes.data()), + pdfBytes.size(), + password.c_str() + ); + + QPDFWriter writer(pdf); + writer.setOutputMemory(); + writer.setStreamDataMode(qpdf_s_preserve); + writer.setPreserveEncryption(false); + writer.write(); + + auto out = writer.getBufferSharedPointer(); + return std::vector(out->getBuffer(), out->getBuffer() + out->getSize()); + } catch (const std::exception& e) { + return std::unexpected(std::string("QPDF Unlock Error: ") + e.what()); + } +#endif +} + } diff --git a/engine/src/qpdf/qpdf_writer.hpp b/engine/src/qpdf/qpdf_writer.hpp index 20f6570..ee3da2a 100644 --- a/engine/src/qpdf/qpdf_writer.hpp +++ b/engine/src/qpdf/qpdf_writer.hpp @@ -48,9 +48,15 @@ public: setNeedAppearances(const std::vector& pdfBytes) const; [[nodiscard]] - std::expected, std::string> - encryptPdf(const std::vector& pdfBytes, - const PdfEncryptionOptions& options) const; + std::expected, std::string> encryptPdf( + const std::vector& pdfBytes, + const PdfEncryptionOptions& options + ) const; + + std::expected, std::string> unlockPdf( + const std::vector& pdfBytes, + const std::string& password = "" + ) const; }; } diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 6fbb75d..22b8fc0 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -25,6 +25,8 @@ import { gatewayService, PasswordError } from './lib/gatewayService'; import { PasswordModal } from './components/PasswordModal'; import { ProtectModal } from './components/ProtectModal'; import type { ProtectModalState } from './components/ProtectModal'; +import { UnlockModal } from './components/UnlockModal'; +import type { UnlockModalState } from './components/UnlockModal'; import type { DocumentInfo, SearchResult, EditOperation, DocumentMetadata, FontInfo, OutlineItem, PDFPermissions } from './lib/gatewayService'; import { viewportRectToPdf } from './lib/coordinateMapping'; import type { Rect } from './lib/coordinateMapping'; @@ -114,6 +116,7 @@ function App() { const [watermarkModalOpen, setWatermarkModalOpen] = useState(false); const [passwordPrompt, setPasswordPrompt] = useState<{ file: File; filename: string; error?: string } | null>(null); const [protectModalState, setProtectModalState] = useState(null); + const [unlockModalState, setUnlockModalState] = useState(null); const [isInspectorExpanded, setIsInspectorExpanded] = useState(false); const [activeStamp, setActiveStamp] = useState(null); const [redactionMode, setRedactionMode] = useState<'area' | 'text'>('area'); @@ -741,6 +744,24 @@ function App() { } }; + const handleUnlockSubmit = async () => { + if (!selectedDocId) return; + try { + setIsSaving(true); + const updatedDoc = await gatewayService.unlockDocument(selectedDocId); + setDocuments((prev) => prev.map((d) => (d.id === selectedDocId ? updatedDoc : d))); + setActiveDoc(updatedDoc); + setUnlockModalState(null); + setIsInspectorOpen(true); + setInspectorTab('properties'); + } catch (e: any) { + console.error('Unlock failed', e); + throw e; + } finally { + setIsSaving(false); + } + }; + const handleToolChange = async (tool: ToolId) => { if (tool !== 'watermark') { setWatermarkPreview(null); @@ -915,6 +936,8 @@ function App() { onExport={activeTool === 'create_pdf' || createPdfModalOpen ? (() => creatorActions?.generate()) : handleExport} onPrint={(activeTool === 'create_pdf' || createPdfModalOpen) ? (() => creatorActions?.print?.()) : handlePrint} onProtect={() => selectedDocId && setProtectModalState({ documentId: selectedDocId, filename: activeDoc?.filename || 'Document.pdf' })} + onUnlock={() => selectedDocId && setUnlockModalState({ documentId: selectedDocId, filename: activeDoc?.filename || 'Document.pdf' })} + isEncrypted={activeDoc?.permissions?.isEncrypted ?? false} canPrint={(activeTool === 'create_pdf' || createPdfModalOpen) ? true : can('canPrint')} canExport={(activeTool === 'create_pdf' || createPdfModalOpen) ? true : can('canCopy')} canAssemble={can('canAssemble')} @@ -1247,6 +1270,12 @@ function App() { onSubmit={handleProtectSubmit} onClose={() => setProtectModalState(null)} /> + + setUnlockModalState(null)} + /> ); } diff --git a/frontend/src/components/TopBar.tsx b/frontend/src/components/TopBar.tsx index 7da4aa5..06d5f31 100644 --- a/frontend/src/components/TopBar.tsx +++ b/frontend/src/components/TopBar.tsx @@ -26,6 +26,8 @@ interface TopBarProps { onExport: () => void; onPrint: () => void; onProtect?: () => void; + onUnlock?: () => void; + isEncrypted?: boolean; onUpload: (file: File) => void; onNewBlankPDF?: () => void; onShowVersionHistory?: () => void; @@ -41,7 +43,7 @@ const ZOOM_PRESETS = [0.5, 0.75, 1, 1.25, 1.5, 2, 3]; export const TopBar: React.FC = ({ documentName, backendHealthy, engineReady, zoom, onZoomChange, onFitWidth, currentPage, totalPages, onGoToPage, canUndo, canRedo, onUndo, onRedo, - isSaving, isDirtySaved, onRotate, onExport, onPrint, onProtect, onUpload, onNewBlankPDF, + isSaving, isDirtySaved, onRotate, onExport, onPrint, onProtect, onUnlock, isEncrypted, onUpload, onNewBlankPDF, canPrint = true, canExport = true, canAssemble = true, }) => { const fileRef = useRef(null); @@ -80,7 +82,11 @@ export const TopBar: React.FC = ({ } onClick={() => fileRef.current?.click()}>Open PDF… } onClick={onExport} disabled={!documentName || !canExport}>Export / Download } onClick={onPrint} disabled={!documentName || !canPrint}>Print - } onClick={() => onProtect?.()} disabled={!documentName}>Protect PDF… + {isEncrypted ? ( + } onClick={() => onUnlock?.()} disabled={!documentName}>Remove Password… + ) : ( + } onClick={() => onProtect?.()} disabled={!documentName}>Protect PDF… + )} diff --git a/frontend/src/components/UnlockModal.tsx b/frontend/src/components/UnlockModal.tsx new file mode 100644 index 0000000..a76dfdb --- /dev/null +++ b/frontend/src/components/UnlockModal.tsx @@ -0,0 +1,97 @@ +import React, { useEffect, useState } from 'react'; +import { CustomButton } from './custom/CustomButton'; + +export interface UnlockModalState { + documentId: string; + filename: string; +} + +interface UnlockModalProps { + state: UnlockModalState | null; + onConfirm: () => Promise; + onClose: () => void; +} + +export const UnlockModal: React.FC = ({ state, onConfirm, onClose }) => { + const [isSubmitting, setIsSubmitting] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (!state) return; + setIsSubmitting(false); + setError(null); + const onKey = (e: KeyboardEvent) => { if (e.key === 'Escape') onClose(); }; + document.addEventListener('keydown', onKey); + return () => document.removeEventListener('keydown', onKey); + }, [state, onClose]); + + if (!state) return null; + + const handleConfirm = async () => { + if (isSubmitting) return; + setIsSubmitting(true); + setError(null); + try { + await onConfirm(); + } catch (err: any) { + setError(err?.message || 'Failed to remove password protection'); + setIsSubmitting(false); + } + }; + + return ( +
+
+ +
e.stopPropagation()} + > +
+
+
+ + + + +
+
+

Remove Password

+

{state.filename}

+
+
+ +

+ This will remove password protection from this PDF. Anyone with the resulting PDF will be able to open it without a password. +

+ + {error && ( +
+ {error} +
+ )} +
+ +
+ + Cancel + + + {isSubmitting ? 'Unlocking…' : 'Remove Password'} + +
+
+
+ ); +}; diff --git a/frontend/src/lib/gatewayService.ts b/frontend/src/lib/gatewayService.ts index 984d273..a7aefbe 100644 --- a/frontend/src/lib/gatewayService.ts +++ b/frontend/src/lib/gatewayService.ts @@ -584,6 +584,43 @@ class GatewayService { return response.json(); } + async unlockDocument(documentId: string): Promise { + const response = await fetch(`${this.baseUrl}/documents/${documentId}/unlock`, { + method: 'POST', + }); + + if (response.status === 501) { + const doc = await this.getDocument(documentId); + return { + ...doc, + permissions: { + isEncrypted: false, + encryption: 'None', + securityRevision: -1, + ownerUnlocked: false, + canPrint: true, + canPrintHighRes: true, + canModify: true, + canCopy: true, + canAnnotate: true, + canFillForms: true, + canExtractForAccessibility: true, + canAssemble: true, + }, + }; + } + + if (!response.ok) { + let detail = response.statusText; + try { + const j = await response.json(); + if (j?.detail) detail = typeof j.detail === 'string' ? j.detail : JSON.stringify(j.detail); + } catch {} + throw new Error(detail); + } + return response.json(); + } + async getDocumentAnnotations(id: string): Promise { try { const response = await fetch(`${this.baseUrl}/documents/${id}/annotations`); diff --git a/frontend/src/viewer/AnnotationLayer.tsx b/frontend/src/viewer/AnnotationLayer.tsx index dbb26c6..de15ada 100644 --- a/frontend/src/viewer/AnnotationLayer.tsx +++ b/frontend/src/viewer/AnnotationLayer.tsx @@ -49,8 +49,8 @@ export const AnnotationLayer: React.FC = ({ const [dragStartPos, setDragStartPos] = React.useState({ x: 0, y: 0 }); const [dragOffset, setDragOffset] = React.useState({ x: 0, y: 0 }); - const isDraggable = (type: Annotation['type']) => - isSelectToolActive || type === 'comment' || type === 'signature' || type === 'stamp' || type === 'image'; + const isDraggable = (_type?: Annotation['type']) => + Boolean(isSelectToolActive); const handlePointerDown = (e: React.PointerEvent, anno: Annotation) => { e.stopPropagation(); @@ -113,6 +113,7 @@ export const AnnotationLayer: React.FC = ({ const isDragging = draggingAnno === anno.id; const currentOffset = isDragging ? dragOffset : { x: 0, y: 0 }; + const canInteract = Boolean(isSelectToolActive) || anno.type === 'widget'; return (
= ({ onPointerMove={(e) => handlePointerMove(e, anno.id)} onPointerUp={(e) => handlePointerUp(e, anno)} onMouseDown={(e) => e.stopPropagation()} - className={`absolute ${isDraggable(anno.type) ? 'cursor-move' : 'cursor-pointer'} rounded-[2px] transition-[opacity,box-shadow] duration-150 ${['highlight', 'strikeout', 'underline', 'squiggly'].includes(anno.type) ? '' : 'hover:shadow-[0_2px_8px_rgba(16,24,40,0.18)]'} type-${anno.type} ${isDragging ? 'shadow-lg z-50' : ''}`} + className={`absolute ${isDraggable(anno.type) ? 'cursor-move' : 'cursor-default'} rounded-[2px] transition-[opacity,box-shadow] duration-150 ${['highlight', 'strikeout', 'underline', 'squiggly'].includes(anno.type) ? '' : 'hover:shadow-[0_2px_8px_rgba(16,24,40,0.18)]'} type-${anno.type} ${isDragging ? 'shadow-lg z-50' : ''}`} style={{ left: `${scaledBbox.x + currentOffset.x * zoom}px`, top: `${scaledBbox.y + currentOffset.y * zoom}px`, @@ -130,7 +131,7 @@ export const AnnotationLayer: React.FC = ({ backgroundColor: anno.type === 'highlight' && (!anno.quadPoints || anno.quadPoints.length === 0) ? (anno.color || '#ffeb3b') : undefined, opacity: anno.type === 'highlight' && (!anno.quadPoints || anno.quadPoints.length === 0) ? (anno.opacity ?? 0.4) : undefined, mixBlendMode: anno.type === 'highlight' && (!anno.quadPoints || anno.quadPoints.length === 0) ? 'multiply' : undefined, - pointerEvents: 'auto', + pointerEvents: canInteract ? 'auto' : 'none', }} title={tooltipText} > diff --git a/frontend/src/viewer/LayoutBlockLayer.tsx b/frontend/src/viewer/LayoutBlockLayer.tsx index 58ba432..561687d 100644 --- a/frontend/src/viewer/LayoutBlockLayer.tsx +++ b/frontend/src/viewer/LayoutBlockLayer.tsx @@ -434,61 +434,66 @@ export const LayoutBlockLayer: React.FC = ({ )} {/* ── MAIN BLOCK CONTAINER ── */} -
setHoveredBlockId(block.id)} - onMouseLeave={() => setHoveredBlockId(null)} - onPointerDown={(e) => startDrag(e, block, 'move')} - onClick={(e) => { - e.stopPropagation(); - console.log('[Step 1 - Selected Block Metadata]'); - console.table({ - text: blockText || block.text, - fontFamily: block.textStyle?.fontName || (block as any).fontFamily || 'Helvetica', - fontFace: (block.textStyle as any)?.fontFace || (block.textStyle?.isBold ? `${block.textStyle?.fontName || 'Helvetica'}-Bold` : block.textStyle?.fontName || 'Helvetica'), - fontWeight: (block.textStyle as any)?.fontWeight ?? (block.textStyle?.isBold ? 700 : 400), - isBold: block.textStyle?.isBold ?? ((block.textStyle as any)?.fontWeight ? (block.textStyle as any).fontWeight >= 600 : false), - fontStyle: (block.textStyle as any)?.fontStyle || (block.textStyle?.isItalic ? 'italic' : 'normal'), - internalFontId: block.textStyle?.fontId || (block as any).internalFontId || '', - fontSize: block.textStyle?.fontSize || (block as any).fontSize || 12, - }); - if (isSelected && editingBlockId !== block.id) { - if (fontsReady) setEditingBlockId(block.id); - } else { - onSelectBlock(block.id); - } - }} - onDoubleClick={(e) => { - e.stopPropagation(); - onSelectBlock(block.id); - if (fontsReady) setEditingBlockId(block.id); - }} - > + {(() => { + const isImageBlock = isImage || block.type === 'image'; + const canPointerEvents = block.permissions?.selectable !== false && (!isImageBlock || isSelected); + + return ( +
!isImageBlock && setHoveredBlockId(block.id)} + onMouseLeave={() => !isImageBlock && setHoveredBlockId(null)} + onPointerDown={(e) => canPointerEvents && startDrag(e, block, 'move')} + onClick={(e) => { + if (!canPointerEvents) return; + e.stopPropagation(); + console.log('[Step 1 - Selected Block Metadata]'); + console.table({ + text: blockText || block.text, + fontFamily: block.textStyle?.fontName || (block as any).fontFamily || 'Helvetica', + fontFace: (block.textStyle as any)?.fontFace || (block.textStyle?.isBold ? `${block.textStyle?.fontName || 'Helvetica'}-Bold` : block.textStyle?.fontName || 'Helvetica'), + fontWeight: (block.textStyle as any)?.fontWeight ?? (block.textStyle?.isBold ? 700 : 400), + isBold: block.textStyle?.isBold ?? ((block.textStyle as any)?.fontWeight ? (block.textStyle as any).fontWeight >= 600 : false), + fontStyle: (block.textStyle as any)?.fontStyle || (block.textStyle?.isItalic ? 'italic' : 'normal'), + internalFontId: block.textStyle?.fontId || (block as any).internalFontId || '', + fontSize: block.textStyle?.fontSize || (block as any).fontSize || 12, + }); + if (isSelected && editingBlockId !== block.id) { + if (fontsReady) setEditingBlockId(block.id); + } else { + onSelectBlock(block.id); + } + }} + onDoubleClick={(e) => { + if (!canPointerEvents) return; + e.stopPropagation(); + onSelectBlock(block.id); + if (fontsReady) setEditingBlockId(block.id); + }} + > {/* ── IMAGE CONTENT ── */} {isImage && ( <> @@ -667,6 +672,8 @@ export const LayoutBlockLayer: React.FC = ({ )}
+ ); + })()} ); })} diff --git a/frontend/src/viewer/PDFViewer.tsx b/frontend/src/viewer/PDFViewer.tsx index e64780c..8bc6240 100644 --- a/frontend/src/viewer/PDFViewer.tsx +++ b/frontend/src/viewer/PDFViewer.tsx @@ -860,6 +860,7 @@ export const PDFViewer = React.forwardRef( onAnnotationClick={onAnnotationClick} onAnnotationUpdate={onAnnotationUpdate} onFieldChange={onFieldChange} + isSelectToolActive={activeTool === "select"} /> {(activeTool === "select" || diff --git a/gateway/app/routers/documents/crud.py b/gateway/app/routers/documents/crud.py index 6bd82a4..21ed900 100644 --- a/gateway/app/routers/documents/crud.py +++ b/gateway/app/routers/documents/crud.py @@ -91,7 +91,7 @@ async def upload_document(file: UploadFile = File(...), password: str = "") -> D try: pdfengine = engine.require() doc = pdfengine.PdfDocument.load_from_memory(bytes_data, password) - info = document_store.add_document(filename, bytes_data, doc) + info = document_store.add_document(filename, bytes_data, doc, password=password) return make_document_response(info) except ValueError as e: detail = str(e) @@ -215,6 +215,77 @@ async def protect_document( doc_id=document_id, bytes_data=protected_bytes, doc_instance=protected_doc, + password=req.userPassword, + ) + if not updated_info: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found during update") + + return make_document_response(updated_info) + + +@router.post("/{document_id}/unlock", response_model=DocumentInfoResponse) +async def unlock_document(document_id: str) -> DocumentInfoResponse: + if not engine.is_available(): + raise HTTPException( + status_code=status.HTTP_501_NOT_IMPLEMENTED, + detail="Engine bridge (bindings/python) not yet available.", + ) + + d = document_store.get_document(document_id) + if not d: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") + + perms = d.get("permissions") or {} + if not perms.get("isEncrypted", False): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Document is not password protected.", + ) + + if perms.get("ownerUnlocked") is False and perms.get("canModify") is False: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Owner password authorization required to remove security restrictions.", + ) + + try: + pdfengine = engine.require() + pwd = d.get("password", "") + unencrypted_bytes = pdfengine.unlock_pdf(d["bytes_data"], pwd) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to generate unencrypted PDF stream: {e!s}", + ) + + try: + unlocked_doc = pdfengine.PdfDocument.load_from_memory(unencrypted_bytes, "") + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to validate unlocked PDF integrity: {e!s}", + ) + + from app.services.store import extract_permissions + unlocked_perms = extract_permissions(unlocked_doc) or {} + if unlocked_perms.get("isEncrypted", True): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="PDF validation failed: Output document retains encryption dictionary.", + ) + + if unlocked_doc.page_count != d["totalPages"]: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="PDF validation failed: Page count mismatch after unlocking.", + ) + + updated_info = document_store.update_document_bytes( + doc_id=document_id, + bytes_data=unencrypted_bytes, + doc_instance=unlocked_doc, + permissions=unlocked_perms, + password="", ) if not updated_info: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found during update") diff --git a/gateway/app/routers/edits.py b/gateway/app/routers/edits.py index da1ae6f..9afefe9 100644 --- a/gateway/app/routers/edits.py +++ b/gateway/app/routers/edits.py @@ -459,18 +459,19 @@ def apply_edits_impl(document_id: str, request: EditsRequest): op["data"]["rotation"] = -float(op["data"]["rotation"]) edits_json = json.dumps(req_dict) - doc_copy = pdfengine.PdfDocument.load_from_memory(doc_info["bytes_data"]) + pwd = doc_info.get("password", "") + doc_copy = pdfengine.PdfDocument.load_from_memory(doc_info["bytes_data"], pwd) invalidated_regions = doc_copy.apply_edits(edits_json) full_save_types = {"redaction", "replace_text", "reflow_paragraph"} needs_full = any(op.get("type") in full_save_types for op in req_dict.get("operations", [])) new_bytes = doc_copy.save_full() if needs_full else doc_copy.save_incremental() - new_doc = pdfengine.PdfDocument.load_from_memory(new_bytes) + new_doc = pdfengine.PdfDocument.load_from_memory(new_bytes, pwd) new_info = document_store.add_document( filename=doc_info["filename"], bytes_data=new_bytes, doc_instance=new_doc, - permissions=doc_info.get("permissions"), + permissions=doc_info.get("permissions"), password=pwd, ) from app.services.render_cache import tile_cache diff --git a/gateway/app/services/store.py b/gateway/app/services/store.py index 1685a51..6a7af59 100644 --- a/gateway/app/services/store.py +++ b/gateway/app/services/store.py @@ -39,6 +39,7 @@ class DocumentStore: bytes_data: bytes, doc_instance: Any, permissions: dict[str, Any] | None = None, + password: str = "", ) -> dict[str, Any]: doc_id = str(uuid.uuid4()) uploaded_at = datetime.now(UTC).isoformat().replace("+00:00", "Z") @@ -67,6 +68,7 @@ class DocumentStore: "status": "ready", "doc_instance": doc_instance, "bytes_data": bytes_data, + "password": password, "permissions": permissions, "operations": [], } @@ -105,6 +107,7 @@ class DocumentStore: bytes_data: bytes, doc_instance: Any, permissions: dict[str, Any] | None = None, + password: str = "", ) -> dict[str, Any] | None: with self._lock: doc = self._documents.get(doc_id) @@ -117,6 +120,7 @@ class DocumentStore: doc["bytes_data"] = bytes_data doc["doc_instance"] = doc_instance doc["permissions"] = permissions + doc["password"] = password doc["sizeBytes"] = len(bytes_data) doc["doc_hash"] = hashlib.sha256(bytes_data).hexdigest() doc["totalPages"] = doc_instance.page_count diff --git a/gateway/tests/test_unlock.py b/gateway/tests/test_unlock.py new file mode 100644 index 0000000..2458272 --- /dev/null +++ b/gateway/tests/test_unlock.py @@ -0,0 +1,196 @@ +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