"""USPS Shipper class.""" from __future__ import annotations import base64 import email import logging import re import shutil from pathlib import Path from typing import Any import anyio from aioimaplib import IMAP4_SSL from bs4 import BeautifulSoup from custom_components.mail_and_packages.const import ( ATTR_COUNT, ATTR_EMAIL, ATTR_GRID_IMAGE_NAME, ATTR_IMAGE_PATH, ATTR_SUBJECT, ATTR_USPS_IMAGE, ATTR_USPS_MAIL, CONF_CUSTOM_IMG_FILE, CONF_DURATION, CONF_FORWARDING_HEADER, CONF_GENERATE_GRID, CONF_GENERATE_MP4, CONF_USPS_PLACEHOLDER, DEFAULT_CUSTOM_IMG_FILE, SENSOR_DATA, ) from custom_components.mail_and_packages.utils.cache import EmailCache from custom_components.mail_and_packages.utils.date import get_formatted_date from custom_components.mail_and_packages.utils.image import ( _generate_mp4, cleanup_images, copy_overlays, generate_delivery_gif, generate_grid_img, io_save_file, random_filename, resize_images, ) from custom_components.mail_and_packages.utils.imap import email_fetch, email_search from .base import Shipper _LOGGER = logging.getLogger(__name__) class USPSShipper(Shipper): """USPS Informed Delivery shipper.""" @property def name(self) -> str: """Return shipper name.""" return "usps" @classmethod def handles_sensor(cls, sensor_type: str) -> bool: """Return True if this shipper handles the given sensor type.""" return sensor_type == "usps_mail" async def process( self, account: IMAP4_SSL, date: str, sensor_type: str, cache: EmailCache | None = None, ) -> dict[str, Any]: """Process USPS Informed Delivery emails.""" config = self._get_usps_config() image_count = 0 images = [] images_delete = [] (server_response, data) = await self._search_informed_delivery(account) # Bail out on error if server_response != "OK" or data[0] is None: return {ATTR_COUNT: image_count} # Setup image directory and overlays if not await self._setup_image_directory(config["image_output_path"]): return {ATTR_COUNT: image_count} all_msg_content = "" if server_response == "OK": _LOGGER.debug("Informed Delivery email found processing...") for num in data[0].split(): (image_count, images, email_content) = await self._process_usps_email( account, num, config["image_output_path"], image_count, images, cache, ) all_msg_content += email_content # Process images images = await self._process_usps_images(all_msg_content, images) image_count = len(images) # Generate filtered list for GIF/MP4/Grid gif_images = images.copy() if not config.get("usps_placeholder", True): placeholder_str = str( Path(__file__).parent.parent / "image-no-mailpieces700.jpg" ) if placeholder_str in gif_images: gif_images.remove(placeholder_str) # Generate camera media await self._create_camera_media(gif_images, config, images_delete) return { ATTR_COUNT: image_count, ATTR_USPS_IMAGE: config["image_name"], ATTR_IMAGE_PATH: config["image_output_path"], ATTR_GRID_IMAGE_NAME: config["image_name"].replace(".gif", "_grid.png"), } async def process_batch( self, account: IMAP4_SSL, date: str, sensors: list[str], cache: EmailCache, since_date: str | None = None, ) -> dict[str, Any]: """Process multiple USPS sensors in batch.""" res = {} for sensor in sensors: res.update(await self.process(account, date, sensor, cache)) # Replicatecoordinator dict structure if sensor not in res: if ATTR_COUNT in res: res[sensor] = res[ATTR_COUNT] # Don't pop ATTR_COUNT because other things might need it, actually # coordinator used to pop it via explicit assignment return res async def _create_camera_media( self, gif_images: list, config: dict, images_delete: list, ): """Create camera media (GIF, MP4, and Grid).""" if len(gif_images) > 0: await self._generate_mail_image( gif_images, config["image_output_path"], config["image_name"], config["gif_duration"], images_delete, ) else: await self._copy_nomail_image( config["image_output_path"], config["image_name"], config["custom_img"], ) if config["gen_mp4"]: await self._generate_mp4_video( config["image_output_path"], config["image_name"], ) if config["gen_grid"]: await self._generate_grid_image( config["image_output_path"], config["image_name"], len(gif_images), ) async def _generate_mp4_video(self, path: str, name: str): """Generate MP4 video from images.""" await self.hass.async_add_executor_job(_generate_mp4, path, name) async def _generate_grid_image(self, path: str, name: str, count: int): """Generate grid image from images.""" await self.hass.async_add_executor_job(generate_grid_img, path, name, count) async def _process_usps_images(self, content: str, images: list) -> list: """Process USPS images (placeholder and filtering).""" # Old USPS format: plain-text email body contained the filename as a reference. # New format is handled in _extract_usps_images on properly decoded HTML. if re.compile(r"\bimage-no-mailpieces?700\.jpg\b").search(content) is not None: placeholder = Path(__file__).parent.parent / "image-no-mailpieces700.jpg" placeholder_str = str(placeholder) if placeholder.exists() and placeholder_str not in images: images.append(placeholder_str) _LOGGER.debug( "Placeholder image found using: image-no-mailpieces700.jpg.", ) # Announcement images removal return self._remove_announcement_images(images) def _remove_announcement_images(self, images: list) -> list: """Remove announcement images.""" return [ el for el in images if not any( ignore in el for ignore in ["mailerProvidedImage", "ra_0", "Mail Attachment.txt"] ) ] async def _generate_mail_image( self, images: list, path: str, name: str, duration: int, delete_list: list, ): """Generate animated GIF from mail images.""" _LOGGER.debug("Resizing images to 724x320...") all_images = await self.hass.async_add_executor_job( resize_images, images, 724, 320, ) delete_list.extend(all_images) try: _LOGGER.debug("Generating animated GIF") gif_path = str(Path(path) / name) await self.hass.async_add_executor_job( generate_delivery_gif, all_images, gif_path, duration * 1000, ) _LOGGER.debug("Mail image generated.") except (OSError, ValueError) as err: _LOGGER.error("Error attempting to generate image: %s", err) for image in delete_list: await self.hass.async_add_executor_job( cleanup_images, f"{Path(image).parent}/", Path(image).name, ) async def _copy_nomail_image(self, path: str, name: str, custom_img: str | None): """Copy the 'no mail' placeholder image.""" def _prepare(): if not Path(path).exists(): Path(path).mkdir(parents=True, exist_ok=True) target = Path(path) / name if target.is_file(): cleanup_images(path, name) src = custom_img or str(Path(__file__).parent.parent / "mail_none.gif") shutil.copyfile(src, str(target)) _LOGGER.debug("No mail found.") try: await self.hass.async_add_executor_job(_prepare) except OSError as err: _LOGGER.error("Error attempting to copy image: %s", err) def _get_usps_config(self) -> dict: """Get USPS specific configuration.""" return { "image_output_path": self.config.get("image_path"), "gif_duration": self.config.get(CONF_DURATION), "image_name": self.config.get("usps_image"), "gen_mp4": self.config.get(CONF_GENERATE_MP4), "custom_img": self.config.get(CONF_CUSTOM_IMG_FILE) or DEFAULT_CUSTOM_IMG_FILE, "gen_grid": self.config.get(CONF_GENERATE_GRID), "usps_placeholder": self.config.get(CONF_USPS_PLACEHOLDER, True), } async def _search_informed_delivery(self, account: IMAP4_SSL) -> tuple: """Search for USPS Informed Delivery emails.""" _LOGGER.debug("Attempting to find Informed Delivery mail") _LOGGER.debug("Informed delivery search date: %s", get_formatted_date()) forwarding_header = self.config.get(CONF_FORWARDING_HEADER, "") if forwarding_header and forwarding_header != "(none)": email_addresses = SENSOR_DATA[ATTR_USPS_MAIL][ATTR_EMAIL] else: forwarding_header = "" forwarded_emails = self.config.get("forwarded_emails", []) if isinstance(forwarded_emails, str): forwarded_emails = [ e.strip() for e in forwarded_emails.split(",") if e.strip() ] if forwarded_emails: email_addresses = ( forwarded_emails + SENSOR_DATA[ATTR_USPS_MAIL][ATTR_EMAIL] ) else: email_addresses = SENSOR_DATA[ATTR_USPS_MAIL][ATTR_EMAIL] return await email_search( account=account, address=email_addresses, date=get_formatted_date(), subject=SENSOR_DATA[ATTR_USPS_MAIL][ATTR_SUBJECT][0], header=forwarding_header, ) async def _setup_image_directory(self, path: str) -> bool: """Ensure image directory exists and is prepared.""" if not await anyio.Path(path).is_dir(): try: await anyio.Path(path).mkdir(parents=True, exist_ok=True) except OSError as err: _LOGGER.error("Error creating directory: %s", err) return False # Clean up and setup overlays await self.hass.async_add_executor_job(cleanup_images, path) await self.hass.async_add_executor_job(copy_overlays, path) return True async def _process_usps_email( self, account: IMAP4_SSL, num: str, image_output_path: str, image_count: int, images: list, cache: EmailCache | None = None, ) -> tuple[int, list, str]: """Process a single USPS Informed Delivery email.""" if cache: msg_parts = (await cache.fetch(num, "(RFC822)"))[1] else: msg_parts = (await email_fetch(account, num, "(RFC822)"))[1] _LOGGER.debug("Processing email number: %s", num) all_content = "" for response_part in msg_parts: if isinstance(response_part, (bytes, bytearray)): all_content += str(response_part, "utf-8", errors="ignore") msg = email.message_from_bytes(response_part) for part in msg.walk(): if part.get_content_type() == "text/html": (image_count, images) = await self._extract_usps_images( part, image_output_path, image_count, images, ) elif part.get_content_type() == "image/jpeg": (image_count, images) = await self._extract_jpeg_attachment( part, image_output_path, image_count, images, ) return image_count, images, all_content async def _extract_usps_images( self, part: email.message.Message, image_output_path: str, image_count: int, images: list, ) -> tuple[int, list]: """Extract images from an email part (HTML/Base64).""" payload = part.get_payload(decode=True) content = ( payload.decode("utf-8", "ignore") if isinstance(payload, (bytes, bytearray)) else str(payload) ) # New USPS format: unscanned mailpieces use a div with a specific id. # Check here on properly decoded HTML — raw RFC822 content is # quoted-printable encoded and soft line breaks could split the string. if "mailpiece-with-no-image-id" in content: placeholder = Path(__file__).parent.parent / "image-no-mailpieces700.jpg" placeholder_str = str(placeholder) if placeholder.exists() and placeholder_str not in images: images.append(placeholder_str) image_count += 1 _LOGGER.debug( "Placeholder image found using: image-no-mailpieces700.jpg.", ) if "data:image/jpeg;base64" not in content: return image_count, images soup = BeautifulSoup(content, "html.parser") found_images = soup.find_all(id="mailpiece-image-src-id") for image in found_images: filename = random_filename() img_data = str(image["src"]).split(",")[1] try: target_path = Path(image_output_path) / filename await self.hass.async_add_executor_job( io_save_file, target_path, base64.b64decode(img_data), ) images.append(str(target_path)) image_count += 1 except (OSError, ValueError, TypeError) as err: _LOGGER.error("Error extracting image: %s", err) return image_count, images async def _extract_jpeg_attachment( self, part: email.message.Message, image_output_path: str, image_count: int, images: list, ) -> tuple[int, list]: """Extract image from JPEG attachment.""" _LOGGER.debug("Extracting image from email attachment") filename = part.get_filename() junkmail = ["mailer", "content", "package"] if filename is None: return image_count, images if any(junk in filename for junk in junkmail): return image_count, images try: target_path = Path(image_output_path) / filename await self.hass.async_add_executor_job( io_save_file, target_path, part.get_payload(decode=True), ) images.append(str(target_path)) image_count += 1 except OSError as err: _LOGGER.critical("Error opening filepath: %s", err) return image_count, images