Totally refactor how tasks work and how dangerzone-container works so that there is a single --convert task

dangerzone/common.py

@@ -11,33 +11,6 @@ class Common(object):
     """
 
     def __init__(self):
-        # Temporary directory to store pixel data and safe PDFs
-        cache_dir = appdirs.user_cache_dir("dangerzone")
-        os.makedirs(cache_dir, exist_ok=True)
-        self.pixel_dir = tempfile.TemporaryDirectory(
-            prefix=os.path.join(cache_dir, "pixel-")
-        )
-        self.safe_dir = tempfile.TemporaryDirectory(
-            prefix=os.path.join(cache_dir, "safe-")
-        )
-
-        try:
-            # Make the folders world-readable to ensure that the container has permission
-            # to access it even if it's owned by root or someone else
-            permissions = (
-                stat.S_IRUSR
-                | stat.S_IWUSR
-                | stat.S_IXUSR
-                | stat.S_IRGRP
-                | stat.S_IXGRP
-                | stat.S_IROTH
-                | stat.S_IXOTH
-            )
-            os.chmod(self.pixel_dir.name, permissions)
-            os.chmod(self.safe_dir.name, permissions)
-        except:
-            pass
-
         # Name of input and out files
-        self.document_filename = None
+        self.input_filename = None
-        self.save_filename = None
+        self.output_filename = None

dangerzone/container.py

@@ -6,6 +6,7 @@ import pipes
 import shutil
 import json
 import os
+import uuid
 
 # What is the container runtime for this platform?
 if platform.system() == "Darwin":
@@ -31,10 +32,6 @@ else:
 
 
 def exec(args):
-    args_str = " ".join(pipes.quote(s) for s in args)
-    print("> " + args_str)
-    sys.stdout.flush()
-
     with subprocess.Popen(
         args,
         stdin=None,
@@ -48,12 +45,8 @@ def exec(args):
         return p.returncode
 
 
-def exec_vm(args, vm_info):
+def vm_ssh_args(vm_info):
-    if container_tech == "dangerzone-vm" and vm_info is None:
+    return [
-        print("--vm-info-path required on this platform")
-        return
-
-    args = [
         "/usr/bin/ssh",
         "-q",
         "-i",
@@ -63,20 +56,60 @@ def exec_vm(args, vm_info):
         "-o",
         "StrictHostKeyChecking=no",
         "user@127.0.0.1",
-    ] + args
+    ]
+
+
+def vm_scp_args(vm_info):
+    return [
+        "/usr/bin/scp",
+        "-i",
+        vm_info["client_key_path"],
+        "-P",
+        str(vm_info["tunnel_port"]),
+        "-o",
+        "StrictHostKeyChecking=no",
+    ]
+
+
+def host_exec(args):
+    args_str = " ".join(pipes.quote(s) for s in args)
+    print("> " + args_str)
+
     return exec(args)
 
 
-def mount_vm(path, vm_info):
+def vm_exec(args, vm_info):
-    basename = os.path.basename(path)
+    if container_tech == "dangerzone-vm" and vm_info is None:
-    normalized_path = f"/home/user/mnt/{basename}"
+        print("--vm-info-path required on this platform")
-    exec_vm(["/usr/bin/sshfs", f"hostbox:{path}", normalized_path], vm_info)
+        return
-    return normalized_path
+
+    args_str = " ".join(pipes.quote(s) for s in args)
+    print("VM > " + args_str)
+
+    args = vm_ssh_args(vm_info) + args
+    return exec(args)
 
 
-def unmount_vm(normalized_path, vm_info):
+def vm_mkdir(vm_info):
-    exec_vm(["/usr/bin/fusermount3", normalized_path], vm_info)
+    guest_path = os.path.join("/home/user/", str(uuid.uuid4()))
-    exec_vm(["/bin/rmdir", normalized_path], vm_info)
+    vm_exec(["/bin/mkdir", guest_path], vm_info)
+    return guest_path
+
+
+def vm_rmdir(guest_path, vm_info):
+    vm_exec(["/bin/rm", "-r", guest_path], vm_info)
+
+
+def vm_upload(host_path, guest_path, vm_info):
+    args = vm_scp_args(vm_info) + [host_path, f"user@127.0.0.1:{guest_path}"]
+    print(f"Uploading '{host_path}' to VM at '{guest_path}'")
+    host_exec(args)
+
+
+def vm_download(guest_path, host_path, vm_info):
+    args = vm_scp_args(vm_info) + [f"user@127.0.0.1:{guest_path}", host_path]
+    print(f"Downloading '{guest_path}' from VM to '{host_path}'")
+    host_exec(args)
 
 
 def exec_container(args, vm_info):
@@ -85,11 +118,11 @@ def exec_container(args, vm_info):
         return
 
     if container_tech == "dangerzone-vm":
-        args = ["podman"] + args
+        args = ["/usr/bin/podman"] + args
-        return exec_vm(args, vm_info)
+        return vm_exec(args, vm_info)
-
+    else:
         args = [container_runtime] + args
-    return exec(args)
+        return host_exec(args)
 
 
 def load_vm_info(vm_info_path):
@@ -103,19 +136,19 @@ def load_vm_info(vm_info_path):
 @click.group()
 def container_main():
     """
-    Dangerzone container commands with elevated privileges.
+    Dangerzone container commands. Humans don't need to run this command by themselves.
-    Humans don't need to run this command by themselves.
     """
     pass
 
 
 @container_main.command()
 @click.option("--vm-info-path", default=None)
-@click.option("--container-name", default="docker.io/flmcode/dangerzone")
+def ls(vm_info_path):
-def ls(vm_info_path, container_name):
     """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)
@@ -124,90 +157,86 @@ def ls(vm_info_path, container_name):
 
 @container_main.command()
 @click.option("--vm-info-path", default=None)
-@click.option("--document-filename", required=True)
+@click.option("--input-filename", required=True)
-@click.option("--pixel-dir", required=True)
+@click.option("--output-filename", required=True)
-@click.option("--container-name", default="docker.io/flmcode/dangerzone")
-def documenttopixels(vm_info_path, document_filename, pixel_dir, container_name):
-    """docker run --network none -v [document_filename]:/tmp/input_file -v [pixel_dir]:/dangerzone [container_name] document-to-pixels"""
-
-    vm_info = load_vm_info(vm_info_path)
-
-    document_dir = os.path.dirname(document_filename)
-    if vm_info:
-        container_name = "localhost/dangerzone"
-        normalized_document_dir = mount_vm(document_dir, vm_info)
-        normalized_document_filename = os.path.join(
-            normalized_document_dir, os.path.basename(document_filename)
-        )
-        normalized_pixel_dir = mount_vm(pixel_dir, vm_info)
-    else:
-        normalized_document_dir = document_dir
-        normalized_document_filename = document_filename
-        normalized_pixel_dir = pixel_dir
-
-    args = ["run", "--network", "none"]
-
-    # docker uses --security-opt, podman doesn't
-    if container_tech == "docker":
-        args += ["--security-opt=no-new-privileges:true"]
-
-    args += [
-        "-v",
-        f"{normalized_document_filename}:/tmp/input_file",
-        "-v",
-        f"{normalized_pixel_dir}:/dangerzone",
-        container_name,
-        "document-to-pixels",
-    ]
-    ret = exec_container(args, load_vm_info(vm_info_path))
-
-    if vm_info:
-        unmount_vm(normalized_document_dir, vm_info)
-        unmount_vm(normalized_pixel_dir, vm_info)
-
-    sys.exit(ret)
-
-
-@container_main.command()
-@click.option("--vm-info-path", default=None)
-@click.option("--pixel-dir", required=True)
-@click.option("--safe-dir", required=True)
-@click.option("--container-name", default="docker.io/flmcode/dangerzone")
 @click.option("--ocr", required=True)
 @click.option("--ocr-lang", required=True)
-def pixelstopdf(vm_info_path, pixel_dir, safe_dir, container_name, ocr, ocr_lang):
+def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
-    """docker run --network none -v [pixel_dir]:/dangerzone -v [safe_dir]:/safezone [container_name] -e OCR=[ocr] -e OCR_LANGUAGE=[ocr_lang] pixels-to-pdf"""
+    # 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
+    #
+    # If there's not a VM
+    # - make tmp pixeldir
+    # - make tmp safedir
+    # - run podman documenttopixels
+    # - run podman pixelstopdf
+    # - delete pixeldir, safedir
+
     vm_info = load_vm_info(vm_info_path)
 
     if vm_info:
-        container_name = "localhost/dangerzone"
+        ssh_args_str = " ".join(pipes.quote(s) for s in vm_ssh_args(vm_info))
-        normalized_pixel_dir = mount_vm(pixel_dir, vm_info)
+        print("If you want to SSH to the VM: " + ssh_args_str)
-        normalized_safe_dir = mount_vm(safe_dir, vm_info)
-    else:
-        normalized_pixel_dir = pixel_dir
-        normalized_safe_dir = safe_dir
 
-    ret = exec_container(
+        container_name = "localhost/dangerzone"
-        [
+
+        input_dir = vm_mkdir(vm_info)
+        pixel_dir = vm_mkdir(vm_info)
+        safe_dir = vm_mkdir(vm_info)
+
+        guest_input_filename = os.path.join(input_dir, "input_file")
+        guest_output_filename = os.path.join(safe_dir, "safe-output-compressed.pdf")
+
+        vm_upload(input_filename, guest_input_filename, vm_info)
+
+        args = [
             "run",
             "--network",
             "none",
             "-v",
-            f"{normalized_pixel_dir}:/dangerzone",
+            f"{guest_input_filename}:/tmp/input_file",
             "-v",
-            f"{normalized_safe_dir}:/safezone",
+            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",
-        ],
+            ]
-        vm_info,
+            ret = exec_container(args, vm_info)
-    )
+            if ret != 0:
+                print("pixels-to-pdf failed")
+            else:
+                vm_download(guest_output_filename, output_filename, vm_info)
 
-    if vm_info:
+        vm_rmdir(input_dir, vm_info)
-        unmount_vm(normalized_pixel_dir, vm_info)
+        vm_rmdir(pixel_dir, vm_info)
-        unmount_vm(normalized_safe_dir, vm_info)
+        vm_rmdir(safe_dir, vm_info)
 
-    sys.exit(ret)
+        return ret
+
+    else:
+        print("not implemented yet")

dangerzone/gui/__init__.py

@@ -50,9 +50,8 @@ class ApplicationWrapper(QtCore.QObject):
 
 
 @click.command()
-@click.option("--custom-container")  # Use this container instead of flmcode/dangerzone
 @click.argument("filename", required=False)
-def gui_main(custom_container, filename):
+def gui_main(filename):
     if platform.system() == "Darwin":
         # Required for macOS Big Sur: https://stackoverflow.com/a/64878899
         os.environ["QT_MAC_WANTS_LAYER"] = "1"
@@ -85,14 +84,6 @@ def gui_main(custom_container, filename):
     global_common = GlobalCommon()
     gui_common = GuiCommon(app, global_common)
 
-    if custom_container:
-        success, error_message = global_common.container_exists(custom_container)
-        if not success:
-            click.echo(error_message)
-            return
-
-        global_common.custom_container = custom_container
-
     # Allow Ctrl-C to smoothly quit the program instead of throwing an exception
     signal.signal(signal.SIGINT, signal.SIG_DFL)
 
@@ -130,7 +121,7 @@ def gui_main(custom_container, filename):
     def select_document(filename=None):
         if (
             len(windows) == 1
-            and windows[list(windows.keys())[0]].common.document_filename == None
+            and windows[list(windows.keys())[0]].common.input_filename == None
         ):
             window = windows[list(windows.keys())[0]]
         else:
@@ -150,7 +141,7 @@ def gui_main(custom_container, filename):
             except PermissionError:
                 click.echo("Permission denied")
                 return False
-            window.common.document_filename = filename
+            window.common.input_filename = filename
             window.doc_selection_widget.document_selected.emit()
 
         return True

dangerzone/gui/main_window.py

@@ -5,7 +5,7 @@ import tempfile
 import subprocess
 from PySide2 import QtCore, QtGui, QtWidgets
 
-from .tasks import ConvertToPixels, ConvertToPDF
+from .tasks import Convert
 from ..common import Common
 
 
@@ -51,7 +51,10 @@ class MainWindow(QtWidgets.QMainWindow):
         self.doc_selection_widget.document_selected.connect(self.document_selected)
 
         # Only use the waiting widget if we have a VM
-        if self.global_common.vm and self.global_common.vm.state != self.global_common.vm.STATE_ON:
+        if (
+            self.global_common.vm
+            and self.global_common.vm.state != self.global_common.vm.STATE_ON
+        ):
             self.waiting_widget.show()
             self.doc_selection_widget.hide()
         else:
@@ -182,7 +185,7 @@ class DocSelectionWidget(QtWidgets.QWidget):
         )
         if filename[0] != "":
             filename = filename[0]
-            self.common.document_filename = filename
+            self.common.input_filename = filename
             self.document_selected.emit()
 
 
@@ -330,26 +333,31 @@ class SettingsWidget(QtWidgets.QWidget):
     def document_selected(self):
         # Update the danger doc label
         self.dangerous_doc_label.setText(
-            f"Dangerous: {os.path.basename(self.common.document_filename)}"
+            f"Untrusted: {os.path.basename(self.common.input_filename)}"
         )
 
         # Update the save location
-        save_filename = f"{os.path.splitext(self.common.document_filename)[0]}-safe.pdf"
+        output_filename = f"{os.path.splitext(self.common.input_filename)[0]}-safe.pdf"
-        self.common.save_filename = save_filename
+        self.common.output_filename = output_filename
-        self.save_lineedit.setText(os.path.basename(save_filename))
+        self.save_lineedit.setText(os.path.basename(output_filename))
 
     def save_browse_button_clicked(self):
         filename = QtWidgets.QFileDialog.getSaveFileName(
             self,
             "Save safe PDF as...",
-            self.common.save_filename,
+            self.common.output_filename,
             filter="Documents (*.pdf)",
         )
         if filename[0] != "":
-            self.common.save_filename = filename[0]
+            self.common.output_filename = filename[0]
-            self.save_lineedit.setText(os.path.basename(self.common.save_filename))
+            self.save_lineedit.setText(os.path.basename(self.common.output_filename))
 
     def start_button_clicked(self):
+        if self.common.output_filename is None:
+            # If not saving, then save it to a temp file instead
+            tmp = tempfile.mkstemp(suffix=".pdf", prefix="dangerzone_")
+            self.common.output_filename = tmp[1]
+
         # Update settings
         self.global_common.settings.set(
             "save", self.save_checkbox.checkState() == QtCore.Qt.Checked
@@ -414,31 +422,21 @@ class TasksWidget(QtWidgets.QWidget):
         layout.addWidget(self.details_scrollarea)
         self.setLayout(layout)
 
-        self.tasks = []
-
     def document_selected(self):
         # Update the danger doc label
         self.dangerous_doc_label.setText(
-            f"Dangerous: {os.path.basename(self.common.document_filename)}"
+            f"Dangerous: {os.path.basename(self.common.input_filename)}"
         )
 
     def start(self):
-        self.tasks += [ConvertToPixels, ConvertToPDF]
-        self.next_task()
-
-    def next_task(self):
-        if len(self.tasks) == 0:
-            self.all_done()
-            return
-
         self.task_details.setText("")
 
-        self.current_task = self.tasks.pop(0)(self.global_common, self.common)
+        self.task = Convert(self.global_common, self.common)
-        self.current_task.update_label.connect(self.update_label)
+        self.task.update_label.connect(self.update_label)
-        self.current_task.update_details.connect(self.update_details)
+        self.task.update_details.connect(self.update_details)
-        self.current_task.task_finished.connect(self.next_task)
+        self.task.task_finished.connect(self.all_done)
-        self.current_task.task_failed.connect(self.task_failed)
+        self.task.task_failed.connect(self.task_failed)
-        self.current_task.start()
+        self.task.start()
 
     def update_label(self, s):
         self.task_label.setText(s)
@@ -449,36 +447,18 @@ class TasksWidget(QtWidgets.QWidget):
     def task_failed(self, err):
         self.task_label.setText("Failed :(")
         self.task_details.setWordWrap(True)
-        text = self.task_details.text()
-        self.task_details.setText(
-            f"{text}\n\n--\n\nDirectory with pixel data: {self.common.pixel_dir.name}\n\n{err}"
-        )
 
     def all_done(self):
-        # Save safe PDF
-        source_filename = f"{self.common.safe_dir.name}/safe-output-compressed.pdf"
-        if self.global_common.settings.get("save"):
-            dest_filename = self.common.save_filename
-        else:
-            # If not saving, then save it to a temp file instead
-            tmp = tempfile.mkstemp(suffix=".pdf", prefix="dangerzone_")
-            dest_filename = tmp[1]
-        shutil.move(source_filename, dest_filename)
-
         # In Windows, open Explorer with the safe PDF in focus
         if platform.system() == "Windows":
-            dest_filename_windows = dest_filename.replace("/", "\\")
+            dest_filename_windows = self.common.output_filename.replace("/", "\\")
             subprocess.Popen(
                 f'explorer.exe /select,"{dest_filename_windows}"', shell=True
             )
 
         # Open
         if self.global_common.settings.get("open"):
-            self.gui_common.open_pdf_viewer(dest_filename)
+            self.gui_common.open_pdf_viewer(self.common.output_filename)
-
-        # Clean up
-        self.common.pixel_dir.cleanup()
-        self.common.safe_dir.cleanup()
 
         # Quit
         if platform.system() == "Darwin":

dangerzone/gui/tasks.py

@@ -46,48 +46,15 @@ class TaskBase(QtCore.QThread):
         return p.returncode, output, stderr
 
 
-class ConvertToPixels(TaskBase):
+class Convert(TaskBase):
     def __init__(self, global_common, common):
-        super(ConvertToPixels, self).__init__()
+        super(Convert, self).__init__()
         self.global_common = global_common
         self.common = common
 
     def run(self):
-        self.update_label.emit("Converting document to pixels")
+        self.update_label.emit("Converting document to safe PDF")
-        args = [
-            "documenttopixels",
-            "--document-filename",
-            self.common.document_filename,
-            "--pixel-dir",
-            self.common.pixel_dir.name,
-            "--container-name",
-            self.global_common.get_container_name(),
-        ]
-        returncode, output, _ = self.exec_container(args)
 
-        if returncode != 0:
-            return
-
-        success, error_message = self.global_common.validate_convert_to_pixel_output(
-            self.common, output
-        )
-        if not success:
-            self.task_failed.emit(error_message)
-            return
-
-        self.task_finished.emit()
-
-
-class ConvertToPDF(TaskBase):
-    def __init__(self, global_common, common):
-        super(ConvertToPDF, self).__init__()
-        self.global_common = global_common
-        self.common = common
-
-    def run(self):
-        self.update_label.emit("Converting pixels to safe PDF")
-
-        # Build environment variables list
         if self.global_common.settings.get("ocr"):
             ocr = "1"
         else:
@@ -97,13 +64,11 @@ class ConvertToPDF(TaskBase):
         ]
 
         args = [
-            "pixelstopdf",
+            "convert",
-            "--pixel-dir",
+            "--input-filename",
-            self.common.pixel_dir.name,
+            self.common.input_filename,
-            "--safe-dir",
+            "--output-filename",
-            self.common.safe_dir.name,
+            self.common.output_filename,
-            "--container-name",
-            self.global_common.get_container_name(),
             "--ocr",
             ocr,
             "--ocr-lang",
@@ -114,4 +79,11 @@ class ConvertToPDF(TaskBase):
         if returncode != 0:
             return
 
+        # success, error_message = self.global_common.validate_convert_to_pixel_output(
+        #     self.common, output
+        # )
+        # if not success:
+        #     self.task_failed.emit(error_message)
+        #     return
+
         self.task_finished.emit()

dangerzone/gui/vm.py

@@ -248,10 +248,6 @@ class Vm(QtCore.QObject):
         self.state = self.STATE_FAIL
         self.vm_state_change.emit(self.state)
 
-    def restart(self):
-        self.stop()
-        self.start()
-
     def stop(self):
         # Kill existing processes
         self.kill_sshd()
@@ -262,6 +258,9 @@ class Vm(QtCore.QObject):
             self.hyperkit_p.terminate()
             self.hyperkit_p = None
 
+        # Just to be extra sure
+        self.kill_hyperkit()
+
     def find_open_port(self):
         with socket.socket() as tmpsock:
             while True:
@@ -286,6 +285,18 @@ class Vm(QtCore.QObject):
                 except Exception:
                     pass
 
+    def kill_hyperkit(self):
+        if os.path.exists(self.hyperkit_pid_path):
+            with open(self.hyperkit_pid_path) as f:
+                hyperkit_pid = int(f.read())
+
+            if psutil.pid_exists(hyperkit_pid):
+                try:
+                    proc = psutil.Process(hyperkit_pid)
+                    proc.kill()
+                except Exception:
+                    pass
+
 
 class WaitForSsh(QtCore.QThread):
     connected = QtCore.Signal()