Enterprise-Email-Scanner-For-Malware.py

import os
import win32com.client
import vt
import time
import matplotlib.pyplot as plt
import tkinter as tk
from tkinter import simpledialog, ttk

# Default configurations
default_config = {
    'vt_api_key': '',
    'attachment_dir': r"",
    'log_file_path': r""
}

config_filename = 'config.txt'

# Function to load config from file
def load_config():
    if os.path.exists(config_filename):
        with open(config_filename, 'r') as config_file:
            return {line.split('=')[0]: line.split('=')[1].strip() for line in config_file.readlines()}
    else:
        return default_config

# Function to save config to file
def save_config(config):
    with open(config_filename, 'w') as config_file:
        for key, value in config.items():
            config_file.write(f"{key}={value}\n")

# Load configuration at start
config = load_config()

def save_attachment(item):
    attachments = []
    if not os.path.exists(config['attachment_dir']):
        os.makedirs(config['attachment_dir'])
    
    for attachment in item.Attachments:
        attachment_path = os.path.join(config['attachment_dir'], attachment.FileName)
        try:
            attachment.SaveAsFile(attachment_path)
            attachments.append(attachment_path)
        except Exception as e:
            print(f"Error saving attachment: {str(e)}")
    return attachments

def scan_with_virustotal(file_paths, status_label):
    client = vt.Client(config['vt_api_key'])
    malicious = False
    try:
        for idx, file_path in enumerate(file_paths):
            if os.path.exists(file_path):
                status_label.config(text=f"Scanning file {idx+1}/{len(file_paths)} with VirusTotal")
                status_label.update()
                start_time = time.time()
                with open(file_path, 'rb') as f:
                    analysis = client.scan_file(f, wait_for_completion=True)
                duration = time.time() - start_time
                
                if analysis.stats['malicious'] > 0:
                    malicious = True
                    break
    finally:
        client.close()
        for file_path in file_paths:
            if os.path.exists(file_path):
                os.remove(file_path)
    return malicious

def move_email_to_spam(message, outlook):
    try:
        # Get the default Junk folder
        junk_folder = outlook.GetDefaultFolder(23)  # 23 is the folder type for Junk in Outlook
        
        # Move the message to the Junk folder
        message.Move(junk_folder)
        print(f"Moved email with subject '{message.Subject}' to Junk.")
    except Exception as e:
        print(f"Error moving email to Junk: {str(e)}")

def main_process(progress_bar, status_label):
    if os.path.exists(config['log_file_path']):
        with open(config['log_file_path'], 'r') as log_file:
            processed_emails = set(log_file.read().splitlines())
    else:
        processed_emails = set()

    num_infected = 0
    num_scanned = 0

    outlook = win32com.client.Dispatch("Outlook.Application").GetNamespace("MAPI")
    inbox = outlook.GetDefaultFolder(6)

    messages = inbox.Items
    messages.Sort("[ReceivedTime]", True)
    total_messages = len(messages)

    progress_bar["maximum"] = total_messages

    for idx, message in enumerate(messages):
        entry_id = message.EntryID
        if entry_id in processed_emails:
            continue
        
        attachments = save_attachment(message)
        num_scanned += len(attachments)
        if attachments and scan_with_virustotal(attachments, status_label):  # Pass the status_label
            num_infected += 1
            status_label.config(text=f"Found infected file in email {idx+1}")
            move_email_to_spam(message, outlook)  # Move email to spam
        else:
            status_label.config(text=f"No infection found in email {idx+1}")
        status_label.update()

        with open(config['log_file_path'], 'a') as log_file:
            log_file.write(entry_id + '\n')
        processed_emails.add(entry_id)

        progress_bar["value"] = idx + 1
        status_label.config(text=f"Processing {idx + 1}/{total_messages} emails")
        progress_bar.update()
        status_label.update()

    categories = ['Infected Files', 'Scanned Files']
    values = [num_infected, num_scanned]

    plt.figure(figsize=(8, 6))
    plt.bar(categories, values, color=['red', 'blue'])
    plt.title('Results of Files Scanning')
    plt.ylabel('Number of Files')
    for i, v in enumerate(values):
        plt.text(i, v + 0.2, str(v), ha='center', color='black')
    plt.show()

def update_config():
    for key in config:
        new_value = simpledialog.askstring("Input", f"Enter {key}: (current: {config[key]})")
        if new_value:
            config[key] = new_value
    save_config(config)  # Save updated configuration to file

def start_gui():
    root = tk.Tk()
    root.title("Email Scanner")

    frame = tk.Frame(root)
    frame.pack(pady=20)

    status_label = tk.Label(frame, text="Press 'Start Scanning' to begin.")
    status_label.pack()

    progress_bar = ttk.Progressbar(frame, length=300)
    progress_bar.pack(pady=10)

    start_button = tk.Button(frame, text="Start Scanning", command=lambda: main_process(progress_bar, status_label))
    start_button.pack(fill='x')

    config_button = tk.Button(frame, text="Configure Settings", command=update_config)
    config_button.pack(fill='x')

    root.mainloop()

if __name__ == "__main__":
    start_gui()