1"""
  2.. include:: README.md
  3"""
  4import datetime
  5import email.header
  6import imaplib
  7import logging
  8import mailbox
  9import os
 10import re
 11import time
 12
 13__all__ = ["IMAPMailbox", "IMAPMessage", "IMAPMessageHeadersOnly"]
 14
 15MESSAGE_HEAD_RE = re.compile(r"(\d+) \(([^\s]+) {(\d+)}$")
 16FOLDER_DATA_RE = re.compile(r"\(([^)]+)\) \"([^\"]+)\" \"?([^\"]+)\"?$")
 17
 18
 19log = logging.getLogger(__name__)
 20log.setLevel(getattr(logging, os.getenv("LOG_LEVEL", "INFO")))
 21
 22
 23def handle_response(response):
 24    """Handle the response from the IMAP server"""
 25    status, data = response
 26    if status != "OK":
 27        raise Exception(data[0])
 28
 29    return data
 30
 31
 32def change_time(time, weeks=0, days=0, hours=0, minutes=0, seconds=0):
 33    """Change the time by a given amount of days, hours, minutes and seconds"""
 34    return time + datetime.timedelta(
 35        weeks=weeks, days=days, hours=hours, minutes=minutes, seconds=seconds
 36    )
 37
 38
 39def imap_date(time):
 40    """Convert a datetime object to an IMAP date string"""
 41    return time.strftime("%d-%b-%Y")
 42
 43
 44def imap_date_range(start, end):
 45    """Create an IMAP date range string for use in a search query"""
 46    return f"(SINCE {imap_date(start)} BEFORE {imap_date(end)})"
 47
 48
 49class IMAPMessage(mailbox.Message):
 50    """A Mailbox Message class that uses an IMAPClient object to fetch the message"""
 51
 52    @classmethod
 53    def from_uid(cls, uid, mailbox):
 54        """Create a new message from a UID"""
 55
 56        # fetch the message from the mailbox
 57        uid, body = next(mailbox.fetch(uid, "RFC822"))
 58        return cls(body)
 59
 60    def __getitem__(self, name: str):
 61        """Get a message header
 62
 63        This method overrides the default implementation of accessing a message headers.
 64        The header is decoded using the email.header.decode_header method. This allows
 65        for the retrieval of headers that contain non-ASCII characters.
 66        """
 67
 68        original_header = super().__getitem__(name)
 69
 70        if original_header is None:
 71            return None
 72
 73        decoded_pairs = email.header.decode_header(original_header)
 74        decoded_chunks = []
 75        for data, charset in decoded_pairs:
 76            if isinstance(data, str):
 77                decoded_chunks.append(data)
 78            elif charset is None:
 79                decoded_chunks.append(data.decode())
 80            elif charset == "unknown-8bit":
 81                decoded_chunks.append(data.decode("utf-8", "replace"))
 82            else:
 83                decoded_chunks.append(data.decode(charset, "replace"))
 84
 85        # decode_chunks = (pair[0] for pair in decoded_pairs)
 86
 87        return " ".join(decoded_chunks)
 88
 89
 90class IMAPMessageHeadersOnly(IMAPMessage):
 91    """A Mailbox Message class that uses an IMAPClient object to fetch the message"""
 92
 93    @classmethod
 94    def from_uid(cls, uid, mailbox):
 95        """Create a new message from a UID"""
 96
 97        # fetch headers only message from the mailbox
 98        uid, body = next(mailbox.fetch(uid, "RFC822.HEADER"))
 99        return cls(body)
100
101
102class IMAPMailbox(mailbox.Mailbox):
103    """A Mailbox class that uses an IMAPClient object as the backend"""
104
105    def __init__(self, host, user, password, folder="INBOX", port=993, security="SSL"):
106        """Create a new IMAPMailbox object"""
107        self.host = host
108        self.user = user
109        self.password = password
110        self.__folder = folder
111        self.__security = security
112        self.__port = port
113
114    def connect(self):
115        """Connect to the IMAP server"""
116        if self.__security == "SSL":
117            log.info("Connecting to IMAP server using SSL")
118            self.__m = imaplib.IMAP4_SSL(self.host, self.__port)
119        elif self.__security == "STARTTLS":
120            log.info("Connecting to IMAP server using STARTTLS")
121            self.__m = imaplib.IMAP4(self.host, self.__port)
122            self.__m.starttls()
123        else:
124            raise ValueError("Invalid security type")
125        self.__m.login(self.user, self.password)
126        self.select(self.__folder)
127
128    def disconnect(self):
129        """Disconnect from the IMAP server"""
130
131        log.info("Disconnecting from IMAP server")
132        self.__m.close()
133        self.__m.logout()
134
135    def __enter__(self):
136        self.connect()
137        return self
138
139    def __exit__(self, *args):
140        self.disconnect()
141
142    def __iter__(self):
143        """Iterate over all messages in the mailbox"""
144        data = handle_response(self.__m.search(None, "ALL"))
145        for uid in data[0].decode().split():
146            yield IMAPMessageHeadersOnly.from_uid(uid, self)
147
148    def values(self):
149        yield from iter(self)
150
151    def keys(self) -> list[str]:
152        """Get a list of all message UIDs in the mailbox"""
153        data = handle_response(self.__m.search(None, "ALL"))
154        return data[0].decode().split()
155
156    def items(self):
157        """Iterate over all messages in the mailbox"""
158        uids = ",".join(self.keys()).encode()
159        return self.fetch(uids, "RFC822")
160
161    @property
162    def capability(self):
163        """Get the server capabilities"""
164        return handle_response(self.__m.capability())[0].decode()
165
166    def add(self, message):
167        """Add a message to the mailbox"""
168
169        self.__m.append(
170            self.current_folder,
171            "",
172            imaplib.Time2Internaldate(time.time()),
173            message.as_bytes(),
174        )
175
176    def copy(self, messageset: bytes, folder: str) -> None:
177        """Copy a message to a different folder"""
178
179        self.__m.copy(messageset, folder)
180
181    def move(self, messageset: bytes, folder: str) -> None:
182        """Move a message to a different folder"""
183
184        self.__m._simple_command("MOVE", messageset, folder)
185
186    def discard(self, messageset: bytes) -> None:
187        """Mark messages for deletion"""
188
189        self.__m.store(messageset, "+FLAGS", "\\Deleted")
190
191    def remove(self, messageset: bytes) -> None:
192        """Remove messages from the mailbox"""
193
194        self.discard(messageset)
195        self.__m.expunge()
196
197    def __delitem__(self, key: str) -> None:
198        raise NotImplementedError("Use discard() instead")
199
200    def __len__(self) -> int:
201        return len(self.keys())
202
203    def fetch(self, messageset: bytes, what):
204        """Fetch messages from the mailbox"""
205
206        messages = handle_response(self.__m.fetch(messageset, what))[::2]
207
208        for head, body in messages:
209            uid, what, size = MESSAGE_HEAD_RE.match(head.decode()).groups()
210            if size != str(len(body)):
211                raise Exception("Size mismatch")
212
213            yield uid, body
214
215    def __expand_search_macros(self, query) -> str:
216        """Expand search macros in the query."""
217
218        today = datetime.date.today()
219        yesterday = today - datetime.timedelta(days=1)
220
221        week_start = today - datetime.timedelta(days=today.weekday())
222        last_week_start = week_start - datetime.timedelta(days=7)
223
224        month_start = datetime.date(today.year, today.month, 1)
225        year_start = datetime.date(today.year, 1, 1)
226
227        if today.month == 1:  # January
228            # last month is December of the previous year
229            last_month_start = datetime.date(today.year - 1, 12, 1)
230        else:
231            last_month_start = datetime.date(today.year, today.month - 1, 1)
232
233        last_year_start = datetime.date(today.year - 1, 1, 1)
234
235        q = query
236        q = q.replace("FIND", "TEXT")
237
238        q = q.replace("TODAY", f"ON {imap_date(today)}")
239        q = q.replace("YESTERDAY", f"ON {imap_date(yesterday)}")
240
241        q = q.replace("THISWEEK", f"SINCE {imap_date(week_start)}")
242        q = q.replace("THISMONTH", f"SINCE {imap_date(month_start)}")
243        q = q.replace("THISYEAR", f"SINCE {imap_date(year_start)}")
244
245        q = q.replace("LASTWEEK", imap_date_range(last_week_start, week_start))
246        q = q.replace("LASTMONTH", imap_date_range(last_month_start, month_start))
247        q = q.replace("LASTYEAR", imap_date_range(last_year_start, year_start))
248
249        # shortcuts
250        q = q.replace("PASTDAY", "PAST1DAY")
251        q = q.replace("PASTWEEK", "PAST1WEEK")
252        q = q.replace("PASTMONTH", "PAST1MONTH")
253        q = q.replace("PASTYEAR", "PAST1YEAR")
254
255        # use regex to match the PASTXDAYS macro
256        q = re.sub(
257            r"PAST(\d+)DAYS?",
258            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1))))}",
259            q,
260        )
261
262        # use regex to match the PASTXWEEKS macro
263        q = re.sub(
264            r"PAST(\d+)WEEKS?",
265            lambda m: f"SINCE {imap_date(change_time(today, weeks=-int(m.group(1))))}",
266            q,
267        )
268
269        # use regex to match the PASTXMONTHS macro
270        q = re.sub(
271            r"PAST(\d+)MONTHS?",
272            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1)) * 30))}",
273            q,
274        )
275
276        # use regex to match the PASTXYEARS macro
277        q = re.sub(
278            r"PAST(\d+)YEARS?",
279            lambda m: f"SINCE {imap_date(change_time(today, days=-int(m.group(1)) * 365))}",
280            q,
281        )
282
283        return q
284
285    def search(self, query):
286        """Search for messages matching the query
287
288        We support extra search macros in the search query in addition to
289        the standard IMAP search macros.
290
291        One search macro is FIND <text>, which is an alias for TEXT.
292        The rest of the macros deal with date ranges.
293
294        The date range macros are expanded to the appropriate date range and
295        are relative to the current date.
296        Example: TODAY expands to ON <date>, where <date> is today's date.
297
298        Note that some of these macros will expand to multiple search terms.
299        Expansions that result in multiple search terms are wrapped in parentheses.
300        Example: LASTWEEK expands to (SINCE <date1> BEFORE <date2>).
301
302        The following extra macros are supported:
303
304
305        - FIND <text> - alias for TEXT, searches the message headers and body
306
307        Current period:
308        - TODAY - messages from today
309        - THISWEEK - messages since the start of the week, Monday to Sunday
310        - THISMONTH - messages since the start of the month
311        - THISYEAR - messages since the start of the year
312
313        Previous period:
314        - YESTERDAY - messages from yesterday
315        - LASTWEEK - messages from the week before
316        - LASTMONTH - messages from the month before
317        - LASTYEAR - messages from the year before
318
319        Periods starting from now:
320
321        _These are just shortcuts_
322        - PASTDAY - messages from the past 1 day, same as PAST1DAY
323        - PASTWEEK - messages from the past 1 week, same as PAST1WEEK
324        - PASTMONTH - messages from the past 30 days, same as PAST1MONTH
325        - PASTYEAR - messages from the past 365 days, same as PAST1YEAR
326
327        _These are pattern matching macros_
328        - PASTXDAYS - messages from the past X days
329        - PASTXWEEKS - messages from the past X weeks
330        - PASTXMONTHS - messages from the past X * 30 days
331        - PASTXYEARS - messages from the past X * 365 days
332
333        These macros can be combined with other search macros, and can be
334        negated with NOT. For example, to search and archive or delete messages with a short
335        relevance period, you can use `NOT PAST3DAYS`, use `NOT PAST3MONTHS` to search for
336        messages older than a quarter, or use `NOT PAST2YEAR` to search for messages older than
337        two years.
338
339        _The `NOT` modifier is very useful for mailbox maintenance_
340
341        _There are no options for hours, because the range seletion does not have time of day precision._
342
343        Returns:
344            bytes: A comma-separated list of message UIDs
345        """
346
347        expanded_query = self.__expand_search_macros(query)
348        data = handle_response(self.__m.search(None, expanded_query))
349        num_results = len(data[0].split(b" "))
350
351        log.info(f"Searching for messages matching: {query}")
352        if expanded_query != query:
353            log.info(f"Expanded search query to: {expanded_query}")
354        log.info(f"Found {num_results} results")
355
356        return data[0].replace(b" ", b",")
357
358    def list_folders(self) -> tuple:
359        """List all folders in the mailbox
360
361        Returns:
362            tuple: A tuple of flags, delimiter, folder name, and folder display name
363        """
364
365        folders_data = handle_response(self.__m.list())
366        for data in folders_data:
367            flags, delimiter, folder = FOLDER_DATA_RE.match(data.decode()).groups()
368            display_name = folder.split(delimiter)[-1]
369            yield (flags, delimiter, folder, display_name)
370
371    @property
372    def current_folder(self):
373        """Get the currently selected folder"""
374        return self.__folder
375
376    def select(self, folder):
377        """Select a folder"""
378        self.__folder = folder
379        self.__m.select(folder)
380        return self