import email
import imaplib
import sys, os
import re
import datetime


def get_all_attach(username, password, server):
    detach_dir = '.'
    if 'attachments' not in os.listdir(detach_dir):
        os.mkdir('attachments')

    mail = imaplib.IMAP4_SSL(server)
    typ, accountDetails = mail.login(username, password)
    try:
        if typ != 'OK':
            raise Exception("Unable to sign-in")
            
        mail.select('INBOX')
        result, data = mail.uid('search', None, 'UNSEEN')
        mail_ids = data[0]
        if not mail_ids:
            raise Exception("no unread messages")
        else:
            msg_ids = str(mail_ids, encoding="utf-8").split()

            # this is where the message iteration begins
            
            for em_id in msg_ids:
                result, data = mail.uid('fetch', em_id, '(RFC822)')
                raw_email = data[0][1]
                email_message = email.message_from_bytes(raw_email)
                

                if '<' not in email_message['From']:
                    sender = email_message['From']
                else:
                    sender_pat = re.compile(r'^.*<([\w\.-]+@[\w\.-]+)>.*$')
                    received_from = sender_pat.match(email_message['From'])
                    try:
                        sender = received_from.groups()[0] 
                    except AttributeError:
                        sender = "unknown"

                if "@" in sender:
                    sender = sender.split('@')[0]
                else:
                    pass

         
                for part in email_message.walk():
                    if part.get_content_maintype() == 'multipart':
                        continue
                    if part.get('Content-Disposition') is None:
                        continue
                    filename = part.get_filename()
                    got_attach = False
                    if filename:
                        # We will need to test is filename exists and if so we add a timestamp to it which should make it unique
                        now = str(datetime.datetime.now())[5:] #strips off the year
                        timestamp = now.replace(':', '_').replace('.', '-').replace(' ', '_')
                        timestamp += "_"
                        filename = sender + '_' + timestamp + filename
                        got_attach = True
                        filePath = os.path.join(detach_dir, 'attachments', filename)
                        if os.path.isfile(filePath):
                            filename = "B-" + filename
                            fp = open(filePath, 'wb')
                            fp.write(part.get_payload(decode=True))
                            fp.close()
                        else:
                            fp = open(filePath, 'wb')
                            fp.write(part.get_payload(decode=True))
                            fp.close()

                if got_attach:
                     mail.uid('store', em_id, '+FLAGS', '\Seen')
                            
        mail.close()
        mail.logout()
        return True

    except:
        print('An error occurred when attempting to download attachments')
        mail.close()
        mail.logout()
        return False


if __name__ == '__main__':

    server = #server-name
    user = #email-username
    password = #emal-password
    test_case = get_all_attach(user, password, server)
    if test_case:
        print("All attachments should have been successfully downloaded. \nCheck your attachments folder for all attachments.\n\n")
    else:
        print("There was an error. Check the inbox and see if there are any unread messages\n\n")

