mirror of
https://github.com/freedomofpress/dangerzone.git
synced 2025-05-01 11:12:24 +02:00
Wrap dangerzone.py back into a class to keep track of percentage
This commit is contained in:
Guthrie McAfee Armstrong
deeplow
parent
eaa08c9c3d
commit
17939cb70c
1 changed files with 293 additions and 308 deletions
|
|
@ -18,10 +18,7 @@ import os
|
|||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
<<<<<<< HEAD
|
||||
from typing import Dict, Optional
|
||||
=======
|
||||
>>>>>>> d990cfb (refactor dangerzone.py, raise exceptions instead of returning int)
|
||||
|
||||
import magic
|
||||
from PIL import Image
|
||||
|
|
@ -55,13 +52,12 @@ def run_command(
|
|||
raise TimeoutError(timeout_message) from e
|
||||
|
||||
|
||||
def output(self, error: bool, text: str, percentage: float) -> None:
|
||||
print(json.dumps({"error": error, "text": text, "percentage": int(percentage)}))
|
||||
sys.stdout.flush()
|
||||
class ConversionJob(object):
|
||||
def __init__(self) -> None:
|
||||
self.percentage: float = 0.0 # TODO Optional[float], but this default value will be overwritten immediately
|
||||
|
||||
|
||||
def document_to_pixels() -> None:
|
||||
percentage: float = 0.0
|
||||
def document_to_pixels(self) -> None:
|
||||
self.percentage: float = 0.0
|
||||
|
||||
conversions: Dict[str, Dict[str, Optional[str]]] = {
|
||||
# .pdf
|
||||
|
|
@ -145,7 +141,7 @@ def document_to_pixels() -> None:
|
|||
if conversion["type"] is None:
|
||||
pdf_filename = "/tmp/input_file"
|
||||
elif conversion["type"] == "libreoffice":
|
||||
output(False, "Converting to PDF using LibreOffice", percentage)
|
||||
self.update_progress("Converting to PDF using LibreOffice")
|
||||
args = [
|
||||
"libreoffice",
|
||||
"--headless",
|
||||
|
|
@ -162,7 +158,7 @@ def document_to_pixels() -> None:
|
|||
)
|
||||
pdf_filename = "/tmp/input_file.pdf"
|
||||
elif conversion["type"] == "convert":
|
||||
output(False, "Converting to PDF using GraphicsMagick", percentage)
|
||||
self.update_progress("Converting to PDF using GraphicsMagick")
|
||||
args = [
|
||||
"gm",
|
||||
"convert",
|
||||
|
|
@ -179,14 +175,10 @@ def document_to_pixels() -> None:
|
|||
raise ValueError(
|
||||
f"Invalid conversion type {conversion['type']} for MIME type {mime_type}"
|
||||
)
|
||||
percentage += 3
|
||||
self.percentage += 3
|
||||
|
||||
# Separate PDF into pages
|
||||
output(
|
||||
False,
|
||||
"Separating document into pages",
|
||||
percentage,
|
||||
)
|
||||
self.update_progress("Separating document into pages"),
|
||||
args = ["pdftk", pdf_filename, "burst", "output", "/tmp/page-%d.pdf"]
|
||||
run_command(
|
||||
args,
|
||||
|
|
@ -196,7 +188,7 @@ def document_to_pixels() -> None:
|
|||
|
||||
page_filenames = glob.glob("/tmp/page-*.pdf")
|
||||
|
||||
percentage += 2
|
||||
self.percentage += 2
|
||||
|
||||
# Convert to RGB pixel data
|
||||
percentage_per_page = 45.0 / len(page_filenames)
|
||||
|
|
@ -208,10 +200,8 @@ def document_to_pixels() -> None:
|
|||
height_filename = f"/tmp/page-{page}.height"
|
||||
filename_base = f"/tmp/page-{page}"
|
||||
|
||||
output(
|
||||
False,
|
||||
f"Converting page {page}/{len(page_filenames)} to pixels",
|
||||
percentage,
|
||||
self.update_progress(
|
||||
f"Converting page {page}/{len(page_filenames)} to pixels"
|
||||
)
|
||||
|
||||
# Convert to png
|
||||
|
|
@ -245,15 +235,11 @@ def document_to_pixels() -> None:
|
|||
|
||||
# Delete the png
|
||||
os.remove(png_filename)
|
||||
percentage += percentage_per_page
|
||||
self.percentage += percentage_per_page
|
||||
|
||||
# END OF FOR LOOP
|
||||
|
||||
output(
|
||||
False,
|
||||
"Converted document to pixels",
|
||||
percentage,
|
||||
)
|
||||
self.update_progress("Converted document to pixels")
|
||||
|
||||
# Move converted files into /dangerzone
|
||||
for filename in (
|
||||
|
|
@ -263,9 +249,8 @@ def document_to_pixels() -> None:
|
|||
):
|
||||
shutil.move(filename, "/dangerzone")
|
||||
|
||||
|
||||
def pixels_to_pdf() -> None:
|
||||
percentage = 50.0
|
||||
def pixels_to_pdf(self) -> None:
|
||||
self.percentage = 50.0
|
||||
|
||||
num_pages = len(glob.glob("/dangerzone/page-*.rgb"))
|
||||
|
||||
|
|
@ -286,10 +271,8 @@ def pixels_to_pdf() -> None:
|
|||
height = f.read().strip()
|
||||
|
||||
if os.environ.get("OCR") == "1": # OCR the document
|
||||
output(
|
||||
False,
|
||||
f"Converting page {page}/{num_pages} from pixels to searchable PDF",
|
||||
percentage,
|
||||
self.update_progress(
|
||||
f"Converting page {page}/{num_pages} from pixels to searchable PDF"
|
||||
)
|
||||
run_command(
|
||||
[
|
||||
|
|
@ -321,10 +304,8 @@ def pixels_to_pdf() -> None:
|
|||
)
|
||||
|
||||
else: # Don't OCR
|
||||
output(
|
||||
False,
|
||||
f"Converting page {page}/{num_pages} from pixels to PDF",
|
||||
percentage,
|
||||
self.update_progress(
|
||||
f"Converting page {page}/{num_pages} from pixels to PDF"
|
||||
)
|
||||
run_command(
|
||||
[
|
||||
|
|
@ -341,16 +322,12 @@ def pixels_to_pdf() -> None:
|
|||
timeout_message=f"Error converting RGB to PDF, convert timed out after {DEFAULT_TIMEOUT} seconds",
|
||||
)
|
||||
|
||||
percentage += percentage_per_page
|
||||
self.percentage += percentage_per_page
|
||||
|
||||
# END OF FOR LOOP
|
||||
|
||||
# Merge pages into a single PDF
|
||||
output(
|
||||
False,
|
||||
f"Merging {num_pages} pages into a single PDF",
|
||||
percentage,
|
||||
)
|
||||
self.update_progress(f"Merging {num_pages} pages into a single PDF")
|
||||
args = ["pdfunite"]
|
||||
for page in range(1, num_pages + 1):
|
||||
args.append(f"/tmp/page-{page}.pdf")
|
||||
|
|
@ -361,14 +338,10 @@ def pixels_to_pdf() -> None:
|
|||
timeout_message=f"Error merging pages into a single PDF, pdfunite timed out after {DEFAULT_TIMEOUT} seconds",
|
||||
)
|
||||
|
||||
percentage += 2
|
||||
self.percentage += 2
|
||||
|
||||
# Compress
|
||||
output(
|
||||
False,
|
||||
f"Compressing PDF",
|
||||
percentage,
|
||||
)
|
||||
self.update_progress(f"Compressing PDF")
|
||||
compress_timeout = num_pages * 3
|
||||
run_command(
|
||||
["ps2pdf", "/tmp/safe-output.pdf", "/tmp/safe-output-compressed.pdf"],
|
||||
|
|
@ -377,34 +350,46 @@ def pixels_to_pdf() -> None:
|
|||
timeout=compress_timeout,
|
||||
)
|
||||
|
||||
percentage = 100.0
|
||||
output(False, "Safe PDF created", percentage)
|
||||
self.percentage = 100.0
|
||||
self.update_progress("Safe PDF created")
|
||||
|
||||
# Move converted files into /safezone
|
||||
shutil.move("/tmp/safe-output.pdf", "/safezone")
|
||||
shutil.move("/tmp/safe-output-compressed.pdf", "/safezone")
|
||||
|
||||
def update_progress(self, text, *, error: bool = False):
|
||||
print(
|
||||
json.dumps(
|
||||
{"error": error, "text": text, "percentage": int(self.percentage)}
|
||||
)
|
||||
)
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) != 2:
|
||||
print(f"Usage: {sys.argv[0]} [document-to-pixels]|[pixels-to-pdf]")
|
||||
return -1
|
||||
|
||||
job = ConversionJob()
|
||||
|
||||
if sys.argv[1] == "document-to-pixels":
|
||||
try:
|
||||
document_to_pixels()
|
||||
except:
|
||||
job.document_to_pixels()
|
||||
except (RuntimeError, TimeoutError) as e:
|
||||
job.update_progress(str(e), error=True)
|
||||
return 1
|
||||
else:
|
||||
return 0
|
||||
return 0 # Success!
|
||||
|
||||
if sys.argv[1] == "pixels-to-pdf":
|
||||
elif sys.argv[1] == "pixels-to-pdf":
|
||||
try:
|
||||
pixels_to_pdf()
|
||||
except:
|
||||
job.pixels_to_pdf()
|
||||
except (RuntimeError, TimeoutError) as e:
|
||||
job.update_progress(str(e), error=True)
|
||||
return 1
|
||||
else:
|
||||
return 0
|
||||
return 0 # Success!
|
||||
|
||||
return -1
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue