fix: validate input is a real PDF before encryption (closes #24)

passifypdf/encryptpdf.py

@@ -5,6 +5,7 @@
 from typing import Union
 
 from pypdf import PdfReader, PdfWriter
+from pypdf.errors import PdfReadError
 
 from .cli import get_arg_parser
 
@@ -21,7 +22,7 @@ def encrypt_pdf(input_pdf: Union[str, Path], output_pdf: Union[str, Path], passw
     Raises:
         FileNotFoundError: If the input PDF file does not exist.
         IsADirectoryError: If the input path is a directory.
-        Exception: For other errors during processing.
+        ValueError: If the input file is not a valid PDF.
     """
     input_path = Path(input_pdf)
     if not input_path.exists():
@@ -30,23 +31,28 @@ def encrypt_pdf(input_pdf: Union[str, Path], output_pdf: Union[str, Path], passw
     if not input_path.is_file():
         raise IsADirectoryError(f"Input path '{input_pdf}' is not a file.")
 
+    invalid_pdf_error = ValueError(f"Input file '{input_pdf}' is not a valid PDF file.")
+    with input_path.open("rb") as source_file:
+        if source_file.read(5) != b"%PDF-":
+            raise invalid_pdf_error
+
     try:
         reader = PdfReader(input_path)
-        writer = PdfWriter()
+    except PdfReadError as exc:
+        raise invalid_pdf_error from exc
 
-        # Add all pages from the reader to the writer
-        for page in reader.pages:
-            writer.add_page(page)
+    writer = PdfWriter()
 
-        # Encrypt with a password
-        writer.encrypt(password)
+    # Add all pages from the reader to the writer
+    for page in reader.pages:
+        writer.add_page(page)
 
-        # Write the encrypted PDF to a new PDF file passed as param
-        with open(output_pdf, "wb") as f:
-            writer.write(f)
+    # Encrypt with a password
+    writer.encrypt(password)
 
-    except Exception as e:
-        raise Exception(f"Failed to encrypt PDF: {e}")
+    # Write the encrypted PDF to a new PDF file passed as param
+    with open(output_pdf, "wb") as f:
+        writer.write(f)
 
 
 def main() -> int:
@@ -69,4 +75,4 @@ def main() -> int:
 
 
 if __name__ == "__main__":
-    sys.exit(main())
+    sys.exit(main())

tests/integrationtests/test_encryptpdf_integration.py

@@ -1,4 +1,8 @@
 import os
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
 from unittest import TestCase
 
 from pypdf import PdfReader, PdfWriter
@@ -12,30 +16,86 @@ def setUp(self):
         self.output_pdf = "test_output.pdf"
         self.password = "strongpassword"
 
-        # Create a dummy PDF
         writer = PdfWriter()
         writer.add_blank_page(width=100, height=100)
         with open(self.input_pdf, "wb") as f:
             writer.write(f)
 
     def tearDown(self):
-        # Cleanup files
         if os.path.exists(self.input_pdf):
             os.remove(self.input_pdf)
         if os.path.exists(self.output_pdf):
             os.remove(self.output_pdf)
 
     def test_encrypt_pdf_integration(self):
-        # Encrypt the PDF
         encrypt_pdf(self.input_pdf, self.output_pdf, self.password)
 
-        # Verify output exists
         self.assertTrue(os.path.exists(self.output_pdf))
 
-        # Verify it is encrypted
         reader = PdfReader(self.output_pdf)
         self.assertTrue(reader.is_encrypted)
 
-        # Verify we can decrypt it
         self.assertTrue(reader.decrypt(self.password))
         self.assertEqual(len(reader.pages), 1)
+
+    def test_cli_rejects_fake_pdf_extension(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            input_path = temp_path / "fake.pdf"
+            output_path = temp_path / "out.pdf"
+            input_path.write_bytes(b"not a pdf at all")
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "passifypdf.encryptpdf",
+                    "-i",
+                    str(input_path),
+                    "-o",
+                    str(output_path),
+                    "-p",
+                    "secret123",
+                ],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+
+            self.assertEqual(result.returncode, 1)
+            self.assertFalse(output_path.exists())
+            self.assertIn(
+                f"Error: Input file '{input_path}' is not a valid PDF file.",
+                result.stderr,
+            )
+
+    def test_cli_rejects_fake_pdf_without_extension(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            input_path = temp_path / "fake_no_ext"
+            output_path = temp_path / "out.pdf"
+            input_path.write_bytes(b"not a pdf at all")
+
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "passifypdf.encryptpdf",
+                    "-i",
+                    str(input_path),
+                    "-o",
+                    str(output_path),
+                    "-p",
+                    "secret123",
+                ],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+
+            self.assertEqual(result.returncode, 1)
+            self.assertFalse(output_path.exists())
+            self.assertIn(
+                f"Error: Input file '{input_path}' is not a valid PDF file.",
+                result.stderr,
+            )

tests/unittests/test_encryptpdf.py

@@ -1,61 +1,89 @@
 from unittest import TestCase
-from unittest.mock import patch, mock_open
+from unittest.mock import MagicMock, mock_open, patch
+
+from pypdf.errors import PdfReadError
+
 from passifypdf.encryptpdf import encrypt_pdf
 
 
 class TestPdfUnitTests(TestCase):
 
-    @patch('passifypdf.encryptpdf.PdfReader')
-    @patch('passifypdf.encryptpdf.PdfWriter')
-    @patch('passifypdf.encryptpdf.Path')  # Mock Path
-    @patch('builtins.open', new_callable=mock_open)
+    @patch("passifypdf.encryptpdf.PdfReader")
+    @patch("passifypdf.encryptpdf.PdfWriter")
+    @patch("passifypdf.encryptpdf.Path")
+    @patch("builtins.open", new_callable=mock_open)
     def test_encrypt_pdf(self, mock_file, mock_path_cls, mock_writer_cls, mock_reader_cls):
-        # Setup mocks
         mock_path_instance = mock_path_cls.return_value
         mock_path_instance.exists.return_value = True
         mock_path_instance.is_file.return_value = True
 
+        mock_source_file = MagicMock()
+        mock_source_file.read.return_value = b"%PDF-"
+        mock_path_instance.open.return_value.__enter__.return_value = mock_source_file
+
         mock_reader_instance = mock_reader_cls.return_value
-        mock_reader_instance.pages = ['page1', 'page2']
-        
+        mock_reader_instance.pages = ["page1", "page2"]
+
         mock_writer_instance = mock_writer_cls.return_value
-
-        # Call the function
+
         encrypt_pdf("input.pdf", "output.pdf", "secret")
-
-        # Verify Path existence check
+
         mock_path_cls.assert_called_with("input.pdf")
-        mock_path_instance.exists.assert_called()
-        mock_path_instance.is_file.assert_called()
+        mock_path_instance.exists.assert_called_once()
+        mock_path_instance.is_file.assert_called_once()
+        mock_path_instance.open.assert_called_once_with("rb")
 
-        # Verify PdfReader was called with the path object
         mock_reader_cls.assert_called_with(mock_path_instance)
-
-        # Verify pages were added
+
         self.assertEqual(mock_writer_instance.add_page.call_count, 2)
-        mock_writer_instance.add_page.assert_any_call('page1')
-        mock_writer_instance.add_page.assert_any_call('page2')
-
-        # Verify encryption
+        mock_writer_instance.add_page.assert_any_call("page1")
+        mock_writer_instance.add_page.assert_any_call("page2")
+
         mock_writer_instance.encrypt.assert_called_with("secret")
-
-        # Verify file write
+
         mock_file.assert_called_with("output.pdf", "wb")
         mock_writer_instance.write.assert_called_with(mock_file())
 
-    @patch('passifypdf.encryptpdf.Path')
+    @patch("passifypdf.encryptpdf.Path")
     def test_encrypt_pdf_file_not_found(self, mock_path_cls):
         mock_path_instance = mock_path_cls.return_value
         mock_path_instance.exists.return_value = False
-        
+
         with self.assertRaises(FileNotFoundError):
             encrypt_pdf("nonexistent.pdf", "output.pdf", "secret")
 
-    @patch('passifypdf.encryptpdf.Path')
+    @patch("passifypdf.encryptpdf.Path")
     def test_encrypt_pdf_is_directory(self, mock_path_cls):
         mock_path_instance = mock_path_cls.return_value
         mock_path_instance.exists.return_value = True
         mock_path_instance.is_file.return_value = False
-        
+
         with self.assertRaises(IsADirectoryError):
             encrypt_pdf("directory", "output.pdf", "secret")
+
+    @patch("passifypdf.encryptpdf.Path")
+    def test_encrypt_pdf_rejects_non_pdf_magic_bytes(self, mock_path_cls):
+        mock_path_instance = mock_path_cls.return_value
+        mock_path_instance.exists.return_value = True
+        mock_path_instance.is_file.return_value = True
+
+        mock_source_file = MagicMock()
+        mock_source_file.read.return_value = b"NOTPD"
+        mock_path_instance.open.return_value.__enter__.return_value = mock_source_file
+
+        with self.assertRaisesRegex(ValueError, "Input file 'input.pdf' is not a valid PDF file\\."):
+            encrypt_pdf("input.pdf", "output.pdf", "secret")
+
+    @patch("passifypdf.encryptpdf.PdfReader", side_effect=PdfReadError("broken pdf"))
+    @patch("passifypdf.encryptpdf.Path")
+    def test_encrypt_pdf_rejects_pdfreader_parse_failure(self, mock_path_cls, _mock_reader):
+        mock_path_instance = mock_path_cls.return_value
+        mock_path_instance.exists.return_value = True
+        mock_path_instance.is_file.return_value = True
+
+        mock_source_file = MagicMock()
+        mock_source_file.read.return_value = b"%PDF-"
+        mock_path_instance.open.return_value.__enter__.return_value = mock_source_file
+
+        with self.assertRaisesRegex(ValueError, "Input file 'input.pdf' is not a valid PDF file\\."):
+            encrypt_pdf("input.pdf", "output.pdf", "secret")