speed up emails by being selective about including images

Author: cccs-rs

Dockerfile

@@ -27,5 +27,6 @@ COPY . .
 ARG version=4.0.0.dev1
 USER root
 RUN sed -i -e "s/\$SERVICE_TAG/$version/g" service_manifest.yml
+RUN unoconv --listener &
 
 USER assemblyline

document_preview/document_preview.py

@@ -9,6 +9,7 @@
 
 from assemblyline_v4_service.common.base import ServiceBase
 from assemblyline_v4_service.common.result import Result, ResultImageSection
+from assemblyline_v4_service.common.request import ServiceRequest as Request
 
 from document_preview.helper.emlrender import processEml as eml2image
 from PIL import Image
@@ -21,7 +22,6 @@ def __init__(self, config=None):
         super(DocumentPreview, self).__init__(config)
 
     def start(self):
-        subprocess.Popen(["unoconv", "--listener"])
         self.log.debug("Document preview service started")
 
     def stop(self):
@@ -59,29 +59,32 @@ def pdf_to_images(self, file):
             page.save(self.working_directory + "/output_" + str(i) + ".jpeg")
             i += 1
 
-    def render_documents(self, file_type, file, file_contents, max_pages=1):
+    def render_documents(self, request: Request, max_pages=1):
         # Word/Excel/Powerpoint
-        if any(file_type == f'document/office/{ms_product}' for ms_product in ['word', 'excel', 'powerpoint']):
-            orientation = "landscape" if file_type.endswith('excel') else "portrait"
-            converted = self.office_conversion(file, orientation, max_pages)
+        if any(request.file_type == f'document/office/{ms_product}' for ms_product in ['word', 'excel', 'powerpoint']):
+            orientation = "landscape" if any(request.file_type.endswith(type)
+                                             for type in ['excel', 'powerpoint']) else "portrait"
+            converted = self.office_conversion(request.file_path, orientation, max_pages)
             if converted[0]:
                 self.pdf_to_images(self.working_directory + "/" + converted[1])
         # PDF
-        elif file_type == 'document/pdf':
-            self.pdf_to_images(file)
+        elif request.file_type == 'document/pdf':
+            self.pdf_to_images(request.file_path)
         # EML/MSG
-        elif file_type.endswith('email'):
+        elif request.file_type.endswith('email'):
+            file_contents = request.file_contents
             # Convert MSG to EML where applicable
-            if file_type == 'document/office/email':
+            if request.file_type == 'document/office/email':
                 with tempfile.NamedTemporaryFile() as tmp:
-                    subprocess.run(['msgconvert', '-outfile', tmp.name, file])
+                    subprocess.run(['msgconvert', '-outfile', tmp.name, request.file_path])
                     tmp.seek(0)
                     file_contents = tmp.read()
 
             # Render EML as PNG
             # If we have internet access, we'll attempt to load external images
             output_image = eml2image(file_contents, self.working_directory, self.log,
-                                     load_images=self.service_attributes.docker_config.allow_internet_access)
+                                     load_ext_images=self.service_attributes.docker_config.allow_internet_access,
+                                     load_images=request.get_param('load_email_images'))
             img = Image.open(output_image)
             img_dim = img.size
             if img_dim[1] > WEBP_MAX_SIZE:
@@ -93,14 +96,14 @@ def render_documents(self, file_type, file, file_contents, max_pages=1):
                         height = img_dim[1] - pos_y
                     box = (0, pos_y, img_dim[0], pos_y + height)
                     slice = img.crop(box)
-                    slice.save(f"{output_image}_{index}", "PNG")
+                    slice.save(os.path.join(self.working_directory, f"output_{index}.png"), "PNG")
                     index += 1
                     pos_y = index * WEBP_MAX_SIZE
 
                 os.remove(output_image)
 
-        elif file_type.endswith('emf'):
-            self.libreoffice_conversion(file, convert_to="png")
+        elif request.file_type.endswith('emf'):
+            self.libreoffice_conversion(request.file_path, convert_to="png")
 
     def execute(self, request):
         start = time()
@@ -109,7 +112,7 @@ def execute(self, request):
         # Attempt to render documents given and dump them to the working directory
         max_pages = request.get_param('max_pages_rendered')
         try:
-            self.render_documents(request.file_type, request.file_path, request.file_contents, max_pages)
+            self.render_documents(request, max_pages)
         except Exception as e:
             # Unable to complete analysis after unexpected error, give up
             self.log.error(e)
@@ -119,9 +122,10 @@ def execute(self, request):
         if any("output" in s for s in os.listdir(self.working_directory)):
             previews = [s for s in os.listdir(self.working_directory) if "output" in s]
             image_section = ResultImageSection(request,  "Successfully extracted the preview.")
+            heur_id = 1 if request.deep_scan or request.get_param('run_ocr') else None
             [image_section.add_image(f"{self.working_directory}/{preview}",
                                      name=f"page_{str(i).zfill(3)}.jpeg", description=f"Here's the preview for page {i}",
-                                     ocr_heuristic_id=1)
+                                     ocr_heuristic_id=heur_id)
              for i, preview in enumerate(natsorted(previews))]
 
             result.add_section(image_section)

document_preview/helper/emlrender.py

@@ -15,9 +15,9 @@
 import email
 import email.header
 import quopri
-import hashlib
 import base64
 import regex
+from tempfile import NamedTemporaryFile
 
 try:
     import imgkit
@@ -58,7 +58,7 @@ def appendImages(images):
     return new_im
 
 
-def processEml(data, dumpDir, logger, load_images=False):
+def processEml(data, output_dir, logger, load_ext_images=False, load_images=False):
     '''
     Process the email (bytes), extract MIME parts and useful headers.
     Generate a PNG picture of the mail
@@ -104,7 +104,7 @@ def processEml(data, dumpDir, logger, load_images=False):
     idField = idField.replace('<', '&lt;').replace('>', '&gt;')
 
     imgkitOptions = {'load-error-handling': 'skip'}
-    if not load_images:
+    if not load_ext_images:
         imgkitOptions.update({'no-images': None})
     # imgkitOptions.update({ 'quiet': None })
     imagesList = []
@@ -121,15 +121,13 @@ def processEml(data, dumpDir, logger, load_images=False):
     </table>
     <hr></p>
     ''' % (dateField, fromField, toField, subjectField, idField)
-    m = hashlib.md5()
-    m.update(headers.encode('utf-8'))
-    imagePath = f'output_{m.hexdigest()}.png'
     try:
-        imgkit.from_string(headers, dumpDir + '/' + imagePath, options=imgkitOptions)
-        logger.info('Created headers %s' % imagePath)
-        imagesList.append(dumpDir + '/' + imagePath)
-    except:
-        logger.warning('Creation of headers failed')
+        header_path = NamedTemporaryFile(suffix=".png").name
+        imgkit.from_string(headers, header_path, options=imgkitOptions)
+        logger.info('Created headers %s' % header_path)
+        imagesList.append(header_path)
+    except Exception as e:
+        logger.warning(f'Creation of headers failed: {e}')
 
     #
     # Main loop - process the MIME parts
@@ -157,38 +155,33 @@ def processEml(data, dumpDir, logger, load_images=False):
             # for char in dirtyChars:
             #     payload = payload.replace(char, '')
 
-            # Generate MD5 hash of the payload
-            m = hashlib.md5()
-            m.update(payload.encode('utf-8'))
-            imagePath = f'output_{m.hexdigest()}.png'
             try:
-                imgkit.from_string(payload, dumpDir + '/' + imagePath, options=imgkitOptions)
-                logger.info('Decoded %s' % imagePath)
-                imagesList.append(dumpDir + '/' + imagePath)
+                payload_path = NamedTemporaryFile(suffix=".png").name
+                imgkit.from_string(payload, payload_path, options=imgkitOptions)
+                logger.info('Decoded %s' % payload_path)
+                imagesList.append(payload_path)
             except Exception as e:
                 logger.warning(f'Decoding this MIME part returned error: {e}')
-        elif mimeType in imageTypes:
+
+        elif mimeType in imageTypes and load_images:
             payload = part.get_payload(decode=False)
+            payload_path = NamedTemporaryFile(suffix=".png").name
             imgdata = base64.b64decode(payload)
-            # Generate MD5 hash of the payload
-            m = hashlib.md5()
-            m.update(payload.encode('utf-8'))
-            imagePath = m.hexdigest() + '.' + mimeType.split('/')[1]
             try:
-                with open(dumpDir + '/' + imagePath, 'wb') as f:
+                with open(payload_path, 'wb') as f:
                     f.write(imgdata)
-                logger.info('Decoded %s' % imagePath)
-                imagesList.append(dumpDir + '/' + imagePath)
+                logger.info('Decoded %s' % payload_path)
+                imagesList.append(payload_path)
             except Exception as e:
                 logger.warning(f'Decoding this MIME part returned error: {e}')
-        else:
-            fileName = part.get_filename()
-            if not fileName:
-                fileName = "Unknown"
-            attachments.append("%s (%s)" % (fileName, mimeType))
-            logger.info('Skipped attachment %s (%s)' % (fileName, mimeType))
-
-    resultImage = dumpDir + '/' + 'output.png'
+        # else:
+        #     fileName = part.get_filename()
+        #     if not fileName:
+        #         fileName = "Unknown"
+        #     attachments.append("%s (%s)" % (fileName, mimeType))
+        #     logger.info('Skipped attachment %s (%s)' % (fileName, mimeType))
+
+    resultImage = os.path.join(output_dir, 'output.png')
     if len(imagesList) > 0:
         images = list(map(Image.open, imagesList))
         combo = appendImages(images)

service_manifest.yml

@@ -23,6 +23,16 @@ submission_params:
     value: 1
     default: 1
 
+  - name: run_ocr
+    type: bool
+    value: false
+    default: false
+
+  - name: load_email_images
+    type: bool
+    value: false
+    default: false
+
 heuristics:
   - heur_id: 1
     name: OCR Detection found