1import logging
2import time
3import socket
4from ssl import CERT_NONE, SSLError, CertificateError, create_default_context
5
6import imapclient
7import imapclient.exceptions
8
9import mailsuite.utils
10
11logger = logging.getLogger(__name__)
12
13
14class MaxRetriesExceeded(RuntimeError):
15    """Raised when the maximum number of retries in exceeded"""
16
17
18def _chunks(list_like_object, n):
19    """Yield successive n-sized chunks from l."""
20    for i in range(0, len(list_like_object), n):
21        yield list_like_object[i:i + n]
22
23
24class IMAPClient(imapclient.IMAPClient):
25    """A simplified IMAP client"""
26
27    def _normalise_folder(self, folder_name):
28        """
29        Returns an appropriate path based on the namespace (if any) and
30        hierarchy separator
31
32        Args:
33            folder_name (str): The path to correct
34
35        Returns:
36            str: A corrected path
37        """
38        if folder_name in ["", "*", "INBOX"]:
39            return imapclient.IMAPClient._normalise_folder(self, folder_name)
40        folder_name = folder_name.rstrip("/")
41        folder_name = folder_name.replace(self._path_prefix, "")
42        if not self._hierarchy_separator == "/":
43            folder_name = folder_name.replace(self._hierarchy_separator, "")
44            folder_name = folder_name.replace("/", self._hierarchy_separator)
45        folder_name = "{0}{1}".format(self._path_prefix, folder_name)
46
47        return imapclient.IMAPClient._normalise_folder(self, folder_name)
48
49    def _start_idle(self, idle_callback, idle_timeout=30):
50        """
51        Starts an IMAP IDLE session
52
53        Args:
54            idle_callback (function: A callback function
55            idle_timeout (int): Number of seconds to wait for an IDLE response
56        """
57        if self._idle_supported is False:
58            raise imapclient.exceptions.IMAPClientError(
59                "IDLE is not supported by the server")
60        idle_callback(self)
61        idle_start_time = time.monotonic()
62        self.idle()
63        while True:
64            try:
65                # Refresh the IDLE session every 5 minutes to stay connected
66                if time.monotonic() - idle_start_time > 5 * 60:
67                    logger.info("IMAP: Refreshing IDLE session")
68                    self.idle_done()
69                    idle_start_time = time.monotonic()
70                    self.idle(self)
71                responses = self.idle_check(timeout=idle_timeout)
72                if responses is not None:
73                    if len(responses) == 0:
74                        # Gmail/G-Suite returns an empty list
75                        self.idle_done()
76                        idle_callback(self)
77                        idle_start_time = time.monotonic()
78                        self.idle()
79                    else:
80                        for r in responses:
81                            if r[0] != 0 and r[1] == b'RECENT':
82                                self.idle_done()
83                                idle_callback(self)
84                                idle_start_time = time.monotonic()
85                                self.idle()
86                                break
87            except (KeyError, socket.error, BrokenPipeError,
88                    ConnectionResetError):
89                logger.debug("IMAP error: Connection reset")
90                self.reset_connection()
91            except imapclient.exceptions.IMAPClientError as error:
92                error = error.__str__().lstrip("b'").rstrip("'").rstrip(".")
93                # Workaround for random Exchange/Office365 IMAP errors
94                if "unexpected response" in error or "BAD" in error:
95                    self.reset_connection()
96            except KeyboardInterrupt:
97                break
98        try:
99            self.idle_done()
100        except BrokenPipeError:
101            pass
102
103    def __init__(self, host, username, password, port=None, ssl=True,
104                 verify=True, timeout=30, max_retries=4,
105                 initial_folder="INBOX", idle_callback=None, idle_timeout=30):
106        """
107        Connects to an IMAP server
108
109        Args:
110            host (str): The server hostname or IP address
111            username (str): The username
112            password (str): The password
113            port (int): The port
114            ssl (bool): Use SSL or TLS
115            verify (bool): Verify the SSL/TLS certificate
116            timeout (float): Number of seconds to wait for an operation
117            max_retries (int): The maximum number of retries after a timeout
118            initial_folder (str): The initial folder to select
119            idle_callback: The function to call when new messages are detected
120            idle_timeout (int): Number of seconds to wait for an IDLE
121                                  response
122        """
123        ssl_context = create_default_context()
124        if verify is False:
125            ssl_context.check_hostname = False
126            ssl_context.verify_mode = CERT_NONE
127        self._init_args = dict(host=host, username=username,
128                               password=password, port=port, ssl=ssl,
129                               verify=verify,
130                               timeout=timeout,
131                               max_retries=max_retries,
132                               initial_folder=initial_folder,
133                               idle_callback=idle_callback,
134                               idle_timeout=idle_timeout)
135        self.max_retries = max_retries
136        self.idle_callback = idle_callback
137        self.idle_timeout = idle_timeout
138        self._path_prefix = ""
139        self._hierarchy_separator = ""
140        if not ssl:
141            logger.info("Connecting to IMAP over plain text")
142        imapclient.IMAPClient.__init__(self,
143                                       host=host,
144                                       port=port,
145                                       ssl=ssl,
146                                       ssl_context=ssl_context,
147                                       use_uid=True,
148                                       timeout=timeout)
149        try:
150            self.login(username, password)
151            self.server_capabilities = self.capabilities()
152            self._move_supported = b"MOVE" in self.server_capabilities
153            self._idle_supported = b"IDLE" in self.server_capabilities
154            self._namespace = b"NAMESPACE" in self.server_capabilities
155            self._hierarchy_separator = self.list_folders()[0][1]
156            if type(self._hierarchy_separator == bytes):
157                self._hierarchy_separator = self._hierarchy_separator.decode(
158                    "utf-8")
159            if self._namespace:
160                self._namespace = self.namespace()
161                personal_namespace = self._namespace.personal
162                if len(personal_namespace) > 0:
163                    self._hierarchy_separator = personal_namespace[0][1]
164                    if not personal_namespace[0][0] == "":
165                        self._path_prefix = personal_namespace[0][0]
166                        if type(self._path_prefix) == bytes:
167                            self._path_prefix = self._path_prefix.decode(
168                                "utf-8")
169            else:
170                self._namespace = None
171            self.select_folder(initial_folder)
172        except (ConnectionResetError, socket.error,
173                TimeoutError,
174                imapclient.exceptions.IMAPClientError) as error:
175            error = error.__str__().lstrip("b'").rstrip("'").rstrip(
176                ".")
177            raise imapclient.exceptions.IMAPClientError(error)
178        except ConnectionAbortedError:
179            raise imapclient.exceptions.IMAPClientError("Connection aborted")
180        except TimeoutError:
181            raise imapclient.exceptions.IMAPClientError("Connection timed out")
182        except SSLError as error:
183            raise imapclient.exceptions.IMAPClientError(
184                "SSL error: {0}".format(error.__str__()))
185        except CertificateError as error:
186            raise imapclient.exceptions.IMAPClientError(
187                "Certificate error: {0}".format(error.__str__()))
188        except BrokenPipeError:
189            raise imapclient.exceptions.IMAPClientError("Broken pipe")
190
191        if idle_callback is not None:
192            self._start_idle(idle_callback, idle_timeout=idle_timeout)
193
194    def reset_connection(self):
195        """Resets the connection to the IMAP server"""
196        logger.info("Reconnecting to IMAP")
197        try:
198            self.shutdown()
199        except Exception as e:
200            logger.info(
201                "Failed to log out: {0}".format(e.__str__()))
202        self.__init__(self._init_args["host"],
203                      self._init_args["username"],
204                      self._init_args["password"],
205                      port=self._init_args["port"],
206                      ssl=self._init_args["ssl"],
207                      verify=self._init_args["verify"],
208                      timeout=self._init_args["timeout"],
209                      max_retries=self._init_args["max_retries"],
210                      initial_folder=self._init_args["initial_folder"],
211                      idle_callback=self._init_args["idle_callback"],
212                      idle_timeout=self._init_args["idle_timeout"],
213                      )
214
215    def fetch_message(self, msg_uid, parse=False, _attempt=1):
216        """
217        Fetch a message by UID, and optionally parse it
218
219        Args:
220            msg_uid (int): The message UID
221            parse (bool): Return parsed results from mailparser
222            _attempt (int): The attempt number
223
224        Returns:
225            str: The raw mail message, including headers
226            dict: A parsed email message
227        """
228        try:
229            raw_msg = self.fetch(msg_uid, ["RFC822"])[msg_uid]
230        except socket.timeout:
231            _attempt = _attempt + 1
232            if _attempt > self.max_retries:
233                raise MaxRetriesExceeded("Maximum retries exceeded")
234            logger.info("Attempt {0} of {1} timed out. Retrying...".format(
235                _attempt,
236                self.max_retries))
237            self.reset_connection()
238            return self.fetch_message(msg_uid, parse=parse, _attempt=_attempt)
239        msg_keys = [b'RFC822', b'BODY[NULL]', b'BODY[]']
240        msg_key = ''
241        for key in msg_keys:
242            if key in raw_msg.keys():
243                msg_key = key
244                break
245        message = raw_msg[msg_key].decode("utf-8", "replace")
246        if parse:
247            message = mailsuite.utils.parse_email(message)
248        return message
249
250    def delete_messages(self, msg_uids, silent=True, _attempt=1):
251        """
252        Deletes the given messages by Message UIDs
253
254        Args:
255            msg_uids (list): A list of UIDs of messages to delete
256            silent (bool): Do it silently
257            _attempt (int): The attempt number
258        """
259        logger.info("Deleting message UID(s) {0}".format(",".join(
260            str(uid) for uid in msg_uids)))
261        if type(msg_uids) == str or type(msg_uids) == int:
262            msg_uids = [int(msg_uids)]
263        try:
264            imapclient.IMAPClient.delete_messages(self, msg_uids,
265                                                  silent=silent)
266            imapclient.IMAPClient.expunge(self, msg_uids)
267        except socket.timeout:
268            _attempt = _attempt + 1
269            if _attempt > self.max_retries:
270                raise MaxRetriesExceeded("Maximum retries exceeded")
271            logger.info("Attempt {0} of {1} timed out. Retrying...".format(
272                _attempt,
273                self.max_retries))
274            self.reset_connection()
275            self.delete_messages(msg_uids, silent=silent, _attempt=_attempt)
276
277    def create_folder(self, folder_path, _attempt=1):
278        """
279        Creates an IMAP folder at the given path
280
281        Args:
282            folder_path (str): The path of the folder to create
283            _attempt (int): The attempt number
284        """
285        if not self.folder_exists(folder_path):
286            logger.info("Creating folder: {0}".format(folder_path))
287            try:
288                imapclient.IMAPClient.create_folder(self, folder_path)
289            except socket.timeout:
290                _attempt = _attempt + 1
291                if _attempt > self.max_retries:
292                    raise MaxRetriesExceeded("Maximum retries exceeded")
293                logger.info("Attempt {0} of {1} timed out. Retrying...".format(
294                    _attempt,
295                    self.max_retries))
296                self.reset_connection()
297                self.create_folder(folder_path, _attempt=_attempt)
298
299    def _move_messages(self, msg_uids, folder_path):
300        """
301        Move the emails with the given UIDs to the given folder
302
303        Args:
304            msg_uids: A UID or list of UIDs of messages to move
305            folder_path (str): The path of the destination folder
306        """
307        folder_path = folder_path.replace("\\", "/").rstrip("/")
308        if type(msg_uids) == str or type(msg_uids) == int:
309            msg_uids = [int(msg_uids)]
310        for chunk in _chunks(msg_uids, 100):
311            if self._move_supported:
312                logger.info("Moving message UID(s) {0} to {1}".format(
313                    ",".join(str(uid) for uid in chunk), folder_path
314                ))
315                try:
316                    self.move(chunk, folder_path)
317                except imapclient.exceptions.IMAPClientError as e:
318                    e = e.__str__().lstrip("b'").rstrip(
319                        "'").rstrip(".")
320                    message = "Error moving message UIDs"
321                    e = "{0} {1}: " "{2}".format(message, msg_uids, e)
322                    logger.info("IMAP error: {0}".format(e))
323                    logger.info(
324                        "Copying message UID(s) {0} to {1} by copy".format(
325                            ",".join(str(uid) for uid in chunk), folder_path
326                        ))
327                    self.copy(msg_uids, folder_path)
328                    self.delete_messages(msg_uids)
329            else:
330                logger.info("Moving message UID(s) {0} to {1} by copy".format(
331                    ",".join(str(uid) for uid in chunk), folder_path
332                ))
333                self.copy(msg_uids, folder_path)
334                self.delete_messages(msg_uids)
335
336    def move_messages(self, msg_uids, folder_path, _attempt=1):
337        """
338        Move the emails with the given UIDs to the given folder
339
340        Args:
341            msg_uids: A UID or list of UIDs of messages to move
342            folder_path (str): The path of the destination folder
343            _attempt (int): The attempt number
344        """
345        try:
346            self._move_messages(msg_uids, folder_path)
347        except socket.timeout:
348            _attempt = _attempt + 1
349            if _attempt > self.max_retries:
350                raise MaxRetriesExceeded("Maximum retries exceeded")
351            logger.info("Attempt {0} of {1} timed out. Retrying...".format(
352                _attempt,
353                self.max_retries))
354            self.reset_connection()
355            self._move_messages(msg_uids, folder_path)
356