email-organizer/scripts/populate_imap.py

#!/usr/bin/env python3
"""
Script to automatically populate the Dovecot IMAP server with test emails for both users.
This script will create various emails in different folders for testing purposes.
"""

import os
import sys
import email
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.utils import formatdate, make_msgid
import imaplib
import time

# Add the project root to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/..')

from app import create_app, db
from app.models import User

def create_test_emails():
    """Create various test emails for different folders"""
    
    # Define email templates with different content types and folders
    emails = [
        # INBOX emails
        {
            "subject": "Welcome to Our Service",
            "from": "noreply@company.com",
            "to": "user1@example.com",
            "folder": "INBOX",
            "body": "Thank you for signing up! Here's your welcome package.",
            "priority": "normal"
        },
        {
            "subject": "Meeting Reminder - Project Alpha",
            "from": "manager@company.com",
            "to": "user1@example.com",
            "folder": "INBOX",
            "body": "Don't forget about our meeting tomorrow at 10am.",
            "priority": "high"
        },
        {
            "subject": "Weekly Newsletter",
            "from": "newsletter@company.com",
            "to": "user2@example.com",
            "folder": "INBOX",
            "body": "Your weekly digest of industry news and updates.",
            "priority": "normal"
        },
        {
            "subject": "Password Reset Request",
            "from": "security@company.com",
            "to": "user2@example.com",
            "folder": "INBOX",
            "body": "You requested a password reset. Click the link below to reset your password.",
            "priority": "high"
        },
        
        # Work emails
        {
            "subject": "Q3 Budget Review",
            "from": "finance@company.com",
            "to": "user1@example.com",
            "folder": "Work",
            "body": "Please review the Q3 budget proposal attached.",
            "priority": "high"
        },
        {
            "subject": "Project Status Update",
            "from": "teamlead@company.com",
            "to": "user1@example.com",
            "folder": "Work",
            "body": "Here's the latest status on our current sprint.",
            "priority": "normal"
        },
        
        # Personal emails
        {
            "subject": "Dinner Party This Weekend",
            "from": "friend@friends.com",
            "to": "user1@example.com",
            "folder": "Personal",
            "body": "Are you free for dinner this Saturday? Let me know!",
            "priority": "normal"
        },
        {
            "subject": "Vacation Plans Discussion",
            "from": "family@family.com",
            "to": "user2@example.com",
            "folder": "Personal",
            "body": "Let's discuss our vacation plans for next month.",
            "priority": "normal"
        },
        
        # Receipts
        {
            "subject": "Amazon Order Confirmation",
            "from": "orders@amazon.com",
            "to": "user1@example.com",
            "folder": "Receipts",
            "body": "Your order has been shipped and will arrive by Friday.",
            "priority": "normal"
        },
        {
            "subject": "Gas Station Receipt",
            "from": "receipts@shell.com",
            "to": "user2@example.com",
            "folder": "Receipts",
            "body": "Thank you for your purchase at Shell station.",
            "priority": "normal"
        },
        
        # Marketing emails
        {
            "subject": "Special Offer - 20% Off Today Only!",
            "from": "deals@retailer.com",
            "to": "user1@example.com",
            "folder": "Marketing",
            "body": "Limited time offer! Get 20% off your next purchase.",
            "priority": "normal"
        },
        {
            "subject": "New Product Launch",
            "from": "news@techcompany.com",
            "to": "user2@example.com",
            "folder": "Marketing",
            "body": "We're excited to announce our latest product launch!",
            "priority": "normal"
        },
        
        # Pending folder (for testing)
        {
            "subject": "Pending Review - Contract Document",
            "from": "legal@company.com",
            "to": "user1@example.com",
            "folder": "Pending",
            "body": "Please review this contract document and provide feedback.",
            "priority": "high"
        },
        {
            "subject": "New Feature Request",
            "from": "support@product.com",
            "to": "user2@example.com",
            "folder": "Pending",
            "body": "We've received a new feature request from a customer.",
            "priority": "normal"
        },
        
        # Archived folder
        {
            "subject": "Old Meeting Notes - Q2",
            "from": "archive@company.com",
            "to": "user1@example.com",
            "folder": "Archived",
            "body": "Archived meeting notes from the previous quarter.",
            "priority": "normal"
        },
        {
            "subject": "Previous Newsletter",
            "from": "old-newsletter@company.com",
            "to": "user2@example.com",
            "folder": "Archived",
            "body": "Previous edition of our newsletter.",
            "priority": "normal"
        }
    ]
    
    return emails

def create_email_message(subject, sender, recipient, body, msg_id=None):
    """Create an email message with proper headers"""
    msg = MIMEMultipart()
    msg['From'] = sender
    msg['To'] = recipient
    msg['Date'] = formatdate(localtime=True)
    msg['Message-ID'] = msg_id or make_msgid()
    msg['Subject'] = subject
    
    # Add body to email
    msg.attach(MIMEText(body, 'plain'))
    
    return msg

def connect_to_imap():
    """Connect to the local IMAP server"""
    try:
        # Connect to the local IMAP server (running on port 5143)
        mail = imaplib.IMAP4('localhost', 5143)
        return mail
    except Exception as e:
        print(f"Failed to connect to IMAP server: {e}")
        return None

def login_to_imap(mail, username, password):
    """Login to the IMAP server"""
    try:
        mail.login(username, password)
        return True
    except Exception as e:
        print(f"Failed to login to IMAP server for {username}: {e}")
        return False

def create_mailbox_if_not_exists(mail, folder_name):
    """Create a mailbox if it doesn't exist"""
    try:
        # Check if mailbox exists
        status, messages = mail.list("", f'"{folder_name}"')
        if status != 'OK':
            # Create the mailbox
            mail.create(f'"{folder_name}"')
            print(f"Created mailbox: {folder_name}")
        return True
    except Exception as e:
        print(f"Error creating mailbox {folder_name}: {e}")
        return False

def store_email_in_folder(mail, folder_name, email_message):
    """Store an email in a specific folder"""
    try:
        # Select the folder
        mail.select(f'"{folder_name}"')
        
        # Store the email using APPEND
        raw_email = email_message.as_bytes()
        mail.append(f'"{folder_name}"', '', imaplib.Time2Internaldate(time.time()), raw_email)
        
        return True
    except Exception as e:
        print(f"Error storing email in {folder_name}: {e}")
        return False

def populate_imap():
    """Main function to populate IMAP server with test emails"""
    print("Populating IMAP server with test emails...")
    
    # Create the test emails
    test_emails = create_test_emails()
    
    # Connect to IMAP server
    mail = connect_to_imap()
    if not mail:
        print("Could not connect to IMAP server")
        return False
    
    # Login with both users
    users = [
        {"email": "user1@example.com", "password": "password1"},
        {"email": "user2@example.com", "password": "password2"}
    ]
    
    for user in users:
        print(f"Processing emails for {user['email']}")
        
        # Login to the IMAP server
        if not login_to_imap(mail, user['email'], user['password']):
            continue
            
        # Create mailboxes if they don't exist
        folders = ["INBOX", "Pending", "Work", "Personal", "Receipts", "Marketing", "Archived"]
        for folder in folders:
            create_mailbox_if_not_exists(mail, folder)
        
        # Store emails for this user
        user_emails = [e for e in test_emails if e["to"] == user["email"]]
        for email_data in user_emails:
            print(f"  Storing email: {email_data['subject']} in {email_data['folder']}")
            
            # Create the email message
            msg = create_email_message(
                subject=email_data["subject"],
                sender=email_data["from"],
                recipient=email_data["to"],
                body=email_data["body"]
            )
            
            # Store in the appropriate folder
            store_email_in_folder(mail, email_data["folder"], msg)
    
    # Logout
    mail.logout()
    print("IMAP server populated with test emails successfully!")
    return True

if __name__ == "__main__":
    populate_imap()