diff --git a/dangerzone/document.py b/dangerzone/document.py index d1e311c..415b92d 100644 --- a/dangerzone/document.py +++ b/dangerzone/document.py @@ -30,12 +30,18 @@ class Document: STATE_SAFE = enum.auto() STATE_FAILED = enum.auto() - def __init__(self, input_filename: str = None, output_filename: str = None) -> None: + def __init__( + self, + input_filename: str = None, + output_filename: str = None, + suffix: str = SAFE_EXTENSION, + ) -> None: # NOTE: See https://github.com/freedomofpress/dangerzone/pull/216#discussion_r1015449418 self.id = secrets.token_urlsafe(6)[0:6] self._input_filename: Optional[str] = None self._output_filename: Optional[str] = None + self._suffix = suffix if input_filename: self.input_filename = input_filename @@ -97,13 +103,35 @@ class Document: self.validate_output_filename(filename) self._output_filename = filename + @property + def suffix(self) -> str: + return self._suffix + + @suffix.setter + def suffix(self, suf: str) -> None: + self._suffix = suf + @property def default_output_filename(self) -> str: - return f"{os.path.splitext(self.input_filename)[0]}{SAFE_EXTENSION}" + return f"{os.path.splitext(self.input_filename)[0]}{self.suffix}" def announce_id(self) -> None: log.info(f"Assigning ID '{self.id}' to doc '{self.input_filename}'") + def set_output_dir(self, path: str) -> None: + # keep the same name + old_filename = os.path.basename(self.output_filename) + + new_path = os.path.abspath(path) + if not os.path.exists(new_path): + raise errors.NonExistantOutputDirException() + if not os.path.isdir(new_path): + raise errors.OutputDirIsNotDirException() + if not os.access(new_path, os.W_OK): + raise errors.UnwriteableOutputDirException() + + self._output_filename = os.path.join(new_path, old_filename) + def is_unconverted(self) -> bool: return self.state is Document.STATE_UNCONVERTED diff --git a/dangerzone/errors.py b/dangerzone/errors.py index 1e6811a..e012235 100644 --- a/dangerzone/errors.py +++ b/dangerzone/errors.py @@ -57,6 +57,20 @@ class NotSetOutputFilenameException(DocumentFilenameException): super().__init__("Output filename has not been set yet.") +class NonExistantOutputDirException(DocumentFilenameException): + """Exception for when the output dir does not exist.""" + + def __init__(self) -> None: + super().__init__("Output directory does not exist") + + +class OutputDirIsNotDirException(DocumentFilenameException): + """Exception for when the specified output dir is not actually a dir.""" + + def __init__(self) -> None: + super().__init__("Specified output directory is actually not a directory") + + def handle_document_errors(func: F) -> F: """Log document-related errors and exit gracefully.""" diff --git a/tests/test_document.py b/tests/test_document.py index 3b0bdba..264bcdd 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -1,12 +1,13 @@ import os import platform +import stat import tempfile from pathlib import Path import pytest from dangerzone import errors -from dangerzone.document import Document +from dangerzone.document import SAFE_EXTENSION, Document from . import sample_doc, unreadable_pdf @@ -68,7 +69,7 @@ def test_output_file_none() -> None: def test_output_file_not_pdf(tmp_path: Path) -> None: - docx_file = str(tmp_path.joinpath("document.docx")) + docx_file = str(tmp_path / "document.docx") d = Document() with pytest.raises(errors.NonPDFOutputFileException) as e: @@ -77,7 +78,43 @@ def test_output_file_not_pdf(tmp_path: Path) -> None: assert not os.path.exists(docx_file) -def test_is_unconverted_by_default(sample_doc: None) -> None: +def test_set_output_dir(sample_doc: str, tmp_path: Path) -> None: + d = Document(sample_doc) + d.set_output_dir(str(tmp_path)) + assert os.path.dirname(d.output_filename) == str(tmp_path) + + +def test_set_output_dir_non_existant(sample_doc: str, tmp_path: Path) -> None: + non_existant_path = str(tmp_path / "fake-dir") + d = Document(sample_doc) + with pytest.raises(errors.NonExistantOutputDirException): + d.set_output_dir(non_existant_path) + + +def test_set_output_dir_is_file(sample_doc: str, tmp_path: Path) -> None: + # create a file + file_path = str(tmp_path / "file") + with open(file_path, "w"): + pass + + d = Document(sample_doc) + with pytest.raises(errors.OutputDirIsNotDirException): + d.set_output_dir(file_path) + + +def test_default_output_filename(sample_doc: str) -> None: + d = Document(sample_doc) + assert d.output_filename.endswith(SAFE_EXTENSION) + + +def test_set_output_filename_suffix(sample_doc: str) -> None: + d = Document(sample_doc) + safe_extension = "-trusted.pdf" + d.suffix = safe_extension + assert d.output_filename.endswith(safe_extension) + + +def test_is_unconverted_by_default(sample_doc: str) -> None: d = Document(sample_doc) assert d.is_unconverted()