From 6a27f41de1c1f6bc9bd2cf9b8eca62ed6653587e Mon Sep 17 00:00:00 2001
From: Sean Whalen <whalenster@gmail.com>
Date: Sun, 30 Jun 2019 19:39:57 -0400
Subject: [PATCH] Add IMAP message download timeout

---
 parsedmarc/__init__.py | 27 ++++++++++++++++++---------
 requirements.txt       |  1 +
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py
index 8656ba0..c466485 100644
--- a/parsedmarc/__init__.py
+++ b/parsedmarc/__init__.py
@@ -32,6 +32,7 @@ import xmltodict
 import imapclient
 import imapclient.exceptions
 import mailparser
+import timeout_decorator
 
 from parsedmarc.utils import get_base_domain, get_ip_address_info
 from parsedmarc.utils import is_outlook_msg, convert_outlook_msg
@@ -948,6 +949,20 @@ def get_dmarc_reports_from_inbox(host=None,
             server_capabilities = get_imap_capabilities(server)
             move_supported = "MOVE" in server_capabilities
 
+        @timeout_decorator.timeout(5)
+        def fetch_message(msg_uid):
+            raw_msg = server.fetch(msg_uid,
+                                   ["RFC822"])[msg_uid]
+            msg_keys = [b'RFC822', b'BODY[NULL]', b'BODY[]']
+            msg_key = ''
+            for key in msg_keys:
+                if key in raw_msg.keys():
+                    msg_key = key
+                    break
+            raw_msg = raw_msg[msg_key]
+
+            return raw_msg
+
         def delete_messages(msg_uids):
             logger.debug("Deleting message UID(s) {0}".format(",".join(
                 str(uid) for uid in msg_uids)))
@@ -1028,16 +1043,10 @@ def get_dmarc_reports_from_inbox(host=None,
             ))
             try:
                 try:
-                    raw_msg = server.fetch(msg_uid,
-                                           ["RFC822"])[msg_uid]
-                    msg_keys = [b'RFC822', b'BODY[NULL]', b'BODY[]']
-                    msg_key = ''
-                    for key in msg_keys:
-                        if key in raw_msg.keys():
-                            msg_key = key
-                            break
-                    raw_msg = raw_msg[msg_key]
+                    raw_msg = fetch_message(msg_uid)
 
+                except timeout_decorator.timeout_decorator.TimeoutError:
+                    logger.debug("Download timed out. Retrying...")
                 except (ConnectionResetError, socket.error,
                         TimeoutError,
                         imapclient.exceptions.IMAPClientError) as error:
diff --git a/requirements.txt b/requirements.txt
index f1d9df2..de9ac0b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,6 +14,7 @@ dateparser>=0.7.1
 elasticsearch>=6.3.1,<7.0.0
 elasticsearch-dsl>=6.3.1,<7.0.0
 kafka-python>=1.4.4
+timeout-decorator>=0.4.1
 nose
 pygments
 flake8