refactor: rewrite for IMAP support

- Replaced POP3 with IMAP
- #TODO: custom server URL
- Polish user prompts

3 files changed, 84 insertions, 61 deletions

diff --git a/bot.py b/bot.py
index b562768..c533d51 100644
--- a/bot.py
+++ b/bot.py
@@ -3,22 +3,26 @@ import os
 import sys
 from telegram import ParseMode, Update
 from telegram.constants import MAX_MESSAGE_LENGTH
-from telegram.ext import (Updater, CommandHandler, CallbackContext)
+from telegram.ext import Updater, CommandHandler, CallbackContext
 from utils.client import EmailClient
 
 
-logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s',
-                    stream=sys.stdout,
-                    level=logging.INFO)
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s",
+    stream=sys.stdout,
+    level=logging.INFO,
+)
 logger = logging.getLogger(__name__)
 
-bot_token = os.environ['TELEGRAM_TOKEN']
+bot_token = os.environ["TELEGRAM_TOKEN"]
+
+owner_chat_id = int(os.environ["OWNER_CHAT_ID"])
 
-owner_chat_id = int(os.environ['OWNER_CHAT_ID'])
 
 def is_owner(update: Update) -> bool:
     return update.message.chat_id == owner_chat_id
 
+
 def handle_large_text(text):
     while text:
         if len(text) < MAX_MESSAGE_LENGTH:
@@ -29,10 +33,12 @@ def handle_large_text(text):
             yield out
             text = text.lstrip(out)
 
+
 def error(update: Update, context: CallbackContext) -> None:
     """Log Errors caused by Updates."""
     logger.warning('Update "%s" caused error "%s"', update, context.error)
 
+
 def start_callback(update: Update, context: CallbackContext) -> None:
     if not is_owner(update):
         return
@@ -40,6 +46,7 @@ def start_callback(update: Update, context: CallbackContext) -> None:
     # print(update)
     update.message.reply_text(msg)
 
+
 def _help(update: Update, context: CallbackContext) -> None:
     if not is_owner(update):
         return
@@ -53,9 +60,12 @@ def _help(update: Update, context: CallbackContext) -> None:
     #            "/setting [email protected] password\n" \
     #            "/inbox\n" \
     #            "/get mail_index"
-    context.bot.send_message(update.message.chat_id, 
-                    # parse_mode=ParseMode.MARKDOWN,
-                    text=help_str)
+    context.bot.send_message(
+        update.message.chat_id,
+        # parse_mode=ParseMode.MARKDOWN,
+        text=help_str,
+    )
+
 
 def setting_email(update: Update, context: CallbackContext) -> None:
     if not is_owner(update):
@@ -67,7 +77,9 @@ def setting_email(update: Update, context: CallbackContext) -> None:
     update.message.reply_text("Configure email success!")
     with EmailClient(email_addr, email_passwd) as client:
         inbox_num = client.get_mails_count()
-    context.job_queue.run_repeating(periodic_task, interval=60, context=update.message.chat_id)
+    context.job_queue.run_repeating(
+        periodic_task, interval=60, context=update.message.chat_id
+    )
     # chat_data['job'] = job
     logger.info("periodic task scheduled.")
 
@@ -76,14 +88,15 @@ def periodic_task(context: CallbackContext) -> None:
     global inbox_num
     logger.info("entering periodic task.")
     with EmailClient(email_addr, email_passwd) as client:
-        new_inbox_num = client.get_mails_count()
-        if new_inbox_num > inbox_num:
-            mail = client.get_mail_by_index(new_inbox_num)
-            content = mail.__repr__()
-            for text in handle_large_text(content):
-                context.bot.send_message(context.job.context,
-                                text=text)
-            inbox_num = new_inbox_num
+        new_unseen_num = client.get_mails_unseen_count()
+        if new_unseen_num:
+            for i in range(inbox_num, inbox_num + new_unseen_num + 1):
+                mail = client.get_mail_by_index(i)
+                content = mail.__repr__()
+                for text in handle_large_text(content):
+                    context.bot.send_message(context.job.context, text=text)
+            inbox_num += new_unseen_num
+
 
 def inbox(update: Update, context: CallbackContext) -> None:
     if not is_owner(update):
@@ -92,26 +105,32 @@ def inbox(update: Update, context: CallbackContext) -> None:
     with EmailClient(email_addr, email_passwd) as client:
         global inbox_num
         new_num = client.get_mails_count()
-        reply_text = "The index of newest mail is *%d*," \
-                     " received *%d* new mails since last" \
-                     " time you checked." % \
-                     (new_num, new_num - inbox_num)
+        reply_text = (
+            "The index of newest mail is *%d*,"
+            " received *%d* new mails since last"
+            " time you checked." % (new_num, new_num - inbox_num)
+        )
         inbox_num = new_num
-        context.bot.send_message(update.message.chat_id,
-                        parse_mode=ParseMode.MARKDOWN,
-                        text=reply_text)
+        context.bot.send_message(
+            update.message.chat_id, parse_mode=ParseMode.MARKDOWN, text=reply_text
+        )
+
 
 def get_email(update: Update, context: CallbackContext) -> None:
     if not is_owner(update):
         return
     index = context.args[0]
+    if not index:
+        context.bot.send_message(
+            update.message.chat_id, text="$index should be a positive number!"
+        )
     logger.info("received get command.")
     with EmailClient(email_addr, email_passwd) as client:
         mail = client.get_mail_by_index(index)
         content = mail.__repr__()
         for text in handle_large_text(content):
-            context.bot.send_message(update.message.chat_id,
-                             text=text)
+            context.bot.send_message(update.message.chat_id, text=text)
+
 
 def main():
     # Create the EventHandler and pass it your bot's token.
@@ -133,7 +152,6 @@ def main():
 
     dp.add_handler(CommandHandler("get", get_email))
 
-
     dp.add_error_handler(error)
 
     # Start the Bot
@@ -145,5 +163,5 @@ def main():
     updater.idle()
 
 
-if __name__ == '__main__':
-    main()
\ No newline at end of file
+if __name__ == "__main__":
+    main()
diff --git a/utils/client.py b/utils/client.py
index 1dc9205..884a922 100644
--- a/utils/client.py
+++ b/utils/client.py
@@ -1,63 +1,66 @@
 import logging
-import poplib
+from imapclient import IMAPClient
 from utils.mail import Email
 
 
 logger = logging.getLogger(__name__)
 
 
-
 class EmailClient(object):
     def __init__(self, email_account, passwd):
         self.email_account = email_account
         self.password = passwd
-        self.server = self.connect(self)
+        self.server, self.number = self.connect(self)
 
     @staticmethod
     def connect(self):
         # parse the server's hostname from email account
-        pop3_server = 'pop.'+self.email_account.split('@')[-1]
-        server = poplib.POP3_SSL(pop3_server)
+        imap4_server = "mx1." + self.email_account.split("@")[-1]
+        server = IMAPClient(host=imap4_server, use_uid=True)
         # display the welcome info received from server,
         # indicating the connection is set up properly
-        logger.info(server.getwelcome().decode('utf8'))
+        logger.info(server.welcome)
         # authenticating
-        server.user(self.email_account)
-        server.pass_(self.password)
-        return server
+        server.login(self.email_account, self.password)
+        return server, server.select_folder("INBOX")[b"EXISTS"]
 
     def get_mails_list(self):
-        _, mails, _ = self.server.list()
-        return mails
+        messages = self.server.search()
+        dict = self.server.fetch(messages, "RFC822")
+        return sorted(dict.items(), key=lambda e: e[0])  # Do resorts
 
     def get_mails_count(self):
-        mails = self.get_mails_list()
-        return len(mails)
+        return self.number
+
+    def get_mails_unseen_count(self):
+        unseen_messages = self.server.search("UNSEEN")
+        return len(unseen_messages)
 
     def get_mail_by_index(self, index):
-        resp_status, mail_lines, mail_octets = self.server.retr(index)
-        return Email(mail_lines)
+        list = self.get_mails_list()
+        msg_data = list[int(index)-1][1]
+        return Email(msg_data[b"RFC822"])
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is None:
-            logger.info('exited normally\n')
-            self.server.quit()
+            logger.info("exited normally\n")
+            self.server.shutdown()
         else:
-            logger.error('raise an exception! ' + str(exc_type))
-            self.server.close()
-            return False # Propagate
+            logger.error("raise an exception! " + str(exc_type))
+            self.server.logout()
+            return False  # Propagate
 
 
-
-if __name__ == '__main__':
+if __name__ == "__main__":
     useraccount = "XXXXX"
     password = "XXXXXX"
 
     client = EmailClient(useraccount, password)
     num = client.get_mails_count()
     print(num)
-    for i in range(1, num):
-        print(client.get_mail_by_index(i))
\ No newline at end of file
+
+    for uid, msg_data in client.get_mails_list().items():
+        print(client.get_mail_by_index(uid))
diff --git a/utils/mail.py b/utils/mail.py
index 491faa8..3df068b 100644
--- a/utils/mail.py
+++ b/utils/mail.py
@@ -1,18 +1,20 @@
 from pyzmail import PyzMessage, decode_text
 
+
 class Email(object):
     def __init__(self, raw_mail_lines):
-        msg_content = b'\r\n'.join(raw_mail_lines)
-        msg =  PyzMessage.factory(msg_content)
+        msg = PyzMessage.factory(raw_mail_lines)
 
         self.subject = msg.get_subject()
-        self.sender = msg.get_address('from')
-        self.date = msg.get_decoded_header('date', '')
-        self.id = msg.get_decoded_header('message-id', '')
+        self.sender = msg.get_address("from")
+        self.date = msg.get_decoded_header("date", "")
+        self.id = msg.get_decoded_header("message-id", "")
 
         for mailpart in msg.mailparts:
-            if mailpart.is_body=='text/plain':
-                payload, used_charset=decode_text(mailpart.get_payload(), mailpart.charset, None)
+            if mailpart.is_body == "text/plain":
+                payload, used_charset = decode_text(
+                    mailpart.get_payload(), mailpart.charset, None
+                )
                 self.charset = used_charset
                 self.text = payload
                 return