diff --git a/passifypdf/encryptpdf.py b/passifypdf/encryptpdf.py index e326cab..7422d71 100644 --- a/passifypdf/encryptpdf.py +++ b/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()) \ No newline at end of file + sys.exit(main()) diff --git a/tests/integrationtests/test_encryptpdf_integration.py b/tests/integrationtests/test_encryptpdf_integration.py index 7cebc00..5cc1a55 100644 --- a/tests/integrationtests/test_encryptpdf_integration.py +++ b/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, + ) diff --git a/tests/unittests/test_encryptpdf.py b/tests/unittests/test_encryptpdf.py index d8d2aa9..3b2c359 100644 --- a/tests/unittests/test_encryptpdf.py +++ b/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")