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

This commit is contained in:
Micah Lee 2021-07-02 13:32:23 -07:00
parent fe63689320
commit 488dca4a71
No known key found for this signature in database
GPG key ID: 403C2657CD994F73
6 changed files with 192 additions and 236 deletions

View file

@ -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.save_filename = None
self.input_filename = None
self.output_filename = None

View file

@ -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):
if container_tech == "dangerzone-vm" and vm_info is None:
print("--vm-info-path required on this platform")
return
args = [
def vm_ssh_args(vm_info):
return [
"/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):
basename = os.path.basename(path)
normalized_path = f"/home/user/mnt/{basename}"
exec_vm(["/usr/bin/sshfs", f"hostbox:{path}", normalized_path], vm_info)
return normalized_path
def vm_exec(args, vm_info):
if container_tech == "dangerzone-vm" and vm_info is None:
print("--vm-info-path required on this platform")
return
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):
exec_vm(["/usr/bin/fusermount3", normalized_path], vm_info)
exec_vm(["/bin/rmdir", normalized_path], vm_info)
def vm_mkdir(vm_info):
guest_path = os.path.join("/home/user/", str(uuid.uuid4()))
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
return exec_vm(args, vm_info)
args = ["/usr/bin/podman"] + args
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.
Humans don't need to run this command by themselves.
Dangerzone container commands. 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, container_name):
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)
@ -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("--pixel-dir", 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("--input-filename", required=True)
@click.option("--output-filename", required=True)
@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):
"""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"""
def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
# 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"
normalized_pixel_dir = mount_vm(pixel_dir, vm_info)
normalized_safe_dir = mount_vm(safe_dir, vm_info)
else:
normalized_pixel_dir = pixel_dir
normalized_safe_dir = safe_dir
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)
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:
unmount_vm(normalized_pixel_dir, vm_info)
unmount_vm(normalized_safe_dir, vm_info)
vm_rmdir(input_dir, vm_info)
vm_rmdir(pixel_dir, vm_info)
vm_rmdir(safe_dir, vm_info)
sys.exit(ret)
return ret
else:
print("not implemented yet")

View file

@ -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

View file

@ -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"
self.common.save_filename = save_filename
self.save_lineedit.setText(os.path.basename(save_filename))
output_filename = f"{os.path.splitext(self.common.input_filename)[0]}-safe.pdf"
self.common.output_filename = output_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.save_lineedit.setText(os.path.basename(self.common.save_filename))
self.common.output_filename = filename[0]
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.current_task.update_label.connect(self.update_label)
self.current_task.update_details.connect(self.update_details)
self.current_task.task_finished.connect(self.next_task)
self.current_task.task_failed.connect(self.task_failed)
self.current_task.start()
self.task = Convert(self.global_common, self.common)
self.task.update_label.connect(self.update_label)
self.task.update_details.connect(self.update_details)
self.task.task_finished.connect(self.all_done)
self.task.task_failed.connect(self.task_failed)
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)
# Clean up
self.common.pixel_dir.cleanup()
self.common.safe_dir.cleanup()
self.gui_common.open_pdf_viewer(self.common.output_filename)
# Quit
if platform.system() == "Darwin":

View file

@ -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")
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)
self.update_label.emit("Converting document to safe PDF")
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",
"--pixel-dir",
self.common.pixel_dir.name,
"--safe-dir",
self.common.safe_dir.name,
"--container-name",
self.global_common.get_container_name(),
"convert",
"--input-filename",
self.common.input_filename,
"--output-filename",
self.common.output_filename,
"--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()

View file

@ -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()