diff --git a/dangerzone/__init__.py b/dangerzone/__init__.py index baf50a8..69fb9db 100644 --- a/dangerzone/__init__.py +++ b/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": - mode = "container" - elif basename == "dangerzone-cli" or basename == "dangerzone-cli.exe": + if basename == "dangerzone-cli" or basename == "dangerzone-cli.exe": mode = "cli" else: mode = "gui" -if mode == "container": - from .container import container_main as main -elif mode == "cli": +if mode == "cli": from .cli import cli_main as main else: from .gui import gui_main as main diff --git a/dangerzone/cli.py b/dangerzone/cli.py index 99b2e04..6633bd6 100644 --- a/dangerzone/cli.py +++ b/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") - - if 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, - ], + success = convert( + global_common, common.input_filename, common.output_filename, ocr_lang ) - if returncode != 0: - return - - # success, error_message = global_common.validate_convert_to_pixel_output( - # common, output - # ) - # if not success: - # click.echo(error_message) - # return - - print_header("Safe PDF created successfully") - click.echo(common.output_filename) + if success: + print_header("Safe PDF created successfully") + click.echo(common.output_filename) + else: + print_header("Failed to convert document") diff --git a/dangerzone/container.py b/dangerzone/container.py index a712be6..2f25dac 100644 --- a/dangerzone/container.py +++ b/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 container_main(): - """ - Dangerzone container commands. Humans don't need to run this command by themselves. - """ - pass +def convert(global_common, input_filename, output_filename, ocr_lang): + success = False - -@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) - pixel_dir = vm_mkdir(vm_info) - safe_dir = vm_mkdir(vm_info) + guest_tmpdir = vm_mkdir(vm_info) + input_dir = os.path.join(guest_tmpdir, "input") + 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) - - 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) - + input_filename = guest_input_filename 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 = [ - "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) - if ret != 0: - print("pixels-to-pdf failed") + # Move the final file to the right place + if vm_info: + vm_download(container_output_filename, output_filename, vm_info) 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 diff --git a/dangerzone/global_common.py b/dangerzone/global_common.py index ba413b3..f7e8a69 100644 --- a/dangerzone/global_common.py +++ b/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): - if getattr(sys, "dangerzone_dev", False): - # 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, input_filename, output_filename, ocr_lang): + convert(self, input_filename, output_filename, ocr_lang) - 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. diff --git a/dev_scripts/dangerzone-container b/dev_scripts/dangerzone-container deleted file mode 120000 index 2fe47df..0000000 --- a/dev_scripts/dangerzone-container +++ /dev/null @@ -1 +0,0 @@ -dangerzone \ No newline at end of file diff --git a/setup.py b/setup.py index 67b8ecc..f6f736e 100644 --- a/setup.py +++ b/setup.py @@ -43,7 +43,6 @@ setuptools.setup( entry_points={ "console_scripts": [ "dangerzone = dangerzone:main", - "dangerzone-container = dangerzone:main", "dangerzone-cli = dangerzone:main", ] },