from flask import Flask, request, jsonify, session
from flask_cors import CORS
import imaplib
import email
from email.header import decode_header
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
import traceback
import json
import os

app = Flask(__name__)
CORS(app)
app.secret_key = 'your_secret_key'

SAFE_EMAILS_FILE = 'safe_emails.json'

# Load safe emails from file
def load_safe_emails():
    if os.path.exists(SAFE_EMAILS_FILE):
        with open(SAFE_EMAILS_FILE, 'r') as file:
            return json.load(file)
    return []

# Save safe emails to file
def save_safe_emails(safe_emails):
    with open(SAFE_EMAILS_FILE, 'w') as file:
        json.dump(safe_emails, file)

safe_emails = load_safe_emails()

# Dane treningowe
training_data = [
    ("Urgent account verification", "support@example.com", 1),
    ("Meeting agenda", "boss@example.com", 0),
    ("Password reset request", "no-reply@example.com", 1),
    ("Team lunch schedule", "hr@example.com", 0),
    ("Suspicious login attempt", "security@example.com", 1),
    ("Project update", "colleague@example.com", 0),
    ("Verify your email address", "verification@example.com", 1),
    ("Weekly report", "manager@example.com", 0),
    ("Your account has been suspended", "no-reply@example.com", 1),
    ("Company policy update", "admin@example.com", 0),
    ("Immediate action required", "alert@example.com", 1),
    ("Holiday party invitation", "events@example.com", 0),
    ("Important security update", "security@example.com", 1),
    ("Monthly performance review", "boss@example.com", 0),
    ("Claim your prize now", "lottery@example.com", 1),
    ("Training session details", "training@example.com", 0),
    ("Unauthorized access detected", "alert@example.com", 1),
    ("Office relocation notice", "admin@example.com", 0),
    ("Confirm your subscription", "newsletter@example.com", 1),
    ("Sales team meeting", "sales@example.com", 0),
    ("Your payment is overdue", "billing@example.com", 1),
    ("Client feedback", "client@example.com", 0),
    ("Update your account details", "update@example.com", 1),
    ("Social event invitation", "social@example.com", 0),
    ("Action required: Update password", "security@example.com", 1),
    ("New project assignment", "manager@example.com", 0),
    ("Notice of data breach", "security@example.com", 1),
    ("Weekly newsletter", "newsletter@example.com", 0),
    ("Re: Your recent purchase", "support@example.com", 1),
    ("Performance appraisal meeting", "hr@example.com", 0),
    ("Important account notice", "no-reply@example.com", 1),
    ("Quarterly earnings report", "finance@example.com", 0),
    ("Urgent: Verify your identity", "security@example.com", 1),
    ("Birthday celebration", "events@example.com", 0),
]

subjects = [x[0] for x in training_data]
senders = [x[1] for x in training_data]
labels = [x[2] for x in training_data]

# Połączenie tytułów i nadawców
combined_features = [s + ' ' + senders[i] for i, s in enumerate(subjects)]
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(combined_features)
y = labels

model = MultinomialNB()
model.fit(X, y)

@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    try:
        mail = imaplib.IMAP4_SSL('imap.wp.pl')
        mail.login(username, password)
        session['username'] = username
        session['password'] = password
        return jsonify({'message': 'Login successful'}), 200
    except imaplib.IMAP4.error as e:
        print(f'Login failed: {e}')
        return jsonify({'message': 'Login failed'}), 401
    except Exception as e:
        print('Error during login:', e)
        traceback.print_exc()
        return jsonify({'message': 'Internal server error'}), 500

@app.route('/check_mail', methods=['GET'])
def check_mail():
    if 'username' not in session or 'password' not in session:
        return jsonify({'message': 'Not logged in'}), 401

    username = session['username']
    password = session['password']

    try:
        mail = imaplib.IMAP4_SSL('imap.wp.pl')
        mail.login(username, password)
        mail.select('INBOX')
        result, data = mail.search(None, 'ALL')
        email_ids = data[0].split()[-10:]  # Pobierz ostatnie 10 e-maili
        emails = []

        for e_id in email_ids:
            result, email_data = mail.fetch(e_id, '(RFC822)')
            raw_email = email_data[0][1]
            msg = email.message_from_bytes(raw_email)
            subject = decode_header_value(msg['subject'])
            sender = decode_header_value(msg['from'])
            is_phishing = detect_phishing(subject, sender, e_id.decode())
            emails.append({'subject': subject, 'from': sender, 'is_phishing': is_phishing, 'id': e_id.decode()})

        return jsonify(emails), 200
    except Exception as e:
        print('Error during email check:', e)
        traceback.print_exc()
        return jsonify({'message': 'Internal server error'}), 500

@app.route('/logout', methods=['POST'])
def logout():
    try:
        session.pop('username', None)
        session.pop('password', None)
        return jsonify({'message': 'Logged out'}), 200
    except Exception as e:
        print('Error during logout:', e)
        traceback.print_exc()
        return jsonify({'message': 'Internal server error'}), 500

@app.route('/mark_safe/<email_id>', methods=['POST'])
def mark_safe(email_id):
    global safe_emails
    safe_emails.append(email_id)
    save_safe_emails(safe_emails)
    print(f'Email {email_id} marked as safe')
    return jsonify({"message": f"Email {email_id} marked as safe"}), 200

@app.route('/move_trash/<email_id>', methods=['POST'])
def move_trash(email_id):
    if 'username' not in session or 'password' not in session:
        return jsonify({'message': 'Not logged in'}), 401

    username = session['username']
    password = session['password']

    try:
        mail = imaplib.IMAP4_SSL('imap.wp.pl')
        mail.login(username, password)
        mail.select('INBOX')
        print(f'Trying to move email ID {email_id} to Trash')  # Logging email ID
        mail.store(email_id, '+FLAGS', '\\Deleted')
        mail.expunge()
        print(f'Email {email_id} deleted')  # Logging deletion
        return jsonify({"message": f"Email {email_id} deleted"}), 200
    except Exception as e:
        print(f'Error during moving email to trash: {e}')
        traceback.print_exc()
        return jsonify({'message': 'Internal server error'}), 500

def decode_header_value(value):
    parts = decode_header(value)
    header_parts = []
    for part, encoding in parts:
        if isinstance(part, bytes):
            try:
                if encoding:
                    header_parts.append(part.decode(encoding))
                else:
                    header_parts.append(part.decode('utf-8'))
            except (LookupError, UnicodeDecodeError):
                header_parts.append(part.decode('utf-8', errors='ignore'))
        else:
            header_parts.append(part)
    return ''.join(header_parts)

def detect_phishing(subject, sender, email_id):
    if email_id in safe_emails:
        return False  # If email is marked as safe, it's not phishing

    phishing_keywords = ['urgent', 'verify', 'account', 'suspend', 'login']
    phishing_senders = ['support@example.com', 'no-reply@example.com']
    if any(keyword in subject.lower() for keyword in phishing_keywords) or sender.lower() in phishing_senders:
        return True
    return False

if __name__ == '__main__':
    app.run(port=5000)