#!/usr/bin/env python3
"""
PentestGPT Tab Implementation
Handles the PentestGPT interface and controls for DarkHal 2.0
"""

import tkinter as tk
from tkinter import ttk, messagebox, filedialog
from pathlib import Path
import subprocess
import threading

class PentestGPTTab:
    """PentestGPT tab with penetration testing agent controls."""
    
    def __init__(self, parent: ttk.Frame, settings_manager, main_app=None):
        self.parent = parent
        self.settings = settings_manager
        self.main_app = main_app
        self.pentestgpt_process = None
        
        # Create main frame
        self.main_frame = ttk.Frame(parent)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Create PentestGPT interface
        self._create_pentestgpt_interface()
    
    def _create_pentestgpt_interface(self):
        """Create PentestGPT control and configuration interface."""
        
        # PentestGPT configuration frame
        config_frame = ttk.LabelFrame(self.main_frame, text="PentestGPT Configuration", padding=10)
        config_frame.pack(fill=tk.X, pady=(0, 10))
        
        # Configuration options
        options_frame = ttk.Frame(config_frame)
        options_frame.pack(fill=tk.X, pady=10)
        
        # Model selection
        ttk.Label(options_frame, text="Model:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.pentestgpt_model_var = tk.StringVar(value="gpt-4")
        model_combo = ttk.Combobox(options_frame, textvariable=self.pentestgpt_model_var, 
                                  values=["gpt-4", "gpt-3.5-turbo", "claude-3", "local-llm"], width=15)
        model_combo.grid(row=0, column=1, padx=10)
        
        # Target type
        ttk.Label(options_frame, text="Target Type:").grid(row=0, column=2, sticky=tk.W, padx=(20, 0))
        self.target_type_var = tk.StringVar(value="web")
        target_combo = ttk.Combobox(options_frame, textvariable=self.target_type_var,
                                   values=["web", "network", "mobile", "cloud", "iot"], width=15)
        target_combo.grid(row=0, column=3, padx=10)
        
        # Scan depth
        ttk.Label(options_frame, text="Scan Depth:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.scan_depth_var = tk.StringVar(value="medium")
        depth_combo = ttk.Combobox(options_frame, textvariable=self.scan_depth_var,
                                  values=["light", "medium", "deep", "comprehensive"], width=15)
        depth_combo.grid(row=1, column=1, padx=10)
        
        # Output directory
        ttk.Label(options_frame, text="Output Dir:").grid(row=1, column=2, sticky=tk.W, padx=(20, 0))
        self.output_dir_var = tk.StringVar(value="./pentest_results")
        ttk.Entry(options_frame, textvariable=self.output_dir_var, width=25).grid(row=1, column=3, padx=(10, 5))
        ttk.Button(options_frame, text="Browse", command=self._browse_output_dir).grid(row=1, column=4, padx=5)
        
        # API Configuration
        api_frame = ttk.LabelFrame(config_frame, text="API Configuration", padding=5)
        api_frame.pack(fill=tk.X, pady=(0, 10))
        
        # OpenAI API Key
        ttk.Label(api_frame, text="OpenAI API Key:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.openai_key_var = tk.StringVar()
        ttk.Entry(api_frame, textvariable=self.openai_key_var, show="*", width=40).grid(row=0, column=1, padx=10)
        
        # Claude API Key
        ttk.Label(api_frame, text="Claude API Key:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.claude_key_var = tk.StringVar()
        ttk.Entry(api_frame, textvariable=self.claude_key_var, show="*", width=40).grid(row=1, column=1, padx=10)
        
        # Control buttons
        control_frame = ttk.Frame(config_frame)
        control_frame.pack(fill=tk.X, pady=10)
        
        self.pentestgpt_start_btn = ttk.Button(control_frame, text="Start PentestGPT", 
                      command=self._start_pentestgpt)
        self.pentestgpt_start_btn.pack(side=tk.LEFT, padx=5)
        
        self.pentestgpt_stop_btn = ttk.Button(control_frame, text="Stop PentestGPT", 
                     command=self._stop_pentestgpt, state="disabled")
        self.pentestgpt_stop_btn.pack(side=tk.LEFT, padx=5)
        
        # Configuration management
        config_btn_frame = ttk.Frame(control_frame)
        config_btn_frame.pack(side=tk.RIGHT, padx=5)
        
        ttk.Button(config_btn_frame, text="Save Config", 
                  command=self._save_pentestgpt_config).pack(side=tk.LEFT, padx=2)
        ttk.Button(config_btn_frame, text="Load Config", 
                  command=self._load_pentestgpt_config).pack(side=tk.LEFT, padx=2)
        ttk.Button(config_btn_frame, text="Configure", 
                  command=self._configure_pentestgpt).pack(side=tk.LEFT, padx=2)
        
        # Output area
        output_frame = ttk.LabelFrame(self.main_frame, text="PentestGPT Output", padding=10)
        output_frame.pack(fill=tk.BOTH, expand=True, pady=10)
        
        self.pentestgpt_output = tk.Text(output_frame, height=20, wrap=tk.WORD)
        pentestgpt_scrollbar = ttk.Scrollbar(output_frame, orient=tk.VERTICAL, command=self.pentestgpt_output.yview)
        self.pentestgpt_output.configure(yscrollcommand=pentestgpt_scrollbar.set)
        self.pentestgpt_output.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        pentestgpt_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # Status bar
        self.pentestgpt_status_var = tk.StringVar(value="PentestGPT Status: Ready")
        ttk.Label(output_frame, textvariable=self.pentestgpt_status_var).pack(anchor=tk.W, pady=(5, 0))
    
    def _start_pentestgpt(self):
        """Start PentestGPT process."""
        self.pentestgpt_output.insert(tk.END, "Starting PentestGPT...\n")
        self.pentestgpt_output.insert(tk.END, f"Model: {self.pentestgpt_model_var.get()}\n")
        self.pentestgpt_output.insert(tk.END, f"Target Type: {self.target_type_var.get()}\n")
        self.pentestgpt_output.insert(tk.END, f"Scan Depth: {self.scan_depth_var.get()}\n\n")
        
        # Check if PentestGPT is available
        pentestgpt_path = Path("experimental_agent/PentestGPT")
        if not pentestgpt_path.exists():
            self.pentestgpt_output.insert(tk.END, "Error: PentestGPT not found in experimental_agent directory\n")
            self.pentestgpt_status_var.set("PentestGPT Status: Error - Not Found")
            return
        
        # This would start the actual PentestGPT process
        # For now, we'll simulate the startup
        self.pentestgpt_output.insert(tk.END, "PentestGPT started successfully!\n")
        self.pentestgpt_output.insert(tk.END, "Ready for penetration testing commands...\n\n")
        
        # Update button states
        self.pentestgpt_start_btn.config(state="disabled")
        self.pentestgpt_stop_btn.config(state="normal")
        self.pentestgpt_status_var.set("PentestGPT Status: Running")
    
    def _stop_pentestgpt(self):
        """Stop PentestGPT process."""
        if self.pentestgpt_process:
            self.pentestgpt_process.terminate()
            self.pentestgpt_process = None
        
        self.pentestgpt_output.insert(tk.END, "PentestGPT stopped.\n\n")
        
        # Update button states
        self.pentestgpt_start_btn.config(state="normal")
        self.pentestgpt_stop_btn.config(state="disabled")
        self.pentestgpt_status_var.set("PentestGPT Status: Stopped")
    
    def _browse_output_dir(self):
        """Browse for output directory."""
        directory = filedialog.askdirectory()
        if directory:
            self.output_dir_var.set(directory)
    
    def _save_pentestgpt_config(self):
        """Save PentestGPT configuration."""
        config = {
            "model": self.pentestgpt_model_var.get(),
            "target_type": self.target_type_var.get(),
            "scan_depth": self.scan_depth_var.get(),
            "output_dir": self.output_dir_var.get(),
            "openai_key": self.openai_key_var.get(),
            "claude_key": self.claude_key_var.get()
        }
        
        try:
            config_file = Path("pentestgpt_config.json")
            import json
            with open(config_file, 'w') as f:
                json.dump(config, f, indent=2)
            self.pentestgpt_output.insert(tk.END, f"Configuration saved to {config_file}\n")
        except Exception as e:
            messagebox.showerror("Save Error", f"Failed to save configuration: {str(e)}")
    
    def _load_pentestgpt_config(self):
        """Load PentestGPT configuration."""
        try:
            config_file = Path("pentestgpt_config.json")
            if not config_file.exists():
                messagebox.showwarning("Load Error", "No configuration file found")
                return
            
            import json
            with open(config_file, 'r') as f:
                config = json.load(f)
            
            # Apply configuration
            self.pentestgpt_model_var.set(config.get("model", "gpt-4"))
            self.target_type_var.set(config.get("target_type", "web"))
            self.scan_depth_var.set(config.get("scan_depth", "medium"))
            self.output_dir_var.set(config.get("output_dir", "./pentest_results"))
            self.openai_key_var.set(config.get("openai_key", ""))
            self.claude_key_var.set(config.get("claude_key", ""))
            
            self.pentestgpt_output.insert(tk.END, f"Configuration loaded from {config_file}\n")
            
        except Exception as e:
            messagebox.showerror("Load Error", f"Failed to load configuration: {str(e)}")
    
    def _configure_pentestgpt(self):
        """Open PentestGPT configuration dialog."""
        PentestGPTConfigDialog(self.parent, self)


class PentestGPTConfigDialog:
    """Dialog for configuring PentestGPT API keys and settings."""
    
    def __init__(self, parent, pentestgpt_tab):
        self.pentestgpt_tab = pentestgpt_tab
        self.dialog = tk.Toplevel(parent)
        self.dialog.title("PentestGPT Configuration")
        self.dialog.geometry("500x400")
        self.dialog.transient(parent)
        self.dialog.grab_set()
        
        # Create configuration interface
        self._create_config_interface()
        
        # Center dialog
        self.dialog.update_idletasks()
        x = (self.dialog.winfo_screenwidth() // 2) - (self.dialog.winfo_width() // 2)
        y = (self.dialog.winfo_screenheight() // 2) - (self.dialog.winfo_height() // 2)
        self.dialog.geometry(f"+{x}+{y}")
    
    def _create_config_interface(self):
        """Create the configuration interface."""
        main_frame = ttk.Frame(self.dialog)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=20)
        
        # API Keys section
        api_frame = ttk.LabelFrame(main_frame, text="API Configuration", padding=10)
        api_frame.pack(fill=tk.X, pady=(0, 20))
        
        # OpenAI API Key
        ttk.Label(api_frame, text="OpenAI API Key:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.openai_entry = ttk.Entry(api_frame, show="*", width=40)
        self.openai_entry.grid(row=0, column=1, padx=(10, 0), pady=5)
        self.openai_entry.insert(0, self.pentestgpt_tab.openai_key_var.get())
        
        # Claude API Key
        ttk.Label(api_frame, text="Claude API Key:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.claude_entry = ttk.Entry(api_frame, show="*", width=40)
        self.claude_entry.grid(row=1, column=1, padx=(10, 0), pady=5)
        self.claude_entry.insert(0, self.pentestgpt_tab.claude_key_var.get())
        
        # Advanced Settings section
        advanced_frame = ttk.LabelFrame(main_frame, text="Advanced Settings", padding=10)
        advanced_frame.pack(fill=tk.X, pady=(0, 20))
        
        # Max concurrent scans
        ttk.Label(advanced_frame, text="Max Concurrent Scans:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.max_scans_var = tk.StringVar(value="3")
        ttk.Entry(advanced_frame, textvariable=self.max_scans_var, width=10).grid(row=0, column=1, padx=(10, 0), pady=5)
        
        # Timeout settings
        ttk.Label(advanced_frame, text="Request Timeout (s):").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.timeout_var = tk.StringVar(value="30")
        ttk.Entry(advanced_frame, textvariable=self.timeout_var, width=10).grid(row=1, column=1, padx=(10, 0), pady=5)
        
        # Buttons
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(20, 0))
        
        ttk.Button(button_frame, text="Save", command=self._save_config).pack(side=tk.RIGHT, padx=(10, 0))
        ttk.Button(button_frame, text="Cancel", command=self._cancel).pack(side=tk.RIGHT)
    
    def _save_config(self):
        """Save the configuration."""
        try:
            # Update the main tab variables
            self.pentestgpt_tab.openai_key_var.set(self.openai_entry.get())
            self.pentestgpt_tab.claude_key_var.set(self.claude_entry.get())
            
            # Save to file
            config = {
                "openai_key": self.openai_entry.get(),
                "claude_key": self.claude_entry.get(),
                "max_scans": self.max_scans_var.get(),
                "timeout": self.timeout_var.get()
            }
            
            import json
            config_file = Path("pentestgpt_advanced_config.json")
            with open(config_file, 'w') as f:
                json.dump(config, f, indent=2)
            
            messagebox.showinfo("Success", "Configuration saved successfully!")
            self.dialog.destroy()
            
        except Exception as e:
            messagebox.showerror("Save Error", f"Failed to save configuration: {str(e)}")
    
    def _cancel(self):
        """Cancel the configuration dialog."""
        self.dialog.destroy()