import customtkinter as ctk
from core.process_runner import ProcessRunner
from core import reporting
import threading
from tkinter import filedialog # Import for file dialog

class WebFrame(ctk.CTkFrame):
    def __init__(self, master, set_session_callback=None, heading_font=None, text_font=None):
        super().__init__(master)
        self.set_session_callback = set_session_callback
        self.heading_font = heading_font if heading_font else ctk.CTkFont(size=18, weight="bold")
        self.text_font = text_font if text_font else ctk.CTkFont(size=12)
        self.monospace_font = ctk.CTkFont(family="Fira Code", size=12) # For output textboxes
        self.grid_columnconfigure(0, weight=1)
        self.grid_rowconfigure(1, weight=1)

        title = ctk.CTkLabel(self, text="Sección de Hacking Web", font=self.heading_font)
        title.grid(row=0, column=0, padx=20, pady=(20, 10), sticky="ew")

        self.tab_view = ctk.CTkTabview(self)
        self.tab_view.grid(row=1, column=0, padx=20, pady=10, sticky="nsew")

        self.tab_view.add("Nikto Scanner")
        self.tab_view.add("Gobuster")
        self.tab_view.add("SQLMap")

        self.nikto_process_runner = None
        self.gobuster_process_runner = None
        self.sqlmap_process_runner = None # New attribute for SQLMap
        self.current_session_name = None

        self.setup_nikto_tab()
        self.setup_gobuster_tab()
        self.setup_sqlmap_tab() # Call the new setup method

    def setup_nikto_tab(self):
        tab = self.tab_view.tab("Nikto Scanner")
        tab.grid_columnconfigure(0, weight=1)
        tab.grid_rowconfigure(4, weight=1) # Adjusted for output textbox vertical expansion

        # ---- Input Controls Frame ----
        input_controls_frame = ctk.CTkFrame(tab)
        input_controls_frame.grid(row=0, column=0, padx=10, pady=10, sticky="ew")
        input_controls_frame.grid_columnconfigure(0, weight=1)
        input_controls_frame.grid_columnconfigure(1, weight=2) # Entry wider than label

        self.nikto_target_label = ctk.CTkLabel(input_controls_frame, text="URL del Objetivo:")
        self.nikto_target_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.nikto_target_entry = ctk.CTkEntry(input_controls_frame, placeholder_text="ej: http://example.com")
        self.nikto_target_entry.grid(row=0, column=1, padx=10, pady=8, sticky="ew")

        # ---- Action Buttons Frame ----
        action_buttons_frame = ctk.CTkFrame(tab)
        action_buttons_frame.grid(row=1, column=0, padx=10, pady=10, sticky="ew")
        action_buttons_frame.grid_columnconfigure(0, weight=1)
        action_buttons_frame.grid_columnconfigure(1, weight=1)

        self.nikto_start_button = ctk.CTkButton(action_buttons_frame, text="Iniciar Escaneo Nikto", command=self.start_nikto_scan)
        self.nikto_start_button.grid(row=0, column=0, padx=10, pady=12, sticky="ew")

        self.nikto_stop_button = ctk.CTkButton(action_buttons_frame, text="Detener Escaneo", command=self.stop_nikto_scan, fg_color="red", state="disabled")
        self.nikto_stop_button.grid(row=0, column=1, padx=10, pady=12, sticky="ew")

        # ---- Status and Progress Indicators ----
        self.nikto_status_label = ctk.CTkLabel(tab, text="Listo", font=self.text_font)
        self.nikto_status_label.grid(row=2, column=0, padx=10, pady=(5, 0), sticky="w")
        
        self.nikto_progress_bar = ctk.CTkProgressBar(tab, mode="indeterminate")
        self.nikto_progress_bar.grid(row=3, column=0, padx=10, pady=(0, 5), sticky="ew")
        self.nikto_progress_bar.set(0) # Hide it initially

        # ---- Output ----
        self.nikto_output_textbox = ctk.CTkTextbox(tab, state="disabled", font=self.monospace_font)
        self.nikto_output_textbox.grid(row=4, column=0, padx=10, pady=10, sticky="nsew")
        
    def start_nikto_scan(self):
        target = self.nikto_target_entry.get()
        if not target:
            return

        self.nikto_start_button.configure(state="disabled")
        self.nikto_stop_button.configure(state="normal")
        
        self.nikto_status_label.configure(text="Iniciando Nikto...", text_color="orange")
        self.nikto_progress_bar.start()
        
        # Crea una nueva sesión de reporte
        self.current_session_name = reporting.create_session(target)
        if self.set_session_callback:
            self.set_session_callback(self.current_session_name)
        
        nikto_path = "/usr/bin/nikto"
        command = [nikto_path, "-h", target]
        
        self.nikto_output_textbox.configure(state="normal")
        self.nikto_output_textbox.delete("1.0", "end")
        self.nikto_output_textbox.insert("1.0", f"--- Iniciando escaneo de Nikto en {target} ---\n\n")
        self.nikto_output_textbox.configure(state="disabled")

        # Usar ProcessRunner para ejecutar el comando
        self.nikto_process_runner = ProcessRunner(
            command,
            self.nikto_output_textbox,
            status_label=self.nikto_status_label,
            progress_bar=self.nikto_progress_bar
        )
        self.nikto_process_runner.start()
        
        threading.Thread(target=self.wait_and_save_nikto, daemon=True).start()

    def stop_nikto_scan(self):
        if self.nikto_process_runner:
            self.nikto_process_runner.stop()
            self.nikto_start_button.configure(state="normal")
            self.nikto_stop_button.configure(state="disabled")
            self.nikto_status_label.configure(text="Detenido.", text_color="red")
            self.nikto_progress_bar.stop()

    def wait_and_save_nikto(self):
        # Esperar a que el hilo del proceso termine
        if self.nikto_process_runner:
            self.nikto_process_runner.thread.join()

        # Una vez que el proceso ha terminado
        full_output = self.nikto_output_textbox.get("1.0", "end")
        reporting.log_tool_output(self.current_session_name, "Nikto Scan", full_output, output_format="txt")
        
        # Volver a habilitar los botones en el hilo principal
        self.master.after(0, self._reset_nikto_buttons)

    def _reset_nikto_buttons(self):
        self.nikto_start_button.configure(state="normal")
        self.nikto_stop_button.configure(state="disabled")
        self.nikto_status_label.configure(text="Completado.", text_color="green")
        self.nikto_progress_bar.stop()
        


    def select_gobuster_wordlist(self):
        filepath = filedialog.askopenfilename(
            title="Seleccionar Wordlist para Gobuster",
            filetypes=[("Wordlist files", "*.txt"), ("All files", "*.*")]
        )
        if filepath:
            self.gobuster_wordlist_entry.delete(0, "end")
            self.gobuster_wordlist_entry.insert(0, filepath)

    def start_gobuster_scan(self):
        target = self.gobuster_target_entry.get()
        wordlist = self.gobuster_wordlist_entry.get()
        if not target or not wordlist:
            return

        self.gobuster_start_button.configure(state="disabled")
        self.gobuster_stop_button.configure(state="normal")
        
        self.gobuster_status_label.configure(text="Iniciando Gobuster...", text_color="orange")
        self.gobuster_progress_bar.start()

        self.current_session_name = reporting.create_session(target)
        if self.set_session_callback:
            self.set_session_callback(self.current_session_name)
        
        gobuster_path = "/usr/bin/gobuster"
        command = [gobuster_path, "dir", "-u", target, "-w", wordlist]
        
        self.gobuster_output_textbox.configure(state="normal")
        self.gobuster_output_textbox.delete("1.0", "end")
        self.gobuster_output_textbox.insert("1.0", f"--- Iniciando escaneo de Gobuster en {target} con {wordlist} ---\n\n")
        self.gobuster_output_textbox.configure(state="disabled")

        self.gobuster_process_runner = ProcessRunner(
            command,
            self.gobuster_output_textbox,
            status_label=self.gobuster_status_label,
            progress_bar=self.gobuster_progress_bar
        )
        self.gobuster_process_runner.start()
        
        threading.Thread(target=self.wait_and_save_gobuster, daemon=True).start()

    def stop_gobuster_scan(self):
        if self.gobuster_process_runner:
            self.gobuster_process_runner.stop()
            self.gobuster_start_button.configure(state="normal")
            self.gobuster_stop_button.configure(state="disabled")
            self.gobuster_status_label.configure(text="Detenido.", text_color="red")
            self.gobuster_progress_bar.stop()

    def wait_and_save_gobuster(self):
        if self.gobuster_process_runner:
            self.gobuster_process_runner.thread.join()

        full_output = self.gobuster_output_textbox.get("1.0", "end")
        reporting.log_tool_output(self.current_session_name, "Gobuster Scan", full_output, output_format="txt")
        
        self.master.after(0, self._reset_gobuster_buttons)

    def _reset_gobuster_buttons(self):
        self.gobuster_start_button.configure(state="normal")
        self.gobuster_stop_button.configure(state="disabled")
        self.gobuster_status_label.configure(text="Completado.", text_color="green")
        self.gobuster_progress_bar.stop()

    # ---- SQLMAP TAB ----
    def setup_sqlmap_tab(self):
        tab = self.tab_view.tab("SQLMap")
        tab.grid_columnconfigure(0, weight=1)
        tab.grid_rowconfigure(4, weight=1) # Adjusted for output textbox vertical expansion

        # ---- Input Controls Frame ----
        input_controls_frame = ctk.CTkFrame(tab)
        input_controls_frame.grid(row=0, column=0, padx=10, pady=10, sticky="ew")
        input_controls_frame.grid_columnconfigure(0, weight=1)
        input_controls_frame.grid_columnconfigure(1, weight=2) # Entry wider than label

        # Target
        target_frame = ctk.CTkFrame(input_controls_frame, fg_color="transparent")
        target_frame.grid(row=0, column=0, columnspan=2, padx=5, pady=5, sticky="ew")
        target_frame.grid_columnconfigure(0, weight=1)
        target_frame.grid_columnconfigure(1, weight=3)
        self.sqlmap_target_label = ctk.CTkLabel(target_frame, text="URL/Request File:")
        self.sqlmap_target_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.sqlmap_target_entry = ctk.CTkEntry(target_frame, placeholder_text="URL o ruta a archivo .req")
        self.sqlmap_target_entry.grid(row=0, column=1, padx=10, pady=8, sticky="ew")

        # Risk and Level
        risk_level_frame = ctk.CTkFrame(input_controls_frame, fg_color="transparent")
        risk_level_frame.grid(row=1, column=0, columnspan=2, padx=5, pady=5, sticky="ew")
        risk_level_frame.grid_columnconfigure(0, weight=1)
        risk_level_frame.grid_columnconfigure(1, weight=1)
        risk_level_frame.grid_columnconfigure(2, weight=1)
        risk_level_frame.grid_columnconfigure(3, weight=1)

        self.sqlmap_risk_label = ctk.CTkLabel(risk_level_frame, text="Riesgo (Risk):")
        self.sqlmap_risk_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.sqlmap_risk_option = ctk.CTkOptionMenu(risk_level_frame, values=["1", "2", "3"])
        self.sqlmap_risk_option.set("1")
        self.sqlmap_risk_option.grid(row=0, column=1, padx=10, pady=8, sticky="ew")

        self.sqlmap_level_label = ctk.CTkLabel(risk_level_frame, text="Nivel (Level):")
        self.sqlmap_level_label.grid(row=0, column=2, padx=10, pady=8, sticky="w")
        self.sqlmap_level_option = ctk.CTkOptionMenu(risk_level_frame, values=["1", "2", "3", "4", "5"])
        self.sqlmap_level_option.set("1")
        self.sqlmap_level_option.grid(row=0, column=3, padx=10, pady=8, sticky="ew")

        # Checkboxes for actions
        checkbox_frame = ctk.CTkFrame(input_controls_frame, fg_color="transparent")
        checkbox_frame.grid(row=2, column=0, columnspan=2, padx=5, pady=5, sticky="ew")
        checkbox_frame.grid_columnconfigure(0, weight=1)
        checkbox_frame.grid_columnconfigure(1, weight=1)

        self.sqlmap_dbs_var = ctk.BooleanVar(value=False)
        self.sqlmap_dbs_checkbox = ctk.CTkCheckBox(checkbox_frame, text="Enumerar Bases de Datos (--dbs)", variable=self.sqlmap_dbs_var)
        self.sqlmap_dbs_checkbox.grid(row=0, column=0, padx=10, pady=8, sticky="w")

        self.sqlmap_tables_var = ctk.BooleanVar(value=False)
        self.sqlmap_tables_checkbox = ctk.CTkCheckBox(checkbox_frame, text="Enumerar Tablas (--tables)", variable=self.sqlmap_tables_var)
        self.sqlmap_tables_checkbox.grid(row=0, column=1, padx=10, pady=8, sticky="w")

        self.sqlmap_dump_var = ctk.BooleanVar(value=False)
        self.sqlmap_dump_checkbox = ctk.CTkCheckBox(checkbox_frame, text="Volcar Datos (--dump)", variable=self.sqlmap_dump_var)
        self.sqlmap_dump_checkbox.grid(row=1, column=0, padx=10, pady=8, sticky="w")

        self.sqlmap_batch_var = ctk.BooleanVar(value=True) # Usually good to have this on
        self.sqlmap_batch_checkbox = ctk.CTkCheckBox(checkbox_frame, text="Modo Batch (--batch)", variable=self.sqlmap_batch_var)
        self.sqlmap_batch_checkbox.grid(row=1, column=1, padx=10, pady=8, sticky="w")
        
        self.sqlmap_random_agent_var = ctk.BooleanVar(value=True)
        self.sqlmap_random_agent_checkbox = ctk.CTkCheckBox(checkbox_frame, text="Random User-Agent (--random-agent)", variable=self.sqlmap_random_agent_var)
        self.sqlmap_random_agent_checkbox.grid(row=2, column=0, padx=10, pady=8, sticky="w")
        
        # POST Data
        post_data_frame = ctk.CTkFrame(input_controls_frame, fg_color="transparent")
        post_data_frame.grid(row=3, column=0, columnspan=2, padx=5, pady=5, sticky="ew")
        post_data_frame.grid_columnconfigure(0, weight=1)
        post_data_frame.grid_columnconfigure(1, weight=3)
        self.sqlmap_data_label = ctk.CTkLabel(post_data_frame, text="Datos POST (--data):")
        self.sqlmap_data_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.sqlmap_data_entry = ctk.CTkEntry(post_data_frame, placeholder_text="ej: 'user=admin&pass=password'")
        self.sqlmap_data_entry.grid(row=0, column=1, padx=10, pady=8, sticky="ew")

        # ---- Action Buttons Frame ----
        action_buttons_frame = ctk.CTkFrame(tab)
        action_buttons_frame.grid(row=1, column=0, padx=10, pady=10, sticky="ew")
        action_buttons_frame.grid_columnconfigure(0, weight=1)
        action_buttons_frame.grid_columnconfigure(1, weight=1)

        self.sqlmap_start_button = ctk.CTkButton(action_buttons_frame, text="Iniciar SQLMap", command=self.start_sqlmap_scan)
        self.sqlmap_start_button.grid(row=0, column=0, padx=10, pady=12, sticky="ew")

        self.sqlmap_stop_button = ctk.CTkButton(action_buttons_frame, text="Detener SQLMap", command=self.stop_sqlmap_scan, fg_color="red", state="disabled")
        self.sqlmap_stop_button.grid(row=0, column=1, padx=10, pady=12, sticky="ew")
        
        # ---- Status and Progress Indicators ----
        self.sqlmap_status_label = ctk.CTkLabel(tab, text="Listo", font=self.text_font)
        self.sqlmap_status_label.grid(row=2, column=0, padx=10, pady=(5, 0), sticky="w")

        self.sqlmap_progress_bar = ctk.CTkProgressBar(tab, mode="indeterminate")
        self.sqlmap_progress_bar.grid(row=3, column=0, padx=10, pady=(0, 5), sticky="ew")
        self.sqlmap_progress_bar.set(0) # Hide it initially

        # ---- Output ----
        self.sqlmap_output_textbox = ctk.CTkTextbox(tab, state="disabled", font=self.monospace_font)
        self.sqlmap_output_textbox.grid(row=4, column=0, padx=10, pady=10, sticky="nsew")

    def start_sqlmap_scan(self):
        target = self.sqlmap_target_entry.get()
        if not target:
            return

        self.sqlmap_start_button.configure(state="disabled")
        self.sqlmap_stop_button.configure(state="normal")
        
        self.sqlmap_status_label.configure(text="Iniciando SQLMap...", text_color="orange")
        self.sqlmap_progress_bar.start()

        self.current_session_name = reporting.create_session(target)
        if self.set_session_callback:
            self.set_session_callback(self.current_session_name)
        
        sqlmap_path = "/usr/bin/sqlmap"
        command = [sqlmap_path, "-u", target] # Base command

        # Add options based on user selection
        command.extend(["--risk", self.sqlmap_risk_option.get()])
        command.extend(["--level", self.sqlmap_level_option.get()])

        if self.sqlmap_dbs_var.get():
            command.append("--dbs")
        if self.sqlmap_tables_var.get():
            command.append("--tables")
        if self.sqlmap_dump_var.get():
            command.append("--dump")
        if self.sqlmap_batch_var.get():
            command.append("--batch")
        if self.sqlmap_random_agent_var.get():
            command.append("--random-agent")

        post_data = self.sqlmap_data_entry.get()
        if post_data:
            command.extend(["--data", post_data])

        self.sqlmap_output_textbox.configure(state="normal")
        self.sqlmap_output_textbox.delete("1.0", "end")
        self.sqlmap_output_textbox.insert("1.0", f"--- Iniciando SQLMap en {target} ---\nComando: {' '.join(command)}\n\n")
        self.sqlmap_output_textbox.configure(state="disabled")

        self.sqlmap_process_runner = ProcessRunner(
            command,
            self.sqlmap_output_textbox,
            status_label=self.sqlmap_status_label,
            progress_bar=self.sqlmap_progress_bar
        )
        self.sqlmap_process_runner.start()
        
        threading.Thread(target=self.wait_and_save_sqlmap, daemon=True).start()

    def stop_sqlmap_scan(self):
        if self.sqlmap_process_runner:
            self.sqlmap_process_runner.stop()
            self.sqlmap_start_button.configure(state="normal")
            self.sqlmap_stop_button.configure(state="disabled")
            self.sqlmap_status_label.configure(text="Detenido.", text_color="red")
            self.sqlmap_progress_bar.stop()

    def wait_and_save_sqlmap(self):
        if self.sqlmap_process_runner:
            self.sqlmap_process_runner.thread.join()

        full_output = self.sqlmap_output_textbox.get("1.0", "end")
        reporting.log_tool_output(self.current_session_name, "SQLMap Scan", full_output, output_format="txt")
        
        self.master.after(0, self._reset_sqlmap_buttons)

    def _reset_sqlmap_buttons(self):
        self.sqlmap_start_button.configure(state="normal")
        self.sqlmap_stop_button.configure(state="disabled")
        self.sqlmap_status_label.configure(text="Completado.", text_color="green")
        self.sqlmap_progress_bar.stop()

    def setup_gobuster_tab(self):
        tab = self.tab_view.tab("Gobuster")
        tab.grid_columnconfigure(0, weight=1)
        tab.grid_rowconfigure(5, weight=1) # Adjusted for output textbox vertical expansion

        # ---- Input Controls Frame ----
        input_controls_frame = ctk.CTkFrame(tab)
        input_controls_frame.grid(row=0, column=0, padx=10, pady=10, sticky="ew")
        input_controls_frame.grid_columnconfigure(0, weight=1)
        input_controls_frame.grid_columnconfigure(1, weight=2) # Entry wider than label

        self.gobuster_target_label = ctk.CTkLabel(input_controls_frame, text="URL del Objetivo:")
        self.gobuster_target_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.gobuster_target_entry = ctk.CTkEntry(input_controls_frame, placeholder_text="ej: http://example.com")
        self.gobuster_target_entry.grid(row=0, column=1, padx=10, pady=8, sticky="ew")

        # Wordlist selection
        wordlist_selection_frame = ctk.CTkFrame(tab)
        wordlist_selection_frame.grid(row=1, column=0, padx=10, pady=5, sticky="ew")
        wordlist_selection_frame.grid_columnconfigure(0, weight=1)
        wordlist_selection_frame.grid_columnconfigure(1, weight=0) # Button fixed width

        self.gobuster_wordlist_label = ctk.CTkLabel(wordlist_selection_frame, text="Ruta Wordlist:")
        self.gobuster_wordlist_label.grid(row=0, column=0, padx=10, pady=8, sticky="w")
        self.gobuster_wordlist_entry = ctk.CTkEntry(wordlist_selection_frame, placeholder_text="ej: /usr/share/wordlists/dirbuster/directory-list-2.3-small.txt")
        self.gobuster_wordlist_entry.grid(row=0, column=1, padx=(5, 5), sticky="ew")
        
        self.gobuster_wordlist_button = ctk.CTkButton(wordlist_selection_frame, text="Seleccionar", command=self.select_gobuster_wordlist, width=90)
        self.gobuster_wordlist_button.grid(row=0, column=2, padx=(10, 0))

        # ---- Action Buttons Frame ----
        action_buttons_frame = ctk.CTkFrame(tab)
        action_buttons_frame.grid(row=2, column=0, padx=10, pady=10, sticky="ew")
        action_buttons_frame.grid_columnconfigure(0, weight=1)
        action_buttons_frame.grid_columnconfigure(1, weight=1)

        self.gobuster_start_button = ctk.CTkButton(action_buttons_frame, text="Iniciar Escaneo Gobuster", command=self.start_gobuster_scan)
        self.gobuster_start_button.grid(row=0, column=0, padx=10, pady=12, sticky="ew")

        self.gobuster_stop_button = ctk.CTkButton(action_buttons_frame, text="Detener Escaneo", command=self.stop_gobuster_scan, fg_color="red", state="disabled")
        self.gobuster_stop_button.grid(row=0, column=1, padx=10, pady=12, sticky="ew")

        # ---- Status and Progress Indicators ----
        self.gobuster_status_label = ctk.CTkLabel(tab, text="Listo", font=self.text_font)
        self.gobuster_status_label.grid(row=3, column=0, padx=10, pady=(5, 0), sticky="w")
        
        self.gobuster_progress_bar = ctk.CTkProgressBar(tab, mode="indeterminate")
        self.gobuster_progress_bar.grid(row=4, column=0, padx=10, pady=(0, 5), sticky="ew")
        self.gobuster_progress_bar.set(0) # Hide it initially

        # ---- Output ----
        self.gobuster_output_textbox = ctk.CTkTextbox(tab, state="disabled", font=self.monospace_font)
        self.gobuster_output_textbox.grid(row=5, column=0, padx=10, pady=10, sticky="nsew")
