1"""
  2.. include:: ../README.md
  3"""
  4
  5import datetime
  6import email
  7import email.header
  8import imaplib
  9import logging
 10import mailbox
 11import os
 12import re
 13import time
 14
 15__all__ = ["IMAPMailbox", "IMAPMessage", "IMAPError"]
 16
 17MESSAGE_HEAD_RE = re.compile(r"(\d+) \(([^\s]+) {(\d+)}$")
 18FOLDER_DATA_RE = re.compile(r"\(([^)]+)\) \"([^\"]+)\" \"?([^\"]+)\"?$")
 19
 20
 21log = logging.getLogger(__name__)
 22log.setLevel(getattr(logging, os.getenv("LOG_LEVEL", "INFO")))
 23
 24
 25class IMAPError(Exception):
 26    """Exception raised for IMAP operation errors."""
 27
 28    pass
 29
 30
 31def handle_response(response):
 32    """Handle the response from the IMAP server"""
 33    status, data = response
 34    if status != "OK":
 35        raise IMAPError(data[0])
 36
 37    return data
 38
 39
 40def change_time(time, weeks=0, days=0, hours=0, minutes=0, seconds=0):
 41    """Change the time by a given amount of days, hours, minutes and seconds"""
 42    return time + datetime.timedelta(
 43        weeks=weeks, days=days, hours=hours, minutes=minutes, seconds=seconds
 44    )
 45
 46
 47def imap_date(time):
 48    """Convert a datetime object to an IMAP date string"""
 49    return time.strftime("%d-%b-%Y")
 50
 51
 52def imap_date_range(start, end):
 53    """Create an IMAP date range string for use in a search query"""
 54    return f"(SINCE {imap_date(start)} BEFORE {imap_date(end)})"
 55
 56
 57class IMAPMessage(mailbox.Message):
 58    """A Mailbox Message class that uses an IMAPClient object to fetch the message
 59
 60    Supports lazy loading: messages can be created with headers only, and the full
 61    body is fetched transparently when accessed.
 62    """
 63
 64    def __init__(self, message=None, uid=None, mailbox_ref=None):
 65        """Create a new IMAPMessage
 66
 67        Args:
 68            message: Email message bytes or Message object
 69            uid: Message UID for lazy loading (optional)
 70            mailbox_ref: Reference to IMAPMailbox for lazy loading (optional)
 71        """
 72        super().__init__(message)
 73        self._uid = uid
 74        self._mailbox_ref = mailbox_ref
 75        self._body_loaded = (
 76            mailbox_ref is None
 77        )  # If no mailbox ref, body is already loaded
 78
 79    @classmethod
 80    def from_uid(cls, uid, mailbox, headers_only=False):
 81        """Create a new message from a UID
 82
 83        Args:
 84            uid: Message UID
 85            mailbox: IMAPMailbox instance
 86            headers_only: If True, fetch only headers for lazy loading
 87        """
 88        if headers_only:
 89            # Fetch headers only, store reference for lazy body loading
 90            _, body = next(mailbox.fetch(uid, "RFC822.HEADER"))
 91            return cls(body, uid=uid, mailbox_ref=mailbox)
 92        else:
 93            # Fetch full message immediately
 94            _, body = next(mailbox.fetch(uid, "RFC822"))
 95            return cls(body, uid=uid)
 96
 97    @property
 98    def uid(self):
 99        """Get the message UID"""
100        return self._uid
101
102    def _ensure_body_loaded(self):
103        """Ensure the full message body is loaded
104
105        If the message was created with headers_only=True, this will fetch
106        the full message from the IMAP server.
107
108        Raises:
109            RuntimeError: If the IMAP connection is closed
110        """
111        if self._body_loaded:
112            return
113
114        if self._mailbox_ref is None:
115            raise RuntimeError("Cannot load body: IMAP connection is closed")
116
117        # Fetch the full message
118        _, body = next(self._mailbox_ref.fetch(self._uid, "RFC822"))
119
120        # Parse the full message
121        full_msg = email.message_from_bytes(body)
122
123        # Update our payload from the parsed message
124        self._payload = full_msg._payload
125
126        # Clear the mailbox reference to allow garbage collection
127        self._mailbox_ref = None
128        self._body_loaded = True
129
130    def __getitem__(self, name: str):
131        """Get a message header
132
133        This method overrides the default implementation of accessing a message headers.
134        The header is decoded using the email.header.decode_header method. This allows
135        for the retrieval of headers that contain non-ASCII characters.
136        """
137        original_header = super().__getitem__(name)
138
139        if original_header is None:
140            return None
141
142        decoded_pairs = email.header.decode_header(original_header)
143        decoded_chunks = []
144        for data, charset in decoded_pairs:
145            if isinstance(data, str):
146                decoded_chunks.append(data)
147            elif charset is None:
148                decoded_chunks.append(data.decode())
149            elif charset == "unknown-8bit":
150                decoded_chunks.append(data.decode("utf-8", "replace"))
151            else:
152                decoded_chunks.append(data.decode(charset, "replace"))
153
154        return " ".join(decoded_chunks)
155
156    # Override body-accessing methods to ensure body is loaded
157
158    def get_payload(self, *args, **kwargs):
159        """Get the message payload, ensuring body is loaded"""
160        self._ensure_body_loaded()
161        return super().get_payload(*args, **kwargs)
162
163    def is_multipart(self):
164        """Check if message is multipart, ensuring body is loaded"""
165        self._ensure_body_loaded()
166        return super().is_multipart()
167
168    def walk(self):
169        """Walk the message tree, ensuring body is loaded"""
170        self._ensure_body_loaded()
171        return super().walk()
172
173    def as_string(self, *args, **kwargs):
174        """Return message as string, ensuring body is loaded"""
175        self._ensure_body_loaded()
176        return super().as_string(*args, **kwargs)
177
178    def as_bytes(self, *args, **kwargs):
179        """Return message as bytes, ensuring body is loaded"""
180        self._ensure_body_loaded()
181        return super().as_bytes(*args, **kwargs)
182
183    def set_payload(self, *args, **kwargs):
184        """Set the message payload, ensuring body is loaded"""
185        self._ensure_body_loaded()
186        return super().set_payload(*args, **kwargs)
187
188    def attach(self, *args, **kwargs):
189        """Attach a payload, ensuring body is loaded"""
190        self._ensure_body_loaded()
191        return super().attach(*args, **kwargs)
192
193
194class IMAPMailbox(mailbox.Mailbox):
195    """A Mailbox class that uses an IMAPClient object as the backend"""
196
197    def __init__(self, host, user, password, folder="INBOX", port=993, security="SSL"):
198        """Create a new IMAPMailbox object"""
199        self.host = host
200        self.user = user
201        self.password = password
202        self.__folder = folder
203        self.__security = security
204        self.__port = port
205
206    def connect(self):
207        """Connect to the IMAP server"""
208        if self.__security == "SSL":
209            log.info("Connecting to IMAP server using SSL")
210            self.__m = imaplib.IMAP4_SSL(self.host, self.__port)
211        elif self.__security == "STARTTLS":
212            log.info("Connecting to IMAP server using STARTTLS")
213            self.__m = imaplib.IMAP4(self.host, self.__port)
214            self.__m.starttls()
215        elif self.__security == "PLAIN":
216            log.info("Connecting to IMAP server without encryption (insecure)")
217            self.__m = imaplib.IMAP4(self.host, self.__port)
218        else:
219            raise ValueError("Invalid security type")
220        self.__m.login(self.user, self.password)
221        self.select(self.__folder)
222
223    def disconnect(self):
224        """Disconnect from the IMAP server"""
225
226        log.info("Disconnecting from IMAP server")
227        self.__m.close()
228        self.__m.logout()
229
230    def __enter__(self):
231        self.connect()
232        return self
233
234    def __exit__(self, *args):
235        self.disconnect()
236
237    def __iter__(self):
238        """Iterate over all messages in the mailbox"""
239        data = handle_response(self.__m.search(None, "ALL"))
240        for uid in data[0].decode().split():
241            yield IMAPMessage.from_uid(uid, self, headers_only=True)
242
243    def values(self):
244        yield from iter(self)
245
246    def keys(self) -> list[str]:
247        """Get a list of all message UIDs in the mailbox"""
248        data = handle_response(self.__m.search(None, "ALL"))
249        return data[0].decode().split()
250
251    def iterkeys(self):
252        """Return an iterator over keys."""
253        data = handle_response(self.__m.search(None, "ALL"))
254        yield from data[0].decode().split()
255
256    def __contains__(self, key):
257        """Return True if the keyed message exists, False otherwise."""
258        return str(key) in self.keys()
259
260    def get_bytes(self, key):
261        """Return a byte string representation or raise KeyError."""
262        if key not in self:
263            raise KeyError(key)
264        _, body = next(self.fetch(key, "RFC822"))
265        return body
266
267    def get_file(self, key):
268        """Return a file-like representation or raise KeyError."""
269        import io
270
271        return io.BytesIO(self.get_bytes(key))
272
273    def get_message(self, key):
274        """Return a Message representation or raise KeyError."""
275        if key not in self:
276            raise KeyError(key)
277        return IMAPMessage.from_uid(key, self, headers_only=False)
278
279    def __getitem__(self, key):
280        """Return the keyed message; raise KeyError if it doesn't exist."""
281        return self.get_message(key)
282
283    def __setitem__(self, key, message):
284        """Replace the keyed message; raise KeyError if it doesn't exist."""
285        if key not in self:
286            raise KeyError(key)
287        self.remove(key)
288        self.add(message)
289
290    def items(self):
291        """Iterate over all messages as (uid, message) tuples"""
292        data = handle_response(self.__m.search(None, "ALL"))
293        for uid in data[0].decode().split():
294            msg = IMAPMessage.from_uid(uid, self, headers_only=True)
295            yield uid, msg
296
297    @property
298    def capability(self):
299        """Get the server capabilities"""
300        return handle_response(self.__m.capability())[0].decode()
301
302    def add(self, message):
303        """Add a message to the mailbox"""
304
305        self.__m.append(
306            self.current_folder,
307            "",
308            imaplib.Time2Internaldate(time.time()),
309            message.as_bytes(),
310        )
311
312    def copy(self, messageset: bytes, folder: str) -> None:
313        """Copy a message to a different folder"""
314
315        self.__m.copy(messageset, folder)
316
317    def move(self, messageset: bytes, folder: str) -> None:
318        """Move a message to a different folder"""
319
320        self.__m._simple_command("MOVE", messageset, folder)
321
322    def remove(self, key):
323        """Remove the keyed message; raise KeyError if it doesn't exist."""
324        if key not in self:
325            raise KeyError(key)
326        self.__m.store(key, "+FLAGS", "\\Deleted")
327        self.__m.expunge()
328
329    def discard(self, key):
330        """If the keyed message exists, remove it."""
331        try:
332            self.remove(key)
333        except KeyError:
334            pass
335
336    def __delitem__(self, key):
337        """Remove the keyed message; raise KeyError if it doesn't exist."""
338        self.remove(key)
339
340    def clear(self):
341        """Remove all messages from the mailbox."""
342        for key in self.keys():
343            self.__m.store(key, "+FLAGS", "\\Deleted")
344        self.__m.expunge()
345
346    def __len__(self) -> int:
347        return len(self.keys())
348
349    def fetch(self, messageset: bytes, what):
350        """Fetch messages from the mailbox"""
351
352        response = handle_response(self.__m.fetch(messageset, what))
353
354        # Filter response to only include message data (tuples), not FLAGS (bytes)
355        messages = [item for item in response if isinstance(item, tuple)]
356
357        for head, body in messages:
358            uid, what, size = MESSAGE_HEAD_RE.match(head.decode()).groups()
359            if size != str(len(body)):
360                raise IMAPError("Size mismatch")
361
362            yield uid, body
363
364    def __expand_search_macros(self, query) -> str:
365        """Expand search macros in the query."""
366
367        today = datetime.date.today()
368        yesterday = today - datetime.timedelta(days=1)
369
370        week_start = today - datetime.timedelta(days=today.weekday())
371        last_week_start = week_start - datetime.timedelta(days=7)
372
373        month_start = datetime.date(today.year, today.month, 1)
374        year_start = datetime.date(today.year, 1, 1)
375
376        if today.month == 1:  # January
377            # last month is December of the previous year
378            last_month_start = datetime.date(today.year - 1, 12, 1)
379        else:
380            last_month_start = datetime.date(today.year, today.month - 1, 1)
381
382        last_year_start = datetime.date(today.year - 1, 1, 1)
383
384        q = query
385        q = q.replace("FIND", "TEXT")
386
387        q = q.replace("TODAY", f"ON {imap_date(today)}")
388        q = q.replace("YESTERDAY", f"ON {imap_date(yesterday)}")
389
390        q = q.replace("THISWEEK", f"SINCE {imap_date(week_start)}")
391        q = q.replace("THISMONTH", f"SINCE {imap_date(month_start)}")
392        q = q.replace("THISYEAR", f"SINCE {imap_date(year_start)}")
393
394        q = q.replace("LASTWEEK", imap_date_range(last_week_start, week_start))
395        q = q.replace("LASTMONTH", imap_date_range(last_month_start, month_start))
396        q = q.replace("LASTYEAR", imap_date_range(last_year_start, year_start))
397
398        # shortcuts
399        q = q.replace("PASTDAY", "PAST1DAY")
400        q = q.replace("PASTWEEK", "PAST1WEEK")
401        q = q.replace("PASTMONTH", "PAST1MONTH")
402        q = q.replace("PASTYEAR", "PAST1YEAR")
403
404        # use regex to match the PASTXDAYS macro
405        q = re.sub(
406            r"PAST(\d+)DAYS?",
407            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1))))}",
408            q,
409        )
410
411        # use regex to match the PASTXWEEKS macro
412        q = re.sub(
413            r"PAST(\d+)WEEKS?",
414            lambda m: f"SINCE {imap_date(change_time(today, weeks=-int(m.group(1))))}",
415            q,
416        )
417
418        # use regex to match the PASTXMONTHS macro
419        q = re.sub(
420            r"PAST(\d+)MONTHS?",
421            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1)) * 30))}",
422            q,
423        )
424
425        # use regex to match the PASTXYEARS macro
426        q = re.sub(
427            r"PAST(\d+)YEARS?",
428            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1)) * 365))}",
429            q,
430        )
431
432        return q
433
434    def search(self, query):
435        """Search for messages matching the query
436
437        We support extra search macros in the search query in addition to
438        the standard IMAP search macros.
439
440        One search macro is FIND <text>, which is an alias for TEXT.
441        The rest of the macros deal with date ranges.
442
443        The date range macros are expanded to the appropriate date range and
444        are relative to the current date.
445        Example: TODAY expands to ON <date>, where <date> is today's date.
446
447        Note that some of these macros will expand to multiple search terms.
448        Expansions that result in multiple search terms are wrapped in parentheses.
449        Example: LASTWEEK expands to (SINCE <date1> BEFORE <date2>).
450
451        The following extra macros are supported:
452
453
454        - FIND <text> - alias for TEXT, searches the message headers and body
455
456        Current period:
457        - TODAY - messages from today
458        - THISWEEK - messages since the start of the week, Monday to Sunday
459        - THISMONTH - messages since the start of the month
460        - THISYEAR - messages since the start of the year
461
462        Previous period:
463        - YESTERDAY - messages from yesterday
464        - LASTWEEK - messages from the week before
465        - LASTMONTH - messages from the month before
466        - LASTYEAR - messages from the year before
467
468        Periods starting from now:
469
470        _These are just shortcuts_
471        - PASTDAY - messages from the past 1 day, same as PAST1DAY
472        - PASTWEEK - messages from the past 1 week, same as PAST1WEEK
473        - PASTMONTH - messages from the past 30 days, same as PAST1MONTH
474        - PASTYEAR - messages from the past 365 days, same as PAST1YEAR
475
476        _These are pattern matching macros_
477        - PASTXDAYS - messages from the past X days
478        - PASTXWEEKS - messages from the past X weeks
479        - PASTXMONTHS - messages from the past X * 30 days
480        - PASTXYEARS - messages from the past X * 365 days
481
482        These macros can be combined with other search macros, and can be
483        negated with NOT. For example, to search and archive or delete messages with a short
484        relevance period, you can use `NOT PAST3DAYS`, use `NOT PAST3MONTHS` to search for
485        messages older than a quarter, or use `NOT PAST2YEAR` to search for messages older than
486        two years.
487
488        _The `NOT` modifier is very useful for mailbox maintenance_
489
490        _There are no options for hours, because the range seletion does not have time of day precision._
491
492        Returns:
493            bytes: A comma-separated list of message UIDs
494        """
495
496        expanded_query = self.__expand_search_macros(query)
497        data = handle_response(self.__m.search(None, expanded_query))
498        num_results = len(data[0].split(b" "))
499
500        log.info(f"Searching for messages matching: {query}")
501        if expanded_query != query:
502            log.info(f"Expanded search query to: {expanded_query}")
503        log.info(f"Found {num_results} results")
504
505        return data[0].replace(b" ", b",")
506
507    def list_folders(self) -> tuple:
508        """List all folders in the mailbox
509
510        Returns:
511            tuple: A tuple of flags, delimiter, folder name, and folder display name
512        """
513
514        folders_data = handle_response(self.__m.list())
515        for data in folders_data:
516            flags, delimiter, folder = FOLDER_DATA_RE.match(data.decode()).groups()
517            display_name = folder.split(delimiter)[-1]
518            yield (flags, delimiter, folder, display_name)
519
520    @property
521    def current_folder(self):
522        """Get the currently selected folder"""
523        return self.__folder
524
525    def select(self, folder):
526        """Select a folder"""
527        self.__folder = folder
528        self.__m.select(folder)
529        return self
530
531    def flush(self):
532        """Write any pending changes to the disk."""
533        pass  # IMAP changes are immediate
534
535    def lock(self):
536        """Lock the mailbox."""
537        pass  # IMAP handles locking server-side
538
539    def unlock(self):
540        """Unlock the mailbox if it is locked."""
541        pass  # IMAP handles locking server-side
542
543    def close(self):
544        """Flush and close the mailbox."""
545        self.flush()
546        self.disconnect()