"""Shipper utility functions for Mail and Packages.""" from __future__ import annotations import base64 import email import logging import re from pathlib import Path from typing import Any from aioimaplib import IMAP4_SSL from custom_components.mail_and_packages.utils.cache import EmailCache from .imap import email_fetch _LOGGER = logging.getLogger(__name__) async def get_tracking( sdata: Any, account: type[IMAP4_SSL], the_format: str | None = None, cache: EmailCache | None = None, ) -> list: """Parse tracking numbers from email. Returns list of tracking numbers """ tracking = [] pattern = re.compile(rf"{the_format}") mail_list = sdata.split() _LOGGER.debug("Searching for tracking numbers in %s messages...", len(mail_list)) for i in mail_list: if cache: data = (await cache.fetch(i, "(RFC822)"))[1] else: data = (await email_fetch(account, i, "(RFC822)"))[1] for response_part in data: if not isinstance(response_part, (bytes, bytearray)): continue msg = email.message_from_bytes(response_part) # 1. Search subject if found := _find_tracking_in_subject(msg, pattern): if found not in tracking: tracking.append(found) continue # 2. Search body if the_format == "1Z?[0-9A-Z]{16}": found = _find_ups_tracking_in_raw(response_part, pattern) else: found = _find_tracking_in_body(msg, pattern, the_format) if found and found not in tracking: tracking.append(found) return tracking def _find_tracking_in_subject( msg: email.message.Message, pattern: re.Pattern, ) -> str | None: """Find tracking number in email subject.""" email_subject = msg["subject"] if email_subject: email_subject = str(email_subject) if (found := pattern.findall(email_subject)) and len(found) > 0: _LOGGER.debug("Found tracking number in email subject: %s", found[0]) return found[0] return None def _find_ups_tracking_in_raw( response_part: bytes | bytearray, pattern: re.Pattern, ) -> str | None: """UPS specific tracking search in raw email bytes.""" try: email_content = str(response_part, "utf-8", errors="ignore") if (found := pattern.findall(email_content)) and len(found) > 0: _LOGGER.debug("Found tracking number in email: %s", found[0]) return found[0] except (TypeError, UnicodeError) as err: _LOGGER.debug("Error processing email content: %s", err) return None def _find_tracking_in_body( msg: email.message.Message, pattern: re.Pattern, the_format: str, ) -> str | None: """Search for tracking number in email body parts.""" for part in msg.walk(): if part.get_content_type() not in ["text/html", "text/plain"]: continue email_msg = part.get_payload(decode=True) email_msg = email_msg.decode("utf-8", "ignore") if (found := pattern.findall(email_msg)) and len(found) > 0: tracking_num = found[0] # DHL is special if " " in the_format: tracking_num = tracking_num.split(" ")[1] _LOGGER.debug("Found tracking number in email body: %s", tracking_num) return tracking_num return None def save_image_data_to_disk( shipper_name: str, path: str, image_data: bytes | None ) -> bool: """Write image bytes to disk and verify.""" if not image_data: # Extraction can hand us zero bytes (e.g. a bare/empty base64 data URI # in the email HTML). Writing that produces a 0-byte "photo" the # camera then serves as a broken image — report failure instead so # the caller falls through to the next extraction pass. _LOGGER.debug( "%s - No image data extracted; not writing %s", shipper_name, path, ) return False try: # Ensure directory exists directory = Path(path).parent if not directory.is_dir(): _LOGGER.debug("%s - Creating directory: %s", shipper_name, directory) try: directory.mkdir(parents=True, exist_ok=True) except OSError as err: _LOGGER.error("Error creating directory: %s", err) return False _LOGGER.debug( "%s - Writing %d bytes to file: %s", shipper_name, len(image_data), path, ) with Path(path).open("wb") as the_file: the_file.write(image_data) except OSError as err: _LOGGER.error( "Error saving %s delivery photo to %s: %s", shipper_name, path, err, ) return False else: if Path(path).exists(): return True _LOGGER.error( "%s - ERROR: File write reported success but file doesn't exist: %s", shipper_name, path, ) return False def generic_delivery_image_extraction( sdata: Any, image_path: str, image_name: str, shipper_name: str, image_type: str, cid_name: str | None = None, attachment_filename_pattern: str | None = None, ) -> bool: """Extract delivery photos from email.""" _LOGGER.debug("Attempting to extract %s delivery photo", shipper_name) msg = ( email.message_from_bytes(sdata) if isinstance(sdata, (bytes, bytearray)) else email.message_from_string(sdata) ) normalized_image_path = image_path.rstrip("/") + "/" shipper_path = f"{normalized_image_path}{shipper_name}/" full_path = shipper_path + image_name # Pass 1: CID search if cid_name and ( found := _extract_from_cid(msg, cid_name, shipper_name, full_path, image_type) ): return found # Pass 2: HTML/Base64 search if found := _extract_from_html(msg, cid_name, shipper_name, full_path, image_type): return found # Pass 3: Attachment search return _extract_from_attachments( msg, attachment_filename_pattern, shipper_name, full_path, image_type, ) def _extract_from_cid( msg: email.message.Message, cid_name: str, shipper_name: str, full_path: str, image_type: str, ) -> bool: """Pass 1: Look for CID embedded images.""" content_type = f"image/{image_type}" for part in msg.walk(): if part.get_content_type() == content_type: content_id = part.get("Content-ID") if content_id and content_id.strip("<>") == cid_name: return save_image_data_to_disk( shipper_name, full_path, part.get_payload(decode=True), ) return False def _extract_from_html( msg: email.message.Message, cid_name: str | None, shipper_name: str, full_path: str, image_type: str, ) -> bool: """Pass 2: Look for HTML content with CID references or base64.""" base64_pattern = rf"data:image/{image_type};base64,((?:[A-Za-z0-9+/]{{4}})*(?:[A-Za-z0-9+/]{{2}}==|[A-Za-z0-9+/]{{3}}=)?)" for part in msg.walk(): if part.get_content_type() != "text/html": continue payload = part.get_payload(decode=True) content = ( payload.decode("utf-8", "ignore") if isinstance(payload, bytes) else str(payload) ) # Base64 check. Skip empty matches: a bare "data:image/...;base64," # URI (seen in real FedEx delivered emails) matches zero characters, # and taking it would discard a real photo later in the same part. matches = [m for m in re.findall(base64_pattern, content) if m] if matches: try: base64_data = matches[0].replace(" ", "").replace("=3D", "=") return save_image_data_to_disk( shipper_name, full_path, base64.b64decode(base64_data), ) except (OSError, ValueError, TypeError) as err: _LOGGER.error( "Error saving %s delivery photo from base64: %s", shipper_name, err, ) return False return False def _extract_from_attachments( msg: email.message.Message, pattern: str | None, shipper_name: str, full_path: str, image_type: str, ) -> bool: """Pass 3: Look for attachments.""" content_type = f"image/{image_type}" for part in msg.walk(): if part.get_content_type() == content_type: filename = part.get_filename() if not filename: continue if pattern and pattern.lower() not in filename.lower(): continue try: return save_image_data_to_disk( shipper_name, full_path, part.get_payload(decode=True), ) except (OSError, ValueError, TypeError) as err: _LOGGER.error( "Error saving %s delivery photo to %s: %s", shipper_name, full_path, err, ) return False return False