fetch_and_test_emails.py

"""
Fetch emails using IMAP and test them with your AI phishing detection model.
Works with Gmail, Outlook, Yahoo, etc. (IMAP must be enabled, use app password if 2FA is on)
"""
import imaplib
import email
from email.header import decode_header
import getpass
import string
import joblib

model = joblib.load('phishing_model.joblib')
vectorizer = joblib.load('vectorizer.joblib')

def clean_text(text):
    if not isinstance(text, str):
        return ""
    text = text.lower()
    text = text.translate(str.maketrans('', '', string.punctuation))
    text = text.strip()
    return text

def predict_phishing(message):
    message = clean_text(message)
    message_vector = vectorizer.transform([message])
    prediction = model.predict(message_vector)
    return "⚠️ Phishing detected" if prediction[0] == 1 else "✅ Safe message"

# --- IMAP Email Fetching ---

print("Enter your email credentials (use app password if needed)")
EMAIL = input("Email: ")
PASSWORD = getpass.getpass("Password: ")
IMAP_SERVER = input("IMAP server (e.g. imap.gmail.com): ")

mail = imaplib.IMAP4_SSL(IMAP_SERVER)
mail.login(EMAIL, PASSWORD)

def fetch_and_predict(folder_name, label):
    mail.select(folder_name)
    status, messages = mail.search(None, "ALL")
    email_ids = messages[0].split()
    print(f"\n--- {label} ({len(email_ids)} emails, showing latest 5) ---")
    for i in email_ids[-5:]:
        res, msg = mail.fetch(i, "(RFC822)")
        for response in msg:
            if isinstance(response, tuple):
                msg = email.message_from_bytes(response[1])
                subject, encoding = decode_header(msg["Subject"])[0]
                if isinstance(subject, bytes):
                    subject = subject.decode(encoding if encoding else "utf-8")
                body = ""
                if msg.is_multipart():
                    for part in msg.walk():
                        ctype = part.get_content_type()
                        cdispo = str(part.get('Content-Disposition'))
                        if ctype == 'text/plain' and 'attachment' not in cdispo:
                            body = part.get_payload(decode=True).decode('utf-8', errors='ignore')
                            break
                else:
                    body = msg.get_payload(decode=True).decode('utf-8', errors='ignore')
                print(f"\nSubject: {subject}")
                print(f"Prediction: {predict_phishing(body)}")

# Check inbox
fetch_and_predict("inbox", "Inbox")
# Check spam/junk (Gmail: [Gmail]/Spam, Outlook: Junk)
try:
    fetch_and_predict("[Gmail]/Spam", "Spam")
except:
    try:
        fetch_and_predict("Junk", "Junk")
    except:
        print("No spam/junk folder found or accessible.")

mail.logout()