import tkinter as tk
from tkinter import messagebox
import subprocess
import os
import csv
import inotify_simple
import threading
import time
import re
import requests
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from collections import defaultdict
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
from datetime import datetime

permission_operations = None
# Define event masks manually
IN_CREATE = 0x00000100
IN_DELETE = 0x00000200
IN_MODIFY = 0x00000002
IN_OPEN = 0x00000020
IN_ISDIR = 0x40000000

####################


TEST_DATA_PATH = 'combined_log_summary.csv'
VARIABLE_NAMES_PATH = 'output.txt'
def predict_ransomware():
# Load the trained model
    model = tf.keras.models.load_model('updated_ransomware_classifier.h5')

    # Load and prepare test data
    # Read variable names
    with open(VARIABLE_NAMES_PATH, encoding='utf-8') as f:
        columns = [line.split(';')[1].strip() for line in f]

    # Load test data
    data = pd.read_csv(TEST_DATA_PATH, header=None, names=columns)

    # Check and clean column names
    data.columns = data.columns.str.strip()
    X = data
    # Standardize the features
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    # Make predictions
    predictions = model.predict(X)
    predicted_labels = (predictions > 0.5).astype(int)


    # Convert predictions to "Yes" or "No"
    predicted_labels_text = ['Yes' if label == 1 else 'No' for label in predicted_labels.flatten()]


    # Get current timestamp
    timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')

                                                                                                                  
    # Save predictions and true labels to a CSV file with timestamp
    output_df = pd.DataFrame({
        'Timestamp': [timestamp] * len(predicted_labels_text),  # Add timestamp column
        'Predicted Label': predicted_labels_text
    })

    output_file = f'prediction.csv'
    output_df.to_csv(output_file, index=False)
    print(f"Predictions saved to {output_file} ({timestamp})")


    def send_predictions_to_api(file_path):
        url = "http://142.93.221.85:8000/predict-malware/"
        with open(file_path, 'rb') as f:
            files = {'file': f}
            response = requests.post(url, files=files)
            if response.status_code == 200:
                print(f"Successfully sent {file_path} to API.")
            else:
                print(f"Failed to send {file_path} to API. Status code: {response.status_code}")
    try:
        send_predictions_to_api(output_file)
    except:
        print("Error Connection Server")

####################

ID = 0

is_flip = 0
flipped = False
class PermissionChangeHandler(FileSystemEventHandler):
    def __init__(self):
        super().__init__()
        self.file_types = set()

    def get_file_extension(self, file_path):
        """Extracts the file extension from the file path."""
        _, ext = os.path.splitext(file_path)
        return ext.strip(".")  # Strip the dot from the extension

    def on_modified(self, event):
        if not event.is_directory:
            file_path = event.src_path
            file_extension = self.get_file_extension(file_path)
            
            # Collect all file types
            file_types = set()
            for operations in permission_operations.values():
                for key in operations:
                    match = re.search(r'\.\w+$', key)
                    if match:
                        file_types.add(match.group().strip('.'))

            if file_extension in file_types:
                current_permissions = oct(os.stat(file_path).st_mode & 0o777)
            

                # Check all operations (chmod/chown) for this file type
                for operation, perms in permission_operations.items():
                    for key in perms:
                        if file_extension in key:
                            perms[key] += 1
                            # print(f"Updated {operation} for {file_extension}: {perms[key]}")

class AuditDManagerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("AuditD Manager")
        self.root.geometry("400x350")  # Adjusted for additional widget

        # Create Widgets
        self.install_button = tk.Button(root, text="Install AuditD", command=self.install_auditd)
        self.install_button.pack(pady=10)
       
        self.start_button = tk.Button(root, text="Start AuditD", command=self.start_auditd)
        self.start_button.pack(pady=10)

        self.stop_button = tk.Button(root, text="Stop AuditD", command=self.stop_auditd)
        self.stop_button.pack(pady=10)

        self.status_button = tk.Button(root, text="Check Status", command=self.check_status)
        self.status_button.pack(pady=10)

        # Add Text Entry for Watch Path

        # Initialize monitoring flags and data structures
        self.monitoring = False
        self.log_file = "/var/log/audit/audit.log"
        self.combined_csv_file = "combined_log_summary.csv"
        self.monitored_files_set = {
            'bash.bashrc', 'bash_completion.d', 'environment', 'fstab', 'fwupd', 'group', 'host.conf', 'hosts', 'init.d',
            'inputrc', 'ld.so.cache', 'locale.alias', 'locale.conf', 'login.defs', 'machine-id', 'modprobe.d', 'nsswitch.conf',
            'passwd', 'sensors.d', 'sensors3.conf', 'shadow', 'shells', 'sudo.conf', 'sudoers', 'sudoers.d'
        }
        self.log_counts = {key: 0 for key in [
            'Id','PROCTITLE', 'AVC', 'SYSCALL', 'USER_AUTH', 'USER_ACCT',
            'USER_CMD', 'CRED_REFR', 'USER_START', 'USER_AVC', 'USER_END', 'CRED_DISP', 'CRED_ACQ',
            'LOGIN', 'SERVICE_START', 'SERVICE_STOP']}
        
        # Track file extensions
        self.ext_count = {ext: {'modified': 0, 'created': 0, 'deleted': 0, 'opened': 0} for ext in [
            '.db', '.AR', '.01', '.GIF', '.TXT', '.scc', '.dat', '.bmp', '.STF', '.scf',
            '.exe', '.typelib', '.cl', '.ocx', '.xml', '.json', '.csv', '.html', '.css',
            '.js', '.py', '.log', '.sql', '.pdf', '.doc', '.docx', '.ppt', '.pptx',
            '.xlsx', '.jpg', '.jpeg', '.png', '.mp4', '.mp3', '.zip', '.tar', '.gz', '.rar', '.7z', '.apk', '.iso']}
        
        # Track permission operations
        global permission_operations
        permission_operations = {
            'chmod': {f'chmod{perm}{ext}': 0 for perm in ['644', '755', '777'] for ext in self.ext_count},
            'chown': {f'chown{owner}{ext}': 0 for owner in ['user', 'group'] for ext in self.ext_count},
            'chgrp': {f'chgrp{group}{ext}': 0 for group in ['staff', 'admin'] for ext in self.ext_count}
        }

        # Directory operations tracking
        self.directory_count = {'created': 0, 'deleted': 0, 'modified': 0, 'opened': 0}

        # Initialize inotify
        self.inotify = inotify_simple.INotify()
        self.EVENT_MASKS = IN_CREATE | IN_DELETE | IN_MODIFY | IN_OPEN | IN_ISDIR
        self.watch_path = '/etc'  # Default path, will be updated
        self.watch_descriptor2 = self.inotify.add_watch(self.watch_path, self.EVENT_MASKS)

        # Observer for filesystem events
        self.observer = None
        self.event_handler = None
        self.monitor_thread = threading.Thread(target=self.monitor_logs)

        # Initialize file monitoring data
        self.open_count = defaultdict(int)

    def run_command(self, command, success_message, error_message):
        try:
            result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            messagebox.showinfo("Success", success_message)
        except subprocess.CalledProcessError as e:
            messagebox.showerror("Error", f"{error_message}\n\n{e.stderr.decode()}")

    def prompt_for_password(self, command, success_message, error_message):
        password_window = tk.Toplevel(self.root)
        password_window.title("Enter Sudo Password")

        tk.Label(password_window, text="Enter your sudo password:").pack(pady=10)

        password_entry = tk.Entry(password_window, show="*")
        password_entry.pack(pady=5)
        
        def on_submit():
            password = password_entry.get()
            password_window.destroy()
            if not password:
                messagebox.showwarning("Input Error", "Please enter your sudo password.")
                return

            full_command = f"echo {password} | sudo -S {command}"
            self.run_command(full_command, success_message, error_message)
        tk.Button(password_window, text="Submit", command=on_submit).pack(pady=10)

    def install_auditd(self):
        command = "sudo apt-get update && sudo apt-get install -y auditd"
        self.prompt_for_password(command, "AuditD installed successfully!", "Failed to install AuditD.")

    def start_auditd(self):
        command = "sudo systemctl start auditd"
        self.prompt_for_password(command, "AuditD started successfully!", "Failed to start AuditD.")
        self.start_monitoring()

    def stop_auditd(self):
        command = "sudo systemctl stop auditd"
        self.prompt_for_password(command, "AuditD stopped successfully!", "Failed to stop AuditD.")
        self.stop_monitoring()

    def check_status(self):
        command = "systemctl status auditd"
        try:
            result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            status = result.stdout.decode()
            messagebox.showinfo("AuditD Status", status)
        except subprocess.CalledProcessError as e:
            messagebox.showerror("Error", f"Failed to check status of AuditD.\n\n{e.stderr.decode()}")

    def start_monitoring(self):
        self.monitoring = True
        if not self.monitor_thread.is_alive():
            self.monitor_thread = threading.Thread(target=self.monitor_logs)
            self.monitor_thread.start()

        # Get the user-defined watch path
        self.watch_path = '/etc'  # Default to root if empty
        self.watch_descriptor = self.inotify.add_watch(self.watch_path, self.EVENT_MASKS)

        # Start filesystem event monitoring
        if self.observer is None:
            self.event_handler = PermissionChangeHandler()
            self.observer = Observer()
            self.observer.schedule(self.event_handler, '/home', recursive=True)
            self.observer.start()

    def stop_monitoring(self):
        self.monitoring = False
        if self.monitor_thread.is_alive():
            self.monitor_thread.join()

        # Stop filesystem event monitoring
        if self.observer:
            self.observer.stop()
            self.observer.join()

    def monitor_logs(self):
        while self.monitoring:
            if os.path.exists(self.log_file):
                with open(self.log_file, 'r') as f:
                    lines = f.readlines()

                for line in lines:
                    if 'type=' in line:
                        log_type = line.split('type=')[1].split(' ')[0]
                        if log_type in self.log_counts:
                            self.log_counts[log_type] += 1

                self.update_csv()

            self.monitor_extensions()
            predict_ransomware()
            time.sleep(5)  # Sleep for one second before the next update

    def update_csv(self):
        # headers = [
        #     'Id' ,'PROCTITLE', 'AVC', 'SYSCALL', 'USER_AUTH', 'USER_ACCT',
        #     'USER_CMD', 'CRED_REFR', 'USER_START', 'USER_AVC', 'USER_END', 'CRED_DISP', 'CRED_ACQ',
        #     'LOGIN', 'SERVICE_START', 'SERVICE_STOP'
        # ] + [f'chmod{perm}{ext}' for perm in ['644', '755', '777'] for ext in self.ext_count] + \
        # [f'chown{owner}{ext}' for owner in ['user', 'group'] for ext in self.ext_count] + \
        # [f'chgrp{group}{ext}' for group in ['staff', 'admin'] for ext in self.ext_count] + \
        # [f'Modified({ext})' for ext in self.ext_count] + \
        # [f'Created({ext})' for ext in self.ext_count] + \
        # [f'Deleted({ext})' for ext in self.ext_count] + \
        # [f'Opened({ext})' for ext in self.ext_count] + \
        # ['Directories Created', 'Directories Deleted', 'Directories Modified', 'Directories Opened']+ \
        # list(self.monitored_files_set) 

        global ID
        ID += 1
        global is_flip
        global flipped
        if flipped:
            is_flip = 1
            flipped = False
        else:
            is_flip = 0
            flipped = True
        
        row = [
            ID,
            self.log_counts.get('PROCTITLE', 0),
            self.log_counts.get('AVC', 0),
            self.log_counts.get('SYSCALL', 0),
            self.log_counts.get('USER_AUTH', 0),
            self.log_counts.get('USER_ACCT', 0),
            self.log_counts.get('USER_CMD', 0),
            self.log_counts.get('CRED_REFR', 0),
            self.log_counts.get('USER_START', 0),
            self.log_counts.get('USER_AVC', 0),
            self.log_counts.get('USER_END', 0),
            self.log_counts.get('CRED_DISP', 0),
            self.log_counts.get('CRED_ACQ', 0),
            self.log_counts.get('LOGIN', 0),
            self.log_counts.get('SERVICE_START', 0),
            self.log_counts.get('SERVICE_STOP', 0),
        ]
        
        # print(permission_operations['chmod'])
        # Add permission operations and extensions
        row.extend(permission_operations['chmod'].values())
        row.extend(permission_operations['chown'].values())
        row.extend(permission_operations['chgrp'].values())

        # Add extension counts for modification, creation, deletion, and opening
        for ext in self.ext_count:
            row.extend([
                self.ext_count[ext]['modified'],
                self.ext_count[ext]['created'],
                self.ext_count[ext]['deleted'],
                self.ext_count[ext]['opened'],
            ])

        # Add directory counts
        row.extend([
            self.directory_count['created'],
            self.directory_count['deleted'],
            self.directory_count['modified'],
            self.directory_count['opened']
        ])

        # Add monitored files open counts
        row.extend(self.open_count.get(file, 0) for file in sorted(self.monitored_files_set))

        # Write to CSV, append if file exists
        file_exists = os.path.isfile(self.combined_csv_file)
        with open(self.combined_csv_file, 'a', newline='') as csv_file:
            writer = csv.writer(csv_file)
            if not file_exists:
                pass
            writer.writerow(row)


    def monitor_extensions(self):
        events = self.inotify.read(timeout=100000)
        for event in events:
            (_, event_types, _, filename) = event

            filename = event.name
            ext = os.path.splitext(filename)[1]
            if ext in self.ext_count:
                if event.mask & IN_CREATE:
                    self.ext_count[ext]['created'] += 1
                if event.mask & IN_DELETE:
                    self.ext_count[ext]['deleted'] += 1
                if event.mask & IN_MODIFY:
                    self.ext_count[ext]['modified'] += 1
                if event.mask & IN_OPEN:
                    self.ext_count[ext]['opened'] += 1
                    if filename in self.monitored_files_set:
                        self.open_count[filename] += 1

            if event.mask & IN_ISDIR:
                if event.mask & IN_CREATE:
                    self.directory_count['created'] += 1
                if event.mask & IN_DELETE:
                    self.directory_count['deleted'] += 1
                if event.mask & IN_MODIFY:
                    self.directory_count['modified'] += 1
                if event.mask & IN_OPEN:
                    self.directory_count['opened'] += 1

if __name__ == "__main__":
    root = tk.Tk()
    app = AuditDManagerApp(root)
    root.mainloop()