PhishGuardian/backend/app.py

96 lines
3.1 KiB
Python

from flask import Flask, request, jsonify
from flask_cors import CORS
import imaplib
import email
from email.header import decode_header
import joblib
app = Flask(__name__)
CORS(app)
model = joblib.load('spam_classifier_model.pkl')
vectorizer = joblib.load('vectorizer.pkl')
@app.route('/fetch-emails', methods=['POST'])
def fetch_emails():
data = request.json
username = data['username']
password = data['password']
try:
mail = imaplib.IMAP4_SSL("outlook.office365.com")
mail.login(username, password)
mail.select("inbox")
except imaplib.IMAP4.error:
return jsonify({"error": "Login failed. Check your email and password."}), 401
status, messages = mail.search(None, "ALL")
email_ids = messages[0].split()
emails = []
for email_id in email_ids:
res, msg = mail.fetch(email_id, "(RFC822)")
for response_part in msg:
if isinstance(response_part, tuple):
msg = email.message_from_bytes(response_part[1])
subject, encoding = decode_header(msg["Subject"])[0]
if isinstance(subject, bytes):
subject = subject.decode(encoding if encoding else "utf-8")
from_ = msg.get("From")
name, email_address = email.utils.parseaddr(from_)
body = ""
if msg.is_multipart():
for part in msg.walk():
if part.get_content_type() == "text/plain" and part.get("Content-Disposition") is None:
body += part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8")
else:
body = msg.get_payload(decode=True).decode(msg.get_content_charset() or "utf-8")
emails.append({"id": email_id.decode(), "from": from_, "name": name, "email_address": email_address,
"subject": subject, "body": body})
return jsonify(emails)
@app.route('/classify-email', methods=['POST'])
def classify_email():
data = request.json
email_body = data['body']
email_vectorized = vectorizer.transform([email_body])
prediction = model.predict(email_vectorized)
result = "Suspicious" if prediction == 1 else "Not suspicious"
return jsonify({"result": result})
@app.route('/mark-safe', methods=['POST'])
def mark_safe():
data = request.json
email_id = data['email_id']
# Logic to mark email as safe
return jsonify({"message": f"Email {email_id} marked as safe"})
@app.route('/delete-email', methods=['POST'])
def delete_email():
data = request.json
email_id = data['email_id']
# Connect to the mail server and delete the email
username = data['username']
password = data['password']
try:
mail = imaplib.IMAP4_SSL("outlook.office365.com")
mail.login(username, password)
mail.select("inbox")
mail.store(email_id, '+FLAGS', '\\Deleted')
mail.expunge()
return jsonify({"message": f"Email {email_id} deleted"})
except imaplib.IMAP4.error:
return jsonify({"error": "Failed to delete email"}), 500
if __name__ == '__main__':
app.run(debug=True)