Add output_dir manipulation methods to DocumentHolder

These will be needed in for the GUI's settings. This also adds test
cases for these documents. The methods are the following:

  - set_output_dir()
    For changing the output directory of the safe file

  - suffix setter and getter - for changing the suffix of the file
This commit is contained in:
deeplow 2022-09-28 12:17:11 +01:00
parent fc3cfba450
commit 5a6c72f09e
No known key found for this signature in database
GPG key ID: 577982871529A52A
3 changed files with 84 additions and 5 deletions

View file

@ -30,12 +30,18 @@ class Document:
STATE_SAFE = enum.auto() STATE_SAFE = enum.auto()
STATE_FAILED = 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 # NOTE: See https://github.com/freedomofpress/dangerzone/pull/216#discussion_r1015449418
self.id = secrets.token_urlsafe(6)[0:6] self.id = secrets.token_urlsafe(6)[0:6]
self._input_filename: Optional[str] = None self._input_filename: Optional[str] = None
self._output_filename: Optional[str] = None self._output_filename: Optional[str] = None
self._suffix = suffix
if input_filename: if input_filename:
self.input_filename = input_filename self.input_filename = input_filename
@ -97,13 +103,35 @@ class Document:
self.validate_output_filename(filename) self.validate_output_filename(filename)
self._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 @property
def default_output_filename(self) -> str: 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: def announce_id(self) -> None:
log.info(f"Assigning ID '{self.id}' to doc '{self.input_filename}'") 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: def is_unconverted(self) -> bool:
return self.state is Document.STATE_UNCONVERTED return self.state is Document.STATE_UNCONVERTED

View file

@ -57,6 +57,20 @@ class NotSetOutputFilenameException(DocumentFilenameException):
super().__init__("Output filename has not been set yet.") 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: def handle_document_errors(func: F) -> F:
"""Log document-related errors and exit gracefully.""" """Log document-related errors and exit gracefully."""

View file

@ -1,12 +1,13 @@
import os import os
import platform import platform
import stat
import tempfile import tempfile
from pathlib import Path from pathlib import Path
import pytest import pytest
from dangerzone import errors from dangerzone import errors
from dangerzone.document import Document from dangerzone.document import SAFE_EXTENSION, Document
from . import sample_doc, unreadable_pdf 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: 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() d = Document()
with pytest.raises(errors.NonPDFOutputFileException) as e: 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) 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) d = Document(sample_doc)
assert d.is_unconverted() assert d.is_unconverted()