From db17bd09157348a9ae654b4d7963aee4281c8fff Mon Sep 17 00:00:00 2001 From: deeplow Date: Thu, 6 Oct 2022 10:00:32 +0100 Subject: [PATCH] Validate I/O filenames in Document class Factor out the filename validation logic and move it into the Document class. Previously, the filename validation logic was scattered across the CLI and GUI code. Also, introduce a new errors.py module whose purpose is to handle document-related errors, by providing: * A special exception for them (DocumentFilenameExcpetion) * A decorator that handles DocumentFilenameException, logs it and the underlying cause, and exits the program gracefully. --- dangerzone/cli.py | 39 ++------------------------------------ dangerzone/document.py | 31 ++++++++++++++++++++++++++++-- dangerzone/errors.py | 29 ++++++++++++++++++++++++++++ dangerzone/gui/__init__.py | 27 ++++++++------------------ 4 files changed, 68 insertions(+), 58 deletions(-) create mode 100644 dangerzone/errors.py diff --git a/dangerzone/cli.py b/dangerzone/cli.py index 9b16af0..1c008b2 100644 --- a/dangerzone/cli.py +++ b/dangerzone/cli.py @@ -7,7 +7,7 @@ from typing import Optional import click from colorama import Back, Fore, Style -from . import container +from . import container, errors from .container import convert from .document import Document from .global_common import GlobalCommon @@ -23,6 +23,7 @@ def print_header(s: str) -> None: @click.option("--output-filename", help="Default is filename ending with -safe.pdf") @click.option("--ocr-lang", help="Language to OCR, defaults to none") @click.argument("filename", required=True) +@errors.handle_document_errors def cli_main( output_filename: Optional[str], ocr_lang: Optional[str], filename: str ) -> None: @@ -31,51 +32,15 @@ def cli_main( display_banner() - # Validate filename - valid = True - try: - with open(os.path.abspath(filename), "rb") as f: - pass - except: - valid = False - - if not valid: - click.echo("Invalid filename") - exit(1) - document = Document(os.path.abspath(filename)) # Validate safe PDF output filename if output_filename: - valid = True - if not output_filename.endswith(".pdf"): - click.echo("Safe PDF filename must end in '.pdf'") - exit(1) - - try: - with open(os.path.abspath(output_filename), "wb"): - pass - except: - valid = False - - if not valid: - click.echo("Safe PDF filename is not writable") - exit(1) - document.output_filename = os.path.abspath(output_filename) - else: document.output_filename = ( f"{os.path.splitext(document.input_filename)[0]}-safe.pdf" ) - try: - with open(document.output_filename, "wb"): - pass - except: - click.echo( - f"Output filename {document.output_filename} is not writable, use --output-filename" - ) - exit(1) # Validate OCR language if ocr_lang: diff --git a/dangerzone/document.py b/dangerzone/document.py index c3f4b27..92016cc 100644 --- a/dangerzone/document.py +++ b/dangerzone/document.py @@ -6,6 +6,8 @@ from typing import Optional import appdirs +from .errors import DocumentFilenameException + class Document: """Track the state of a single document. @@ -21,24 +23,49 @@ class Document: if input_filename: self.input_filename = input_filename + @staticmethod + def validate_input_filename(filename: str) -> None: + try: + open(filename, "rb") + except FileNotFoundError as e: + raise DocumentFilenameException( + "Input file not found: make sure you typed it correctly." + ) from e + except PermissionError as e: + raise DocumentFilenameException( + "You don't have permission to open the input file." + ) from e + + @staticmethod + def validate_output_filename(filename: str) -> None: + if not filename.endswith(".pdf"): + raise DocumentFilenameException("Safe PDF filename must end in '.pdf'") + try: + with open(os.path.abspath(filename), "wb"): + pass + except PermissionError as e: + raise DocumentFilenameException("Safe PDF filename is not writable") from e + @property def input_filename(self) -> str: if self._input_filename is None: - raise RuntimeError("Input filename has not been set yet.") + raise DocumentFilenameException("Input filename has not been set yet.") else: return self._input_filename @input_filename.setter def input_filename(self, filename: str) -> None: + self.validate_input_filename(filename) self._input_filename = filename @property def output_filename(self) -> str: if self._output_filename is None: - raise RuntimeError("Output filename has not been set yet.") + raise DocumentFilenameException("Output filename has not been set yet.") else: return self._output_filename @output_filename.setter def output_filename(self, filename: str) -> None: + self.validate_output_filename(filename) self._output_filename = filename diff --git a/dangerzone/errors.py b/dangerzone/errors.py new file mode 100644 index 0000000..5f6c84d --- /dev/null +++ b/dangerzone/errors.py @@ -0,0 +1,29 @@ +import functools +import logging +import sys + +import click + +log = logging.getLogger(__name__) + + +class DocumentFilenameException(Exception): + """Exception for document-related filename errors.""" + + +def handle_document_errors(func): + """Log document-related errors and exit gracefully.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except DocumentFilenameException as e: + if getattr(sys, "dangerzone_dev", False): + # Show the full traceback only on dev environments. + msg = "An exception occured while validating a document filename" + log.exception(msg) + click.echo(str(e)) + exit(1) + + return wrapper diff --git a/dangerzone/gui/__init__.py b/dangerzone/gui/__init__.py index 47619ad..5ca6fce 100644 --- a/dangerzone/gui/__init__.py +++ b/dangerzone/gui/__init__.py @@ -10,6 +10,7 @@ import click import colorama from PySide2 import QtCore, QtGui, QtWidgets +from .. import errors from ..document import Document from ..global_common import GlobalCommon from .common import GuiCommon @@ -50,6 +51,7 @@ class ApplicationWrapper(QtCore.QObject): @click.command() @click.argument("filename", required=False) +@errors.handle_document_errors def gui_main(filename: Optional[str]) -> bool: setup_logging() @@ -86,36 +88,23 @@ def gui_main(filename: Optional[str]) -> bool: del windows[window_id] # Open a document in a window - def new_window(filename: Optional[str] = None) -> bool: - document = Document(filename) + def new_window(input_filename: Optional[str] = None) -> None: + document = Document(input_filename) window_id = uuid.uuid4().hex - window = MainWindow(global_common, gui_common, window_id) + window = MainWindow(global_common, gui_common, window_id, document) window.delete_window.connect(delete_window) windows[window_id] = window - if filename: - # Validate filename - file_path: str = os.path.abspath(os.path.expanduser(filename)) - try: - open(file_path, "rb") - except FileNotFoundError: - click.echo("File not found") - return False - except PermissionError: - click.echo("Permission denied") - return False - window.document.input_filename = file_path + if input_filename: window.content_widget.doc_selection_widget.document_selected.emit() - return True - # Open a new window if not filename is passed if filename is None: new_window() else: # If filename is passed as an argument, open it - if not new_window(filename): - return True + input_filename: str = os.path.abspath(os.path.expanduser(filename)) + new_window(input_filename) # Open a new window, if all windows are closed def application_activated() -> None: