feat: parse new USPS informed delivery digest emails (#739)

* feat: parse new USPS informed delivery digest emails

* linting

custom_components/mail_and_packages/helpers.py

@@ -1,5 +1,6 @@
 """Helper functions for Mail and Packages."""
 
+import base64
 import datetime
 import email
 import hashlib
@@ -17,6 +18,7 @@
 from typing import Any, List, Optional, Type, Union
 
 import aiohttp
+from bs4 import BeautifulSoup
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
@@ -632,31 +634,61 @@ def get_mails(
     if server_response == "OK":
         _LOGGER.debug("Informed Delivery email found processing...")
         for num in data[0].split():
-            msg = email.message_from_string(
-                email_fetch(account, num, "(RFC822)")[1][0][1].decode("utf-8")
-            )
-
-            # walking through the email parts to find images
-            for part in msg.walk():
-                if part.get_content_maintype() == "multipart":
-                    continue
-                if part.get("Content-Disposition") is None:
-                    continue
+            msg = email_fetch(account, num, "(RFC822)")[1]
+            for response_part in msg:
+                if isinstance(response_part, tuple):
+                    msg = email.message_from_bytes(response_part[1])
+                    _LOGGER.debug("msg: %s", msg)
+
+                    # walking through the email parts to find images
+                    for part in msg.walk():
+                        if part.get_content_type() == "text/html":
+                            _LOGGER.debug("Found html email processing...")
+                            part = part.get_payload(decode=True)
+                            part = part.decode("utf-8", "ignore")
+                            soup = BeautifulSoup(part, "html.parser")
+                            found_images = soup.find_all(id="mailpiece-image-src-id")
+                            if not found_images:
+                                continue
+                            _LOGGER.debug("Found images: %s", bool(found_images))
 
-                _LOGGER.debug("Extracting image from email")
+                            # Convert all the images to binary data
+                            for image in found_images:
+                                filename = random_filename()
+                                data = str(image["src"]).split(",")[1]
+                                _LOGGER.debug("Data: %s", data)
+                                try:
+                                    with open(
+                                        image_output_path + filename, "wb"
+                                    ) as the_file:
+                                        the_file.write(base64.b64decode(data))
+                                        images.append(image_output_path + filename)
+                                        image_count = image_count + 1
+                                except Exception as err:
+                                    _LOGGER.critical(
+                                        "Error opening filepath: %s", str(err)
+                                    )
+                                    return image_count
+
+                        # Log error message if we are unable to open the filepath for
+                        # some reason
+                        elif part.get_content_type() == "image/jpeg":
+                            _LOGGER.debug("Extracting image from email")
+                            try:
+                                with open(
+                                    image_output_path + part.get_filename(), "wb"
+                                ) as the_file:
+                                    the_file.write(part.get_payload(decode=True))
+                                    images.append(
+                                        image_output_path + part.get_filename()
+                                    )
+                                    image_count = image_count + 1
+                            except Exception as err:
+                                _LOGGER.critical("Error opening filepath: %s", str(err))
+                                return image_count
 
-                # Log error message if we are unable to open the filepath for
-                # some reason
-                try:
-                    with open(
-                        image_output_path + part.get_filename(), "wb"
-                    ) as the_file:
-                        the_file.write(part.get_payload(decode=True))
-                        images.append(image_output_path + part.get_filename())
-                        image_count = image_count + 1
-                except Exception as err:
-                    _LOGGER.critical("Error opening filepath: %s", str(err))
-                    return image_count
+                        elif part.get_content_type() == "multipart":
+                            continue
 
         # Remove duplicate images
         _LOGGER.debug("Removing duplicate images.")
@@ -738,6 +770,11 @@ def get_mails(
     return image_count
 
 
+def random_filename(ext: str = ".jpg") -> str:
+    """Generate random filename."""
+    return f"{str(uuid.uuid4())}{ext}"
+
+
 def _generate_mp4(path: str, image_file: str) -> None:
     """Generate mp4 from gif.
 

custom_components/mail_and_packages/manifest.json

@@ -9,7 +9,7 @@
     "@firstof9"
   ],
   "config_flow": true,
-  "requirements": ["Pillow>=9.0"],
+  "requirements": ["beautifulsoup4","Pillow>=9.0"],
   "iot_class": "cloud_polling",
   "version": "0.0.0-dev"
 }

requirements.txt

@@ -1 +1,2 @@
-Pillow>=9.0
+Pillow>=9.0
+beautifulsoup4

tests/conftest.py

@@ -265,6 +265,32 @@ def mock_imap_usps_informed_digest():
         yield mock_conn
 
 
+@pytest.fixture()
+def mock_imap_usps_new_informed_digest():
+    """Mock imap class values."""
+    with patch(
+        "custom_components.mail_and_packages.helpers.imaplib"
+    ) as mock_imap_usps_new_informed_digest:
+        mock_conn = mock.Mock(spec=imaplib.IMAP4_SSL)
+        mock_imap_usps_new_informed_digest.IMAP4_SSL.return_value = mock_conn
+
+        mock_conn.login.return_value = (
+            "OK",
+            [b"[email protected] authenticated (Success)"],
+        )
+        mock_conn.list.return_value = (
+            "OK",
+            [b'(\\HasNoChildren) "/" "INBOX"'],
+        )
+        mock_conn.search.return_value = ("OK", [b"1"])
+        mock_conn.uid.return_value = ("OK", [b"1"])
+        f = open("tests/test_emails/new_informed_delivery.eml", "r")
+        email_file = f.read()
+        mock_conn.fetch.return_value = ("OK", [(b"", email_file.encode("utf-8"))])
+        mock_conn.select.return_value = ("OK", [])
+        yield mock_conn
+
+
 @pytest.fixture()
 def mock_imap_usps_informed_digest_missing():
     """Mock imap class values."""