Remove separate dangerzone-container entry point, make CLI work with it, and refactor container code to be more DRY

dangerzone/__init__.py

@@ -5,16 +5,12 @@ if "DANGERZONE_MODE" in os.environ:
     mode = os.environ["DANGERZONE_MODE"]
 else:
     basename = os.path.basename(sys.argv[0])
-    if basename == "dangerzone-container" or basename == "dangerzone-container.exe":
+    if basename == "dangerzone-cli" or basename == "dangerzone-cli.exe":
-        mode = "container"
-    elif basename == "dangerzone-cli" or basename == "dangerzone-cli.exe":
         mode = "cli"
     else:
         mode = "gui"
 
-if mode == "container":
+if mode == "cli":
-    from .container import container_main as main
-elif mode == "cli":
     from .cli import cli_main as main
 else:
     from .gui import gui_main as main

dangerzone/cli.py

@@ -1,10 +1,11 @@
 import os
 import shutil
 import click
-from colorama import Fore, Back, Style
+from colorama import Fore, Style
 
 from .global_common import GlobalCommon
 from .common import Common
+from .container import convert
 
 
 def print_header(s):
@@ -12,36 +13,6 @@ def print_header(s):
     click.echo(Style.BRIGHT + s)
 
 
-def exec_container(global_common, args):
-    output = ""
-
-    with global_common.exec_dangerzone_container(args) as p:
-        for line in p.stdout:
-            output += line.decode()
-
-            # Hack to add colors to the command executing
-            if line.startswith(b"> "):
-                print(
-                    Style.DIM + "> " + Style.NORMAL + Fore.CYAN + line.decode()[2:],
-                    end="",
-                )
-            else:
-                print("  " + line.decode(), end="")
-
-        stderr = p.stderr.read().decode()
-        if len(stderr) > 0:
-            print("")
-            for line in stderr.strip().split("\n"):
-                print("  " + Style.DIM + line)
-
-    if p.returncode != 0:
-        click.echo(f"Return code: {p.returncode}")
-        if p.returncode == 126 or p.returncode == 127:
-            click.echo(f"Authorization failed")
-
-    return p.returncode, output, stderr
-
-
 @click.command()
 @click.option("--output-filename", help="Default is filename ending with -safe.pdf")
 @click.option("--ocr-lang", help="Language to OCR, defaults to none")
@@ -64,7 +35,7 @@ def cli_main(output_filename, ocr_lang, filename):
         click.echo("Invalid filename")
         return
 
-    common.document_filename = os.path.abspath(filename)
+    common.input_filename = os.path.abspath(filename)
 
     # Validate safe PDF output filename
     if output_filename:
@@ -87,7 +58,7 @@ def cli_main(output_filename, ocr_lang, filename):
 
     else:
         common.output_filename = (
-            f"{os.path.splitext(common.document_filename)[0]}-safe.pdf"
+            f"{os.path.splitext(common.input_filename)[0]}-safe.pdf"
         )
         try:
             with open(common.output_filename, "wb") as f:
@@ -117,37 +88,12 @@ def cli_main(output_filename, ocr_lang, filename):
 
     # Convert the document
     print_header("Converting document to safe PDF")
-
+    success = convert(
-    if ocr_lang:
+        global_common, common.input_filename, common.output_filename, ocr_lang
-        ocr = "1"
-    else:
-        ocr = "0"
-        ocr_lang = ""
-
-    returncode, output, _ = exec_container(
-        global_common,
-        [
-            "convert",
-            "--input-filename",
-            common.document_filename,
-            "--output-filename",
-            common.output_filename,
-            "--ocr",
-            ocr,
-            "--ocr-lang",
-            ocr_lang,
-        ],
     )
 
-    if returncode != 0:
+    if success:
-        return
+        print_header("Safe PDF created successfully")
-
+        click.echo(common.output_filename)
-    # success, error_message = global_common.validate_convert_to_pixel_output(
+    else:
-    #     common, output
+        print_header("Failed to convert document")
-    # )
-    # if not success:
-    #     click.echo(error_message)
-    #     return
-
-    print_header("Safe PDF created successfully")
-    click.echo(common.output_filename)

dangerzone/container.py

@@ -1,4 +1,3 @@
-import click
 import platform
 import subprocess
 import sys
@@ -109,7 +108,7 @@ def vm_download(guest_path, host_path, vm_info):
 
 def exec_container(args, vm_info=None):
     if container_tech == "dangerzone-vm" and vm_info is None:
-        print("--vm-info-path required on this platform")
+        print("Invalid VM info")
         return
 
     if container_tech == "dangerzone-vm":
@@ -133,111 +132,37 @@ def load_vm_info(vm_info_path):
         return json.loads(f.read())
 
 
-@click.group()
+def convert(global_common, input_filename, output_filename, ocr_lang):
-def container_main():
+    success = False
-    """
-    Dangerzone container commands. Humans don't need to run this command by themselves.
-    """
-    pass
 
-
-@container_main.command()
-@click.option("--vm-info-path", default=None)
-def ls(vm_info_path):
-    """docker image ls [container_name]"""
-    if vm_info_path:
-        container_name = "localhost/dangerzone"
-    else:
-        container_name = "dangerzone"
-
-    sys.exit(
-        exec_container(["image", "ls", container_name]), load_vm_info(vm_info_path)
-    )
-
-
-@container_main.command()
-@click.option("--vm-info-path", default=None)
-@click.option("--input-filename", required=True)
-@click.option("--output-filename", required=True)
-@click.option("--ocr", required=True)
-@click.option("--ocr-lang", required=True)
-def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
     container_name = "dangerzone.rocks/dangerzone"
-    vm_info = load_vm_info(vm_info_path)
+    if ocr_lang:
+        ocr = "1"
+    else:
+        ocr = "0"
 
+    if global_common.vm:
+        vm_info = load_vm_info(global_common.vm.vm_info_path)
+    else:
+        vm_info = None
+
+    # If we're using the VM, create temp dirs in the guest and upload the input file
+    # Otherwise, create temp dirs
     if vm_info:
-        # If there's a VM:
-        # - make inputdir on VM
-        # - make pixeldir on VM
-        # - make safedir on VM
-        # - scp input file to inputdir
-        # - run podman documenttopixels
-        # - run podman pixelstopdf
-        # - scp output file to host
-        # - delete inputdir, pixeldir, safedir
         ssh_args_str = " ".join(pipes.quote(s) for s in vm_ssh_args(vm_info))
         print("If you want to SSH to the VM: " + ssh_args_str)
 
-        input_dir = vm_mkdir(vm_info)
+        guest_tmpdir = vm_mkdir(vm_info)
-        pixel_dir = vm_mkdir(vm_info)
+        input_dir = os.path.join(guest_tmpdir, "input")
-        safe_dir = vm_mkdir(vm_info)
+        pixel_dir = os.path.join(guest_tmpdir, "pixel")
+        safe_dir = os.path.join(guest_tmpdir, "safe")
 
         guest_input_filename = os.path.join(input_dir, "input_file")
-        guest_output_filename = os.path.join(safe_dir, "safe-output-compressed.pdf")
+        container_output_filename = os.path.join(safe_dir, "safe-output-compressed.pdf")
 
         vm_upload(input_filename, guest_input_filename, vm_info)
-
+        input_filename = guest_input_filename
-        args = [
-            "run",
-            "--network",
-            "none",
-            "-v",
-            f"{guest_input_filename}:/tmp/input_file",
-            "-v",
-            f"{pixel_dir}:/dangerzone",
-            container_name,
-            "document-to-pixels",
-        ]
-        ret = exec_container(args, vm_info)
-        if ret != 0:
-            print("documents-to-pixels failed")
-        else:
-            args = [
-                "run",
-                "--network",
-                "none",
-                "-v",
-                f"{pixel_dir}:/dangerzone",
-                "-v",
-                f"{safe_dir}:/safezone",
-                "-e",
-                f"OCR={ocr}",
-                "-e",
-                f"OCR_LANGUAGE={ocr_lang}",
-                container_name,
-                "pixels-to-pdf",
-            ]
-            ret = exec_container(args, vm_info)
-            if ret != 0:
-                print("pixels-to-pdf failed")
-            else:
-                vm_download(guest_output_filename, output_filename, vm_info)
-
-        vm_rmdir(input_dir, vm_info)
-        vm_rmdir(pixel_dir, vm_info)
-        vm_rmdir(safe_dir, vm_info)
-
-        sys.exit(ret)
-
     else:
-        # If there's not a VM
-        # - make tmp pixeldir
-        # - make tmp safedir
-        # - run podman documenttopixels
-        # - run podman pixelstopdf
-        # - copy safe PDF to output filename
-        # - delete pixeldir, safedir
-
         tmpdir = tempfile.TemporaryDirectory()
         pixel_dir = os.path.join(tmpdir.name, "pixels")
         safe_dir = os.path.join(tmpdir.name, "safe")
@@ -246,40 +171,129 @@ def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
 
         container_output_filename = os.path.join(safe_dir, "safe-output-compressed.pdf")
 
+    # Convert document to pixels
+    args = [
+        "run",
+        "--network",
+        "none",
+        "-v",
+        f"{input_filename}:/tmp/input_file",
+        "-v",
+        f"{pixel_dir}:/dangerzone",
+        container_name,
+        "document-to-pixels",
+    ]
+    ret = exec_container(args, vm_info)
+    if ret != 0:
+        print("documents-to-pixels failed")
+    else:
+        # TODO: validate convert to pixels output
+
+        # Convert pixels to safe PDF
         args = [
             "run",
             "--network",
             "none",
             "-v",
-            f"{input_filename}:/tmp/input_file",
-            "-v",
             f"{pixel_dir}:/dangerzone",
+            "-v",
+            f"{safe_dir}:/safezone",
+            "-e",
+            f"OCR={ocr}",
+            "-e",
+            f"OCR_LANGUAGE={ocr_lang}",
             container_name,
-            "document-to-pixels",
+            "pixels-to-pdf",
         ]
-        ret = exec_container(args)
+        ret = exec_container(args, vm_info)
         if ret != 0:
-            print("documents-to-pixels failed")
+            print("pixels-to-pdf failed")
         else:
-            args = [
+            # Move the final file to the right place
-                "run",
+            if vm_info:
-                "--network",
+                vm_download(container_output_filename, output_filename, vm_info)
-                "none",
-                "-v",
-                f"{pixel_dir}:/dangerzone",
-                "-v",
-                f"{safe_dir}:/safezone",
-                "-e",
-                f"OCR={ocr}",
-                "-e",
-                f"OCR_LANGUAGE={ocr_lang}",
-                container_name,
-                "pixels-to-pdf",
-            ]
-            ret = exec_container(args)
-            if ret != 0:
-                print("pixels-to-pdf failed")
             else:
                 os.rename(container_output_filename, output_filename)
 
+            # We did it
+            success = True
+
+    # Clean up
+    if vm_info:
+        vm_rmdir(guest_tmpdir, vm_info)
+    else:
         shutil.rmtree(tmpdir.name)
+
+    return success
+
+
+# From global_common:
+
+# def validate_convert_to_pixel_output(self, common, output):
+#     """
+#     Take the output from the convert to pixels tasks and validate it. Returns
+#     a tuple like: (success (boolean), error_message (str))
+#     """
+#     max_image_width = 10000
+#     max_image_height = 10000
+
+#     # Did we hit an error?
+#     for line in output.split("\n"):
+#         if (
+#             "failed:" in line
+#             or "The document format is not supported" in line
+#             or "Error" in line
+#         ):
+#             return False, output
+
+#     # How many pages was that?
+#     num_pages = None
+#     for line in output.split("\n"):
+#         if line.startswith("Document has "):
+#             num_pages = line.split(" ")[2]
+#             break
+#     if not num_pages or not num_pages.isdigit() or int(num_pages) <= 0:
+#         return False, "Invalid number of pages returned"
+#     num_pages = int(num_pages)
+
+#     # Make sure we have the files we expect
+#     expected_filenames = []
+#     for i in range(1, num_pages + 1):
+#         expected_filenames += [
+#             f"page-{i}.rgb",
+#             f"page-{i}.width",
+#             f"page-{i}.height",
+#         ]
+#     expected_filenames.sort()
+#     actual_filenames = os.listdir(common.pixel_dir.name)
+#     actual_filenames.sort()
+
+#     if expected_filenames != actual_filenames:
+#         return (
+#             False,
+#             f"We expected these files:\n{expected_filenames}\n\nBut we got these files:\n{actual_filenames}",
+#         )
+
+#     # Make sure the files are the correct sizes
+#     for i in range(1, num_pages + 1):
+#         with open(f"{common.pixel_dir.name}/page-{i}.width") as f:
+#             w_str = f.read().strip()
+#         with open(f"{common.pixel_dir.name}/page-{i}.height") as f:
+#             h_str = f.read().strip()
+#         w = int(w_str)
+#         h = int(h_str)
+#         if (
+#             not w_str.isdigit()
+#             or not h_str.isdigit()
+#             or w <= 0
+#             or w > max_image_width
+#             or h <= 0
+#             or h > max_image_height
+#         ):
+#             return False, f"Page {i} has invalid geometry"
+
+#         # Make sure the RGB file is the correct size
+#         if os.path.getsize(f"{common.pixel_dir.name}/page-{i}.rgb") != w * h * 3:
+#             return False, f"Page {i} has an invalid RGB file size"
+
+#     return True, True

dangerzone/global_common.py

@@ -11,6 +11,7 @@ import colorama
 from colorama import Fore, Back, Style
 
 from .settings import Settings
+from .container import convert
 
 
 class GlobalCommon(object):
@@ -37,9 +38,6 @@ class GlobalCommon(object):
         # In case we have a custom container
         self.custom_container = None
 
-        # dangerzone-container path
-        self.dz_container_path = self.get_dangerzone_container_path()
-
         # VM object, if available
         self.vm = None
 
@@ -416,34 +414,9 @@ class GlobalCommon(object):
         resource_path = os.path.join(prefix, filename)
         return resource_path
 
-    def get_dangerzone_container_path(self):
+    def exec_dangerzone_container(self, input_filename, output_filename, ocr_lang):
-        if getattr(sys, "dangerzone_dev", False):
+        convert(self, input_filename, output_filename, ocr_lang)
-            # Look for resources directory relative to python file
-            path = os.path.join(
-                os.path.dirname(
-                    os.path.dirname(
-                        os.path.abspath(inspect.getfile(inspect.currentframe()))
-                    )
-                ),
-                "dev_scripts",
-                "dangerzone-container",
-            )
-            if platform.system() == "Windows":
-                path = f"{path}.bat"
-            return path
-        else:
-            if platform.system() == "Darwin":
-                return os.path.join(
-                    os.path.dirname(sys.executable), "dangerzone-container"
-                )
-            elif platform.system() == "Windows":
-                return os.path.join(
-                    os.path.dirname(sys.executable), "dangerzone-container.exe"
-                )
-            else:
-                return "/usr/bin/dangerzone-container"
 
-    def exec_dangerzone_container(self, args):
         args = [self.dz_container_path] + args
         if self.vm:
             args += ["--vm-info-path", self.vm.vm_info_path]
@@ -494,75 +467,6 @@ class GlobalCommon(object):
 
         return True, True
 
-    def validate_convert_to_pixel_output(self, common, output):
-        """
-        Take the output from the convert to pixels tasks and validate it. Returns
-        a tuple like: (success (boolean), error_message (str))
-        """
-        max_image_width = 10000
-        max_image_height = 10000
-
-        # Did we hit an error?
-        for line in output.split("\n"):
-            if (
-                "failed:" in line
-                or "The document format is not supported" in line
-                or "Error" in line
-            ):
-                return False, output
-
-        # How many pages was that?
-        num_pages = None
-        for line in output.split("\n"):
-            if line.startswith("Document has "):
-                num_pages = line.split(" ")[2]
-                break
-        if not num_pages or not num_pages.isdigit() or int(num_pages) <= 0:
-            return False, "Invalid number of pages returned"
-        num_pages = int(num_pages)
-
-        # Make sure we have the files we expect
-        expected_filenames = []
-        for i in range(1, num_pages + 1):
-            expected_filenames += [
-                f"page-{i}.rgb",
-                f"page-{i}.width",
-                f"page-{i}.height",
-            ]
-        expected_filenames.sort()
-        actual_filenames = os.listdir(common.pixel_dir.name)
-        actual_filenames.sort()
-
-        if expected_filenames != actual_filenames:
-            return (
-                False,
-                f"We expected these files:\n{expected_filenames}\n\nBut we got these files:\n{actual_filenames}",
-            )
-
-        # Make sure the files are the correct sizes
-        for i in range(1, num_pages + 1):
-            with open(f"{common.pixel_dir.name}/page-{i}.width") as f:
-                w_str = f.read().strip()
-            with open(f"{common.pixel_dir.name}/page-{i}.height") as f:
-                h_str = f.read().strip()
-            w = int(w_str)
-            h = int(h_str)
-            if (
-                not w_str.isdigit()
-                or not h_str.isdigit()
-                or w <= 0
-                or w > max_image_width
-                or h <= 0
-                or h > max_image_height
-            ):
-                return False, f"Page {i} has invalid geometry"
-
-            # Make sure the RGB file is the correct size
-            if os.path.getsize(f"{common.pixel_dir.name}/page-{i}.rgb") != w * h * 3:
-                return False, f"Page {i} has an invalid RGB file size"
-
-        return True, True
-
     def install_container(self):
         """
         Make sure the podman container is installed. Linux only.

dev_scripts/dangerzone-container

@@ -1 +0,0 @@
-dangerzone

setup.py

@@ -43,7 +43,6 @@ setuptools.setup(
     entry_points={
         "console_scripts": [
             "dangerzone = dangerzone:main",
-            "dangerzone-container = dangerzone:main",
             "dangerzone-cli = dangerzone:main",
         ]
     },