From b4849995e390d6247fd1863be6ded9003f5c85ee Mon Sep 17 00:00:00 2001
From: deeplow <deeplower@protonmail.com>
Date: Mon, 21 Nov 2022 15:02:35 +0000
Subject: [PATCH] Add CLI support for archiving original / unsafe PDFs

---
 dangerzone/cli.py    | 23 +++++++++++++++++++----
 dangerzone/errors.py |  2 +-
 dangerzone/logic.py  |  7 +++++--
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/dangerzone/cli.py b/dangerzone/cli.py
index 4596836..6cd2300 100644
--- a/dangerzone/cli.py
+++ b/dangerzone/cli.py
@@ -6,7 +6,7 @@ import click
 from colorama import Back, Fore, Style
 
 from . import args, container, errors
-from .document import SAFE_EXTENSION
+from .document import ARCHIVE_SUBDIR, SAFE_EXTENSION
 from .logic import DangerzoneCore
 from .util import get_version
 
@@ -25,6 +25,12 @@ def print_header(s: str) -> None:
     help=f"Default is filename ending with {SAFE_EXTENSION}",
 )
 @click.option("--ocr-lang", help="Language to OCR, defaults to none")
+@click.option(
+    "--archive",
+    "archive",
+    flag_value=True,
+    help=f"Archives the unsafe version in a subdirectory named '{ARCHIVE_SUBDIR}'",
+)
 @click.argument(
     "filenames",
     required=True,
@@ -34,20 +40,23 @@ def print_header(s: str) -> None:
 )
 @errors.handle_document_errors
 def cli_main(
-    output_filename: Optional[str], ocr_lang: Optional[str], filenames: List[str]
+    output_filename: Optional[str],
+    ocr_lang: Optional[str],
+    filenames: List[str],
+    archive: bool,
 ) -> None:
     setup_logging()
     dangerzone = DangerzoneCore()
 
     display_banner()
     if len(filenames) == 1 and output_filename:
-        dangerzone.add_document_from_filename(filenames[0], output_filename)
+        dangerzone.add_document_from_filename(filenames[0], output_filename, archive)
     elif len(filenames) > 1 and output_filename:
         click.echo("--output-filename can only be used with one input file.")
         exit(1)
     else:
         for filename in filenames:
-            dangerzone.add_document_from_filename(filename)
+            dangerzone.add_document_from_filename(filename, archive=archive)
 
     # Validate OCR language
     if ocr_lang:
@@ -76,6 +85,12 @@ def cli_main(
         print_header("Safe PDF(s) created successfully")
         for document in documents_safe:
             click.echo(document.output_filename)
+
+        if archive:
+            print_header(
+                f"Unsafe (original) documents moved to '{ARCHIVE_SUBDIR}' subdirectory"
+            )
+
     if documents_failed != []:
         print_header("Failed to convert document(s)")
         for document in documents_failed:
diff --git a/dangerzone/errors.py b/dangerzone/errors.py
index f3b50e3..5e38015 100644
--- a/dangerzone/errors.py
+++ b/dangerzone/errors.py
@@ -97,7 +97,7 @@ def handle_document_errors(func: F) -> F:
         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"
+                msg = "An exception occured while validating a document"
                 log.exception(msg)
             click.echo(str(e))
             exit(1)
diff --git a/dangerzone/logic.py b/dangerzone/logic.py
index ca31fea..1fbe2cf 100644
--- a/dangerzone/logic.py
+++ b/dangerzone/logic.py
@@ -42,9 +42,12 @@ class DangerzoneCore(object):
         self.documents: List[Document] = []
 
     def add_document_from_filename(
-        self, input_filename: str, output_filename: Optional[str] = None
+        self,
+        input_filename: str,
+        output_filename: Optional[str] = None,
+        archive: bool = False,
     ) -> None:
-        doc = Document(input_filename, output_filename)
+        doc = Document(input_filename, output_filename, archive=archive)
         self.add_document(doc)
 
     def add_document(self, doc: Document) -> None: