mirror of
https://github.com/freedomofpress/dangerzone.git
synced 2025-04-28 18:02:38 +02:00
Remove separate dangerzone-container entry point, make CLI work with it, and refactor container code to be more DRY
This commit is contained in:
Micah Lee
parent
4a2c92e911
commit
c9c01f6e79
6 changed files with 148 additions and 290 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
if success:
|
||||
print_header("Safe PDF created successfully")
|
||||
click.echo(common.output_filename)
|
||||
else:
|
||||
print_header("Failed to convert document")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
input_filename = guest_input_filename
|
||||
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,6 +171,7 @@ 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",
|
||||
|
|
@ -257,10 +183,13 @@ def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
|
|||
container_name,
|
||||
"document-to-pixels",
|
||||
]
|
||||
ret = exec_container(args)
|
||||
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",
|
||||
|
|
@ -276,10 +205,95 @@ def convert(vm_info_path, input_filename, output_filename, ocr, ocr_lang):
|
|||
container_name,
|
||||
"pixels-to-pdf",
|
||||
]
|
||||
ret = exec_container(args)
|
||||
ret = exec_container(args, vm_info)
|
||||
if ret != 0:
|
||||
print("pixels-to-pdf failed")
|
||||
else:
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -1 +0,0 @@
|
|||
dangerzone
|
||||
1
setup.py
1
setup.py
|
|
@ -43,7 +43,6 @@ setuptools.setup(
|
|||
entry_points={
|
||||
"console_scripts": [
|
||||
"dangerzone = dangerzone:main",
|
||||
"dangerzone-container = dangerzone:main",
|
||||
"dangerzone-cli = dangerzone:main",
|
||||
]
|
||||
},
|
||||
|
|
|
|||
Loading…
Reference in a new issue