EnhancedBurpGPT/EnhancedBurpGpt.py

# -*- coding: utf-8 -*-

from burp import IBurpExtender
from burp import IContextMenuFactory
from burp import IScannerCheck
from burp import ITab
from javax.swing import JMenuItem, JPanel, JTextArea, JScrollPane, BoxLayout, JTabbedPane, JDialog, JProgressBar, JLabel
from javax.swing import JButton, JTextField, JOptionPane, JSplitPane, JCheckBox, JComboBox
from java.awt import BorderLayout, Dimension, Color
from java.io import PrintWriter
from java.util import ArrayList
import json
import urllib2
import ssl
from javax.swing import SwingUtilities
from java.lang import Thread, Runnable, String, System
from java.util.concurrent import TimeUnit, Future, TimeoutException
from java.util.concurrent import ExecutorService, Executors
import java.text
from java.io import ByteArrayOutputStream, OutputStreamWriter
from java.nio.charset import StandardCharsets
import base64
from javax.swing import DefaultListModel, JList
from javax.swing.event import ListSelectionListener
from javax.swing import BorderFactory, Box
from java.awt import GridBagLayout, GridBagConstraints, Insets
from abc import ABCMeta, abstractmethod

# ============================================================================
# Logging System
# ============================================================================

class LogLevel:
    """Log level constants"""
    DEBUG = 0
    INFO = 1
    WARN = 2
    ERROR = 3
    
    @staticmethod
    def to_string(level):
        """Convert log level to string"""
        levels = ["DEBUG", "INFO", "WARN", "ERROR"]
        if 0 <= level < len(levels):
            return levels[level]
        return "UNKNOWN"
    
    @staticmethod
    def to_color(level):
        """Get color for log level"""
        colors = {
            LogLevel.DEBUG: Color(128, 128, 128),  # Gray
            LogLevel.INFO: Color(0, 128, 0),       # Green
            LogLevel.WARN: Color(255, 165, 0),     # Orange
            LogLevel.ERROR: Color(255, 0, 0)       # Red
        }
        return colors.get(level, Color.BLACK)


class Logger:
    """Enhanced logging system with levels and formatting"""
    
    def __init__(self, log_area, stdout):
        self.log_area = log_area
        self.stdout = stdout
        self.log_level = LogLevel.INFO
        self.max_lines = 1000
        self.line_count = 0
        
    def set_log_level(self, level):
        """Set minimum log level to display"""
        self.log_level = level
        
    def log(self, message, level=LogLevel.INFO):
        """Log message with specified level"""
        if level < self.log_level:
            return
        
        # Format timestamp
        timestamp = java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(java.util.Date())
        level_str = LogLevel.to_string(level)
        formatted = "[{}] [{}] {}".format(timestamp, level_str, message)
        
        # Output to stdout
        self.stdout.println(formatted)
        
        # Output to UI with auto-truncation
        if hasattr(self, 'log_area') and self.log_area:
            self._append_to_log_area(formatted)
    
    def _append_to_log_area(self, message):
        """Append message to log area with auto-truncation"""
        def append():
            try:
                self.log_area.append(message + "\n")
                self.line_count += 1
                
                # Auto-truncate if exceeds max lines
                if self.line_count > self.max_lines:
                    text = self.log_area.getText()
                    lines = text.split("\n")
                    # Keep only the last max_lines
                    truncated = "\n".join(lines[-self.max_lines:])
                    self.log_area.setText(truncated)
                    self.line_count = self.max_lines
                
                # Auto-scroll to bottom
                self.log_area.setCaretPosition(self.log_area.getDocument().getLength())
            except Exception as e:
                self.stdout.println("[-] Error appending to log area: {}".format(str(e)))
        
        SwingUtilities.invokeLater(append)
    
    def debug(self, message):
        """Log debug message"""
        self.log(message, LogLevel.DEBUG)
    
    def info(self, message):
        """Log info message"""
        self.log(message, LogLevel.INFO)
    
    def warn(self, message):
        """Log warning message"""
        self.log(message, LogLevel.WARN)
    
    def error(self, message):
        """Log error message"""
        self.log(message, LogLevel.ERROR)
    
    def clear(self):
        """Clear log area"""
        if hasattr(self, 'log_area') and self.log_area:
            SwingUtilities.invokeLater(lambda: self.log_area.setText(""))
            self.line_count = 0

# ============================================================================
# Content Truncation Utilities
# ============================================================================

class ContentTruncator:
    """Smart content truncation that preserves structure"""
    
    @staticmethod
    def smart_truncate(content, max_length, content_name="content"):
        """
        Intelligently truncate content while preserving structure
        
        Args:
            content: Content to truncate
            max_length: Maximum allowed length
            content_name: Name for logging purposes
            
        Returns:
            Truncated content with notification
        """
        if not content:
            return ""
        
        # Convert to string if needed
        if isinstance(content, (bytes, bytearray)):
            try:
                content = content.decode('utf-8', errors='ignore')
            except:
                content = str(content)
        else:
            content = str(content)
        
        original_length = len(content)
        
        # No truncation needed
        if original_length <= max_length:
            return content
        
        # Try to preserve HTTP structure (headers + partial body)
        headers_end = content.find("\r\n\r\n")
        if headers_end != -1:
            headers = content[:headers_end + 4]
            body = content[headers_end + 4:]
            
            # Calculate space for body
            header_size = len(headers)
            truncation_marker = "\n\n[!!! CONTENT TRUNCATED !!!]\n[Original size: {} chars, showing first {} chars of body]\n[Configure 'Max Request/Response Length' in settings to show more]\n\n".format(
                original_length, 
                max_length - header_size - 200
            )
            
            available_body_space = max_length - header_size - len(truncation_marker)
            
            if available_body_space > 100:
                # Keep beginning and end of body
                keep_size = available_body_space // 2
                truncated_body = body[:keep_size] + "\n\n... [middle content omitted] ...\n\n" + body[-keep_size:]
                
                return headers + truncation_marker + truncated_body
            else:
                # Not enough space, just truncate body start
                return headers + truncation_marker + body[:available_body_space]
        
        # No HTTP structure detected, simple truncation with markers
        keep_size = (max_length - 200) // 2
        truncation_marker = "\n\n[!!! CONTENT TRUNCATED !!!]\n[Original size: {} chars, showing first and last {} chars]\n[Configure 'Max Request/Response Length' in settings to show more]\n\n".format(
            original_length,
            keep_size
        )
        
        return content[:keep_size] + truncation_marker + content[-keep_size:]
    
    @staticmethod
    def truncate_with_marker(content, max_length):
        """Simple truncation with a clear marker"""
        if not content or len(content) <= max_length:
            return content
        
        original_length = len(content)
        marker = "\n\n[!!! TRUNCATED - Original length: {} chars !!!]\n".format(original_length)
        usable_length = max_length - len(marker)
        
        return content[:usable_length] + marker

# ============================================================================
# SSL Context Handler
# ============================================================================

class TrustAllSSLContext:
    def __init__(self):
        pass
        
    @staticmethod
    def create():
        trust_all_context = ssl.create_default_context()
        trust_all_context.check_hostname = False
        trust_all_context.verify_mode = ssl.CERT_NONE
        return trust_all_context

# ============================================================================
# Multi-Provider API Adapter Architecture
# ============================================================================

class APIProvider:
    """Abstract base class for API providers"""
    __metaclass__ = ABCMeta
    
    def __init__(self, api_key, api_url, model, timeout=60):
        self.api_key = api_key
        self.api_url = api_url
        self.model = model
        self.timeout = timeout
        self.disable_ssl = False
        self.log_callback = None
        # Proxy settings
        self.proxy_enabled = False
        self.proxy_type = "HTTP"
        self.proxy_host = ""
        self.proxy_port = ""
        self.proxy_username = ""
        self.proxy_password = ""
    
    def set_proxy(self, proxy_type, host, port, username="", password=""):
        """Configure proxy settings"""
        self.proxy_enabled = True
        self.proxy_type = proxy_type.upper()
        self.proxy_host = host
        self.proxy_port = port
        self.proxy_username = username
        self.proxy_password = password
        self.log("[+] Proxy configured: {}://{}:{}".format(proxy_type, host, port))
        
        # CRITICAL: Set Java proxy properties immediately for SOCKS5
        if self.proxy_type == "SOCKS5":
            self._setup_java_proxy()
    
    def _setup_java_proxy(self):
        """Setup Java system properties for SOCKS5 proxy"""
        try:
            System.setProperty("socksProxyHost", self.proxy_host)
            System.setProperty("socksProxyPort", str(self.proxy_port))
            
            if self.proxy_username and self.proxy_password:
                System.setProperty("java.net.socks.username", self.proxy_username)
                System.setProperty("java.net.socks.password", self.proxy_password)
            
            self.log("[+] Java SOCKS5 proxy properties set: {}:{}".format(self.proxy_host, self.proxy_port))
        except Exception as e:
            self.log("[-] Error setting Java proxy properties: {}".format(str(e)))
    
    def _make_request_with_socks_proxy(self, url, data=None, headers=None):
        """Make HTTP request using Java's URLConnection with SOCKS5 proxy"""
        from java.net import URL, Proxy, SocketAddress, InetSocketAddress
        from java.net import Proxy as JavaProxy
        from java.io import InputStreamReader, BufferedReader
        
        self.log("[+] Using Java URLConnection with SOCKS5 proxy: {}:{}".format(self.proxy_host, self.proxy_port))
        
        try:
            # Create SOCKS proxy
            sock_addr = InetSocketAddress(self.proxy_host, int(self.proxy_port))
            proxy = JavaProxy(JavaProxy.Type.SOCKS, sock_addr)
            
            # Create URL connection
            java_url = URL(url)
            conn = java_url.openConnection(proxy)
            
            # Set headers
            if headers:
                for key, value in headers.items():
                    conn.setRequestProperty(key, value)
            
            # Set timeout
            conn.setConnectTimeout(int(self.timeout * 1000))
            conn.setReadTimeout(int(self.timeout * 1000))
            
            # Send data if POST
            if data:
                conn.setDoOutput(True)
                conn.setRequestMethod("POST")
                out = conn.getOutputStream()
                out.write(data)
                out.close()
            
            # Get response
            response_code = conn.getResponseCode()
            if response_code >= 400:
                self.log("[-] HTTP Error {}: {}".format(response_code, conn.getResponseMessage()))
                raise Exception("HTTP Error {}: {}".format(response_code, conn.getResponseMessage()))
            
            # Read response body
            reader = BufferedReader(InputStreamReader(conn.getInputStream()))
            response_text = ""
            line = reader.readLine()
            while line is not None:
                response_text += line
                line = reader.readLine()
            reader.close()
            
            return response_text
            
        except Exception as e:
            self.log("[-] SOCKS5 proxy error: {}".format(str(e)))
            raise Exception("SOCKS5 proxy error: {}".format(str(e)))
    
    def _create_proxy_opener(self):
        """Create urllib2 opener with proxy support"""
        if not self.proxy_enabled:
            return None
        
        # For HTTP/HTTPS, use urllib2.ProxyHandler
        proxy_protocol = self.proxy_type.lower()
        
        # Build proxy URL
        if self.proxy_username and self.proxy_password:
            proxy_url = "{}://{}:{}@{}:{}".format(
                proxy_protocol,
                self.proxy_username,
                self.proxy_password,
                self.proxy_host,
                self.proxy_port
            )
        else:
            proxy_url = "{}://{}:{}".format(
                proxy_protocol,
                self.proxy_host,
                self.proxy_port
            )
        
        self.log("[+] Using {} proxy: {}:{}".format(proxy_protocol.upper(), self.proxy_host, self.proxy_port))
        
        # Create proxy handler
        proxy_handler = urllib2.ProxyHandler({
            'http': proxy_url,
            'https': proxy_url
        })
        
        if self.disable_ssl:
            ssl_context = TrustAllSSLContext.create()
            https_handler = urllib2.HTTPSHandler(context=ssl_context)
            return urllib2.build_opener(proxy_handler, https_handler)
        else:
            return urllib2.build_opener(proxy_handler)
    
    def set_log_callback(self, callback):
        """Set logging callback function"""
        self.log_callback = callback
    
    def log(self, message):
        """Log message if callback is set"""
        if self.log_callback:
            self.log_callback(message)
    
    @abstractmethod
    def build_request_data(self, prompt, max_tokens):
        """Build request data - format differs per provider"""
        pass
    
    @abstractmethod
    def build_headers(self):
        """Build request headers - authentication differs per provider"""
        pass
    
    @abstractmethod
    def parse_response(self, response_data):
        """Parse response data - format differs per provider"""
        pass
    
    @abstractmethod
    def get_models_url(self):
        """Get models list URL - differs per provider"""
        pass
    
    def send_request(self, prompt, max_tokens):
        """Common request sending flow"""
        try:
            data = self.build_request_data(prompt, max_tokens)
            json_data = json.dumps(data).encode('utf-8')
            
            headers = self.build_headers()
            
            # CRITICAL: Use Java URLConnection for SOCKS5, urllib2 for HTTP/HTTPS
            if self.proxy_enabled and self.proxy_type == "SOCKS5":
                response_text = self._make_request_with_socks_proxy(self.api_url, json_data, headers)
                result = json.loads(response_text)
                return self.parse_response(result)
            
            request = urllib2.Request(url=self.api_url, data=json_data, headers=headers)
            
            # Use proxy if configured
            opener = self._create_proxy_opener()
            if opener:
                response = opener.open(request, timeout=self.timeout)
            elif self.disable_ssl:
                self.log("[*] SSL certificate validation is disabled")
                ssl_context = TrustAllSSLContext.create()
                response = urllib2.urlopen(request, context=ssl_context, timeout=self.timeout)
            else:
                response = urllib2.urlopen(request, timeout=self.timeout)
            
            response_data = response.read()
            response_text = str(response_data)
            result = json.loads(response_text)
            
            return self.parse_response(result)
            
        except Exception as e:
            raise Exception("Error calling {} API: {}".format(self.__class__.__name__, str(e)))
    
    def fetch_models(self):
        """Fetch available models list"""
        try:
            models_url = self.get_models_url()
            if not models_url:
                return self.get_default_models()
            
            # CRITICAL: Use Java URLConnection for SOCKS5, urllib2 for HTTP/HTTPS
            if self.proxy_enabled and self.proxy_type == "SOCKS5":
                headers = self.build_headers()
                response_text = self._make_request_with_socks_proxy(models_url, None, headers)
                result = json.loads(response_text)
                return self.parse_models_response(result)
            
            # 使用完整的认证头来获取模型列表
            headers = self.build_headers()
            request = urllib2.Request(url=models_url, headers=headers)
            
            # Use proxy if configured
            opener = self._create_proxy_opener()
            if opener:
                response = opener.open(request, timeout=self.timeout)
            elif self.disable_ssl:
                ssl_context = TrustAllSSLContext.create()
                response = urllib2.urlopen(request, context=ssl_context, timeout=self.timeout)
            else:
                response = urllib2.urlopen(request, timeout=self.timeout)
            
            response_data = response.read()
            response_text = str(response_data)
            result = json.loads(response_text)
            
            return self.parse_models_response(result)
            
        except Exception as e:
            self.log("[-] Error fetching models: {}".format(str(e)))
            return self.get_default_models()
    
    @abstractmethod
    def parse_models_response(self, response_data):
        """Parse models list response"""
        pass
    
    @abstractmethod
    def get_default_models(self):
        """Get default models list if fetch fails"""
        pass


class OpenAIProvider(APIProvider):
    """OpenAI-compatible API provider (OpenAI, DeepSeek, Ollama, etc.)"""
    
    def build_request_data(self, prompt, max_tokens):
        return {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens
        }
    
    def build_headers(self):
        host = self.api_url.split("://")[-1].split("/")[0]
        return {
            "Host": host,
            "Content-Type": "application/json",
            "Authorization": "Bearer " + self.api_key,
            "Accept": "*/*",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
            "Accept-Language": "zh-CN"
        }
    
    def parse_response(self, response_data):
        return {
            "content": response_data["choices"][0]["message"]["content"],
            "usage": response_data.get("usage", {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0
            })
        }
    
    def get_models_url(self):
        return self.api_url.replace("/v1/chat/completions", "/v1/models")
    
    def parse_models_response(self, response_data):
        if 'data' in response_data:
            return [model.get('id') for model in response_data['data'] if model.get('id')]
        return self.get_default_models()
    
    def get_default_models(self):
        return ["gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"]


class GeminiProvider(APIProvider):
    """Google Gemini API provider"""
    
    def __init__(self, api_key, model="gemini-pro", timeout=60):
        api_url = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent".format(model)
        super(GeminiProvider, self).__init__(api_key, api_url, model, timeout)
    
    def build_request_data(self, prompt, max_tokens):
        return {
            "contents": [{"parts": [{"text": prompt}]}],
            "generationConfig": {
                "maxOutputTokens": max_tokens,
                "temperature": 0.7
            }
        }
    
    def build_headers(self):
        # Fixed: Removed redundant x-goog-api-key header
        # Gemini API uses URL parameter authentication (?key=...), not header authentication
        return {
            "Content-Type": "application/json",
            "Accept": "*/*",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
    
    def parse_response(self, response_data):
        try:
            content = response_data["candidates"][0]["content"]["parts"][0]["text"]
            usage_metadata = response_data.get("usageMetadata", {})
            
            return {
                "content": content,
                "usage": {
                    "prompt_tokens": usage_metadata.get("promptTokenCount", 0),
                    "completion_tokens": usage_metadata.get("candidatesTokenCount", 0),
                    "total_tokens": usage_metadata.get("totalTokenCount", 0)
                }
            }
        except (KeyError, IndexError) as e:
            raise Exception("Failed to parse Gemini response: {}".format(str(e)))
    
    def send_request(self, prompt, max_tokens):
        """Override to use URL parameter authentication"""
        try:
            url_with_key = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}".format(
                self.model, self.api_key)
            
            data = self.build_request_data(prompt, max_tokens)
            json_data = json.dumps(data).encode('utf-8')
            
            headers = {
                "Content-Type": "application/json",
                "Accept": "*/*",
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
            }
            
            # CRITICAL: Use Java URLConnection for SOCKS5
            if self.proxy_enabled and self.proxy_type == "SOCKS5":
                response_text = self._make_request_with_socks_proxy(url_with_key, json_data, headers)
                result = json.loads(response_text)
                return self.parse_response(result)
            
            request = urllib2.Request(url=url_with_key, data=json_data, headers=headers)
            
            # Use proxy if configured
            opener = self._create_proxy_opener()
            if opener:
                response = opener.open(request, timeout=self.timeout)
            elif self.disable_ssl:
                ssl_context = TrustAllSSLContext.create()
                response = urllib2.urlopen(request, context=ssl_context, timeout=self.timeout)
            else:
                response = urllib2.urlopen(request, timeout=self.timeout)
            
            response_data = response.read()
            response_text = str(response_data)
            result = json.loads(response_text)
            
            return self.parse_response(result)
            
        except Exception as e:
            raise Exception("Error calling Gemini API: {}".format(str(e)))
    
    def get_models_url(self):
        # Fixed: Added pageSize parameter to support full model list retrieval
        return "https://generativelanguage.googleapis.com/v1beta/models?key={}&pageSize=100".format(self.api_key)
    
    def fetch_models(self):
        """Override fetch_models for Gemini to handle API call differently"""
        try:
            all_models = []
            page_token = None
            
            while True:
                models_url = self.get_models_url()
                
                # Fixed: Added pagination token support
                if page_token:
                    models_url += "&pageToken={}".format(page_token)
                
                self.log("[+] Gemini models URL: {}".format(models_url[:80] + "..."))
                
                # CRITICAL: Use Java URLConnection for SOCKS5
                if self.proxy_enabled and self.proxy_type == "SOCKS5":
                    response_text = self._make_request_with_socks_proxy(models_url, None, {})
                else:
                    # Fixed: Explicitly set request method to GET
                    # Official API requires: GET https://generativelanguage.googleapis.com/v1beta/models
                    request = urllib2.Request(url=models_url)
                    request.get_method = lambda: 'GET'  # Critical fix: Set request method to GET
                    
                    # Use proxy if configured
                    opener = self._create_proxy_opener()
                    if opener:
                        response = opener.open(request, timeout=self.timeout)
                    elif self.disable_ssl:
                        self.log("[*] SSL verification disabled for Gemini")
                        ssl_context = TrustAllSSLContext.create()
                        response = urllib2.urlopen(request, context=ssl_context, timeout=self.timeout)
                    else:
                        response = urllib2.urlopen(request, timeout=self.timeout)
                    
                    response_data = response.read()
                    response_text = str(response_data)
                self.log("[+] Gemini API 响应长度: {} 字节".format(len(response_text)))
                
                result = json.loads(response_text)
                models = self.parse_models_response(result)
                all_models.extend(models)
                
                # Check if there's a next page
                page_token = result.get('nextPageToken')
                if not page_token:
                    break
            
            return all_models if all_models else self.get_default_models()
            
        except urllib2.HTTPError as e:
            self.log("[-] HTTP 错误 {}: {}".format(e.code, str(e)))
            self.log("[*] 使用默认模型列表")
            return self.get_default_models()
        except urllib2.URLError as e:
            self.log("[-] Gemini API 网络错误: {}".format(str(e)))
            self.log("[*] 使用默认模型列表")
            return self.get_default_models()
        except Exception as e:
            self.log("[-] Gemini获取模型失败: {}".format(str(e)))
            self.log("[*] 使用默认模型列表")
            return self.get_default_models()
    
    def parse_models_response(self, response_data):
        try:
            if 'models' in response_data:
                models = []
                all_models_count = len(response_data['models'])
                self.log("[+] Gemini API 返回了 {} 个模型（总数）".format(all_models_count))
                
                for model in response_data['models']:
                    model_name = model.get('name', '')
                    if model_name.startswith('models/'):
                        model_name = model_name[7:]
                    
                    # Get supported methods
                    supported_methods = model.get('supportedGenerationMethods', [])
                    
                    # Fixed: Check if model is deprecated and skip it
                    deprecation_info = model.get('deprecationInfo', {})
                    if deprecation_info:
                        self.log("[*] 跳过已弃用模型: {}".format(model_name))
                        continue
                    
                    # Only add models that support generateContent
                    if model_name and 'generateContent' in supported_methods:
                        models.append(model_name)
                        self.log("[+] 发现可用模型: {}".format(model_name))
                
                self.log("[+] 筛选后可用的 Gemini 模型数量: {}".format(len(models)))
                return models
            else:
                self.log("[-] 响应中没有 'models' 字段")
                return self.get_default_models()
        except Exception as e:
            self.log("[-] 解析Gemini模型列表出错: {}".format(str(e)))
            return self.get_default_models()
    
    def get_default_models(self):
        return ["gemini-pro", "gemini-pro-vision", "gemini-1.5-pro", "gemini-1.5-flash"]


class ClaudeProvider(APIProvider):
    """Anthropic Claude API provider"""
    
    def __init__(self, api_key, model="claude-3-sonnet-20240229", api_url="https://api.anthropic.com/v1/messages", timeout=60):
        super(ClaudeProvider, self).__init__(api_key, api_url, model, timeout)
        self.anthropic_version = "2023-06-01"
    
    def build_request_data(self, prompt, max_tokens):
        return {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens
        }
    
    def build_headers(self):
        return {
            "Content-Type": "application/json",
            "x-api-key": self.api_key,
            "anthropic-version": self.anthropic_version,
            "Accept": "*/*",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
    
    def parse_response(self, response_data):
        try:
            content = response_data["content"][0]["text"]
            usage = response_data.get("usage", {})
            
            return {
                "content": content,
                "usage": {
                    "prompt_tokens": usage.get("input_tokens", 0),
                    "completion_tokens": usage.get("output_tokens", 0),
                    "total_tokens": usage.get("input_tokens", 0) + usage.get("output_tokens", 0)
                }
            }
        except (KeyError, IndexError) as e:
            raise Exception("Failed to parse Claude response: {}".format(str(e)))
    
    def get_models_url(self):
        return None
    
    def parse_models_response(self, response_data):
        return self.get_default_models()
    
    def get_default_models(self):
        return [
            "claude-3-opus-20240229",
            "claude-3-sonnet-20240229",
            "claude-3-haiku-20240307",
            "claude-2.1"
        ]


class ProviderFactory:
    """Factory for creating API provider instances"""
    
    PROVIDER_TYPES = {
        "OpenAI": OpenAIProvider,
        "Gemini": GeminiProvider,
        "Claude": ClaudeProvider,
        "DeepSeek": OpenAIProvider,
        "Ollama": OpenAIProvider,
        "Custom": OpenAIProvider
    }
    
    @staticmethod
    def create_provider(provider_type, api_key, api_url, model, timeout=60):
        """Create API provider instance"""
        provider_class = ProviderFactory.PROVIDER_TYPES.get(provider_type)
        
        if not provider_class:
            raise ValueError("Unsupported provider type: {}".format(provider_type))
        
        if provider_type == "Gemini":
            return GeminiProvider(api_key, model, timeout)
        elif provider_type == "Claude":
            return ClaudeProvider(api_key, model, api_url, timeout)
        else:
            return provider_class(api_key, api_url, model, timeout)
    
    @staticmethod
    def get_provider_types():
        """Get all supported provider types"""
        return list(ProviderFactory.PROVIDER_TYPES.keys())
    
    @staticmethod
    def get_default_config(provider_type):
        """Get default configuration for provider"""
        configs = {
            "OpenAI": {
                "api_url": "https://api.openai.com/v1/chat/completions",
                "model": "gpt-4o",
                "requires_url": True
            },
            "Gemini": {
                "api_url": "",
                "model": "gemini-pro",
                "requires_url": False
            },
            "Claude": {
                "api_url": "https://api.anthropic.com/v1/messages",
                "model": "claude-3-sonnet-20240229",
                "requires_url": True
            },
            "DeepSeek": {
                "api_url": "https://api.deepseek.com/v1/chat/completions",
                "model": "deepseek-chat",
                "requires_url": True
            },
            "Ollama": {
                "api_url": "http://localhost:11434/v1/chat/completions",
                "model": "llama2",
                "requires_url": True
            },
            "Custom": {
                "api_url": "https://api.example.com/v1/chat/completions",
                "model": "gpt-3.5-turbo",
                "requires_url": True
            }
        }
        return configs.get(provider_type, configs["Custom"])

# ============================================================================
# Burp Extension
# ============================================================================

class BurpExtender(IBurpExtender, IContextMenuFactory, IScannerCheck, ITab):
    def registerExtenderCallbacks(self, callbacks):
        self._callbacks = callbacks
        self._helpers = callbacks.getHelpers()
        self.stdout = PrintWriter(callbacks.getStdout(), True)
        
        # Default configuration
        self.provider_type = "OpenAI"
        self.api_key = "Please enter your API key"
        self.api_url = "https://api.openai.com/v1/chat/completions"
        self.model = "gpt-4o"
        self.max_tokens = 3072
        self.timeout_seconds = 60
        self.disable_ssl_verification = False
        self.max_request_length = 1000
        self.max_response_length = 2000
        
        # Proxy configuration
        self.enable_proxy = False
        self.proxy_type = "HTTP"
        self.proxy_host = "127.0.0.1"
        self.proxy_port = "10809"
        self.proxy_username = ""
        self.proxy_password = ""
        
        self.executor = Executors.newCachedThreadPool()
        
        self.tab = JTabbedPane()
        
        self.log_panel = JPanel(BorderLayout())
        self.log_area = JTextArea()
        self.log_area.setEditable(False)
        log_scroll = JScrollPane(self.log_area)
        self.log_panel.add(log_scroll, BorderLayout.CENTER)
        
        config_panel = self.create_config_panel()
        results_panel = self.create_results_panel()
        
        self.tab.addTab("Configuration", config_panel)
        self.tab.addTab("Analysis Results", results_panel)
        self.tab.addTab("Logs", self.log_panel)
        
        callbacks.setExtensionName("Enhanced BurpGPT")
        callbacks.registerContextMenuFactory(self)
        callbacks.registerScannerCheck(self)
        callbacks.addSuiteTab(self)
        
        self.debug = True
        
        default_prompt = """Answer in Chinese. Please analyze this HTTP request and response:

Request:
{REQUEST}

Response:
{RESPONSE}

Please identify any security issues and suggest fixes."""
        
        self.prompt_area.setText(default_prompt)
        
        System.setProperty("jsse.enableSNIExtension", "false")
        System.setProperty("https.protocols", "TLSv1.2")
        System.setProperty("javax.net.ssl.trustStore", "")
        System.setProperty("javax.net.ssl.trustStorePassword", "")
        
    def create_config_panel(self):
        config_panel = JPanel(GridBagLayout())
        constraints = GridBagConstraints()
        constraints.fill = GridBagConstraints.HORIZONTAL
        constraints.insets = Insets(5, 10, 5, 10)
        
        api_panel = JPanel(GridBagLayout())
        api_panel.setBorder(BorderFactory.createTitledBorder("API Settings"))
        
        # Provider Type Selection
        constraints.gridx = 0
        constraints.gridy = 0
        constraints.gridwidth = 1
        constraints.weightx = 0.2
        api_panel.add(JLabel("Provider Type:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.provider_combo = JComboBox(ProviderFactory.get_provider_types())
        self.provider_combo.setSelectedItem(self.provider_type)
        self.provider_combo.setToolTipText("Select API provider type")
        
        def on_provider_change(event):
            selected = str(self.provider_combo.getSelectedItem())
            config = ProviderFactory.get_default_config(selected)
            self.url_field.setText(config["api_url"])
            self.url_field.setEnabled(config["requires_url"])
            self.model_combo.removeAllItems()
            self.model_combo.addItem(config["model"])
            self.log("[*] Provider changed to: {}".format(selected))
        
        self.provider_combo.addActionListener(on_provider_change)
        api_panel.add(self.provider_combo, constraints)
        
        # API URL
        constraints.gridx = 0
        constraints.gridy = 1
        constraints.weightx = 0.2
        api_panel.add(JLabel("API URL:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.url_field = JTextField(self.api_url, 40)
        self.url_field.setToolTipText("Enter the API endpoint URL")
        api_panel.add(self.url_field, constraints)
        
        # API Key
        constraints.gridx = 0
        constraints.gridy = 2
        constraints.weightx = 0.2
        api_panel.add(JLabel("API Key:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.key_field = JTextField(self.api_key, 40)
        self.key_field.setToolTipText("Enter your API key")
        api_panel.add(self.key_field, constraints)
        
        # Model
        constraints.gridx = 0
        constraints.gridy = 3
        constraints.weightx = 0.2
        api_panel.add(JLabel("Model:"), constraints)
        
        self.model_combo = JComboBox()
        self.model_combo.setEditable(True)
        self.model_combo.setToolTipText("Select or enter the model name to use")
        self.model_combo.addItem(self.model)
        
        fetch_models_button = JButton("Fetch Models")
        fetch_models_button.setToolTipText("Fetch available models from API")
        
        model_panel = JPanel(BorderLayout())
        model_panel.add(self.model_combo, BorderLayout.CENTER)
        model_panel.add(fetch_models_button, BorderLayout.EAST)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        api_panel.add(model_panel, constraints)
        
        # SSL Options
        constraints.gridx = 0
        constraints.gridy = 4
        constraints.weightx = 0.2
        api_panel.add(JLabel("SSL Options:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.disable_ssl_check = JCheckBox("Disable SSL Certificate Validation", self.disable_ssl_verification)
        self.disable_ssl_check.setToolTipText("Enable this if you encounter SSL certificate issues")
        api_panel.add(self.disable_ssl_check, constraints)
        
        # Proxy panel
        from javax.swing import JPasswordField
        proxy_panel = JPanel(GridBagLayout())
        proxy_panel.setBorder(BorderFactory.createTitledBorder("Proxy Settings"))
        
        constraints.gridx = 0
        constraints.gridy = 0
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Enable Proxy:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.enable_proxy_check = JCheckBox("Use proxy for API requests", self.enable_proxy)
        self.enable_proxy_check.setToolTipText("Enable if you need to use a proxy")
        proxy_panel.add(self.enable_proxy_check, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 1
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Proxy Type:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.proxy_type_combo = JComboBox(["HTTP", "HTTPS", "SOCKS5"])
        self.proxy_type_combo.setSelectedItem(self.proxy_type)
        self.proxy_type_combo.setToolTipText("Select proxy type ")
        proxy_panel.add(self.proxy_type_combo, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 2
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Proxy Host:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.proxy_host_field = JTextField(self.proxy_host, 20)
        self.proxy_host_field.setToolTipText("Proxy server address (e.g., 127.0.0.1)")
        proxy_panel.add(self.proxy_host_field, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 3
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Proxy Port:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.proxy_port_field = JTextField(self.proxy_port, 10)
        self.proxy_port_field.setToolTipText("Proxy port ")
        proxy_panel.add(self.proxy_port_field, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 4
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Username (optional):"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.proxy_username_field = JTextField(self.proxy_username, 20)
        self.proxy_username_field.setToolTipText("Leave empty if no authentication required")
        proxy_panel.add(self.proxy_username_field, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 5
        constraints.weightx = 0.2
        proxy_panel.add(JLabel("Password (optional):"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.proxy_password_field = JPasswordField(self.proxy_password, 20)
        self.proxy_password_field.setToolTipText("Leave empty if no authentication required")
        proxy_panel.add(self.proxy_password_field, constraints)
        
        # Limits panel
        limits_panel = JPanel(GridBagLayout())
        limits_panel.setBorder(BorderFactory.createTitledBorder("Limits & Timeouts"))
        
        constraints.gridx = 0
        constraints.gridy = 0
        constraints.weightx = 0.2
        limits_panel.add(JLabel("Timeout (seconds):"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.timeout_field = JTextField(str(self.timeout_seconds), 10)
        self.timeout_field.setToolTipText("Maximum time to wait for API response")
        limits_panel.add(self.timeout_field, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 1
        constraints.weightx = 0.2
        limits_panel.add(JLabel("Max Request Length:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.req_length_field = JTextField(str(self.max_request_length), 10)
        self.req_length_field.setToolTipText("Maximum length of request content to analyze")
        limits_panel.add(self.req_length_field, constraints)
        
        constraints.gridx = 0
        constraints.gridy = 2
        constraints.weightx = 0.2
        limits_panel.add(JLabel("Max Response Length:"), constraints)
        
        constraints.gridx = 1
        constraints.weightx = 0.8
        self.resp_length_field = JTextField(str(self.max_response_length), 10)
        self.resp_length_field.setToolTipText("Maximum length of response content to analyze")
        limits_panel.add(self.resp_length_field, constraints)
        
        # Prompt panel
        prompt_panel = JPanel(GridBagLayout())
        prompt_panel.setBorder(BorderFactory.createTitledBorder("Prompt Template"))
        
        constraints.gridx = 0
        constraints.gridy = 0
        constraints.gridwidth = 2
        constraints.weighty = 1.0
        self.prompt_area = JTextArea(5, 40)
        self.prompt_area.setLineWrap(True)
        self.prompt_area.setWrapStyleWord(True)
        self.prompt_area.setToolTipText("Template for analysis prompt")
        prompt_scroll = JScrollPane(self.prompt_area)
        prompt_panel.add(prompt_scroll, constraints)
        
        # Button panel
        button_panel = JPanel()
        button_panel.setBorder(BorderFactory.createEmptyBorder(10, 0, 10, 0))
        
        def fetch_models(event):
            try:
                provider_type = str(self.provider_combo.getSelectedItem())
                api_key = self.key_field.getText()
                api_url = self.url_field.getText()
                model = str(self.model_combo.getSelectedItem()) if self.model_combo.getSelectedItem() else "gpt-4o"
                
                if not api_key:
                    JOptionPane.showMessageDialog(None, "Please enter API Key first!")
                    return
                
                self.log("[+] Fetching models from {} provider".format(provider_type))
                
                # Create progress dialog
                progress_dialog = JDialog()
                progress_dialog.setTitle("Fetching Models...")
                progress_dialog.setSize(300, 100)
                progress_dialog.setLocationRelativeTo(None)
                progress_dialog.setLayout(BorderLayout())
                
                progress_bar = JProgressBar()
                progress_bar.setIndeterminate(True)
                label = JLabel("Fetching available models...", SwingUtilities.CENTER)
                progress_dialog.add(label, BorderLayout.NORTH)
                progress_dialog.add(progress_bar, BorderLayout.CENTER)
                
                # Create async task for fetching models
                class FetchModelsTask(Runnable):
                    def __init__(self, outer):
                        self.outer = outer
                    
                    def run(self):
                        try:
                            SwingUtilities.invokeLater(lambda: progress_dialog.setVisible(True))
                            
                            provider = ProviderFactory.create_provider(
                                provider_type, api_key, api_url, model, self.outer.timeout_seconds
                            )
                            provider.set_log_callback(self.outer.log)
                            provider.disable_ssl = self.outer.disable_ssl_check.isSelected()
                            
                            # Configure proxy if enabled
                            if self.outer.enable_proxy_check.isSelected():
                                provider.set_proxy(
                                    str(self.outer.proxy_type_combo.getSelectedItem()),
                                    self.outer.proxy_host_field.getText(),
                                    self.outer.proxy_port_field.getText(),
                                    self.outer.proxy_username_field.getText(),
                                    String(self.outer.proxy_password_field.getPassword())
                                )
                            
                            models = provider.fetch_models()
                            
                            def update_models_ui():
                                try:
                                    progress_dialog.dispose()
                                    
                                    if models:
                                        self.outer.model_combo.removeAllItems()
                                        for model_name in models:
                                            self.outer.model_combo.addItem(model_name)
                                        self.outer.log("[+] Successfully fetched {} models".format(len(models)))
                                        JOptionPane.showMessageDialog(None, "Successfully fetched {} models!".format(len(models)))
                                    else:
                                        raise Exception("No models returned")
                                except Exception as e:
                                    error_msg = "Error updating models: {}".format(str(e))
                                    self.outer.log("[-] " + error_msg)
                                    JOptionPane.showMessageDialog(None, error_msg)
                            
                            SwingUtilities.invokeLater(update_models_ui)
                            
                        except Exception as e:
                            error_msg = "Error fetching models: {}".format(str(e))
                            self.outer.log("[-] " + error_msg)
                            SwingUtilities.invokeLater(lambda: progress_dialog.dispose())
                            SwingUtilities.invokeLater(lambda: JOptionPane.showMessageDialog(None, error_msg))
                
                Thread(FetchModelsTask(self)).start()
                
            except Exception as e:
                error_msg = "Error initializing fetch: {}".format(str(e))
                self.log("[-] " + error_msg)
                JOptionPane.showMessageDialog(None, error_msg)
        
        fetch_models_button.addActionListener(fetch_models)
        
        def save_config(event):
            try:
                self.provider_type = str(self.provider_combo.getSelectedItem())
                self.api_url = self.url_field.getText()
                self.api_key = self.key_field.getText()
                self.model = str(self.model_combo.getSelectedItem())
                self.timeout_seconds = int(self.timeout_field.getText())
                self.max_request_length = int(self.req_length_field.getText())
                self.max_response_length = int(self.resp_length_field.getText())
                self.disable_ssl_verification = self.disable_ssl_check.isSelected()
                
                # Save proxy settings
                self.enable_proxy = self.enable_proxy_check.isSelected()
                self.proxy_type = str(self.proxy_type_combo.getSelectedItem())
                self.proxy_host = self.proxy_host_field.getText()
                self.proxy_port = self.proxy_port_field.getText()
                self.proxy_username = self.proxy_username_field.getText()
                self.proxy_password = String(self.proxy_password_field.getPassword())
                
                # Validate required fields
                if not self.api_key or not self.model:
                    JOptionPane.showMessageDialog(None, "API Key and Model cannot be empty!")
                    return
                
                # Warn if using Custom provider with example URL
                if self.provider_type == "Custom":
                    if not self.api_url or self.api_url.strip() == "":
                        JOptionPane.showMessageDialog(None, "Warning: API URL cannot be empty for Custom provider!")
                        return
                    
                    if "api.example.com" in self.api_url:
                        result = JOptionPane.showConfirmDialog(None, 
                            "You are using the example API URL.\n\n" +
                            "Please replace it with your actual API endpoint.\n\n" +
                            "Do you want to save anyway?",
                            "Example URL Detected",
                            JOptionPane.YES_NO_OPTION,
                            JOptionPane.WARNING_MESSAGE)
                        if result != JOptionPane.YES_OPTION:
                            return
                    
                    if self.model == "gpt-3.5-turbo":
                        result = JOptionPane.showConfirmDialog(None,
                            "You are using the default model name 'gpt-3.5-turbo'.\n\n" +
                            "Make sure this model exists in your custom API.\n\n" +
                            "Do you want to save anyway?",
                            "Default Model Name",
                            JOptionPane.YES_NO_OPTION,
                            JOptionPane.WARNING_MESSAGE)
                        if result != JOptionPane.YES_OPTION:
                            return
                    
                JOptionPane.showMessageDialog(None, "Configuration saved successfully!")
                
                self.log("[+] Configuration updated:")
                self.log("  - Provider: {}".format(self.provider_type))
                self.log("  - API URL: {}".format(self.api_url))
                self.log("  - Model: {}".format(self.model))
                self.log("  - SSL Verification: {}".format("Disabled" if self.disable_ssl_verification else "Enabled"))
                self.log("  - Proxy: {}".format("Enabled ({}://{}:{})".format(
                    self.proxy_type, self.proxy_host, self.proxy_port) if self.enable_proxy else "Disabled"))
                
            except Exception as e:
                JOptionPane.showMessageDialog(None, "Error saving configuration: " + str(e))
                self.log("[-] Error saving configuration: {}".format(str(e)))
        
        def reset_config(event):
            if JOptionPane.showConfirmDialog(None, 
                "Are you sure you want to reset all settings to default values?",
                "Confirm Reset",
                JOptionPane.YES_NO_OPTION) == JOptionPane.YES_OPTION:
                self.provider_combo.setSelectedItem("OpenAI")
                self.url_field.setText("https://api.openai.com/v1/chat/completions")
                self.key_field.setText("Please enter your API key")
                self.model_combo.removeAllItems()
                self.model_combo.addItem("gpt-4o")
                self.timeout_field.setText("60")
                self.req_length_field.setText("1000")
                self.resp_length_field.setText("2000")
                self.disable_ssl_check.setSelected(False)
                self.enable_proxy_check.setSelected(False)
                self.proxy_type_combo.setSelectedItem("HTTP")
                self.proxy_host_field.setText("127.0.0.1")
                self.proxy_port_field.setText("10809")
                self.proxy_username_field.setText("")
                self.proxy_password_field.setText("")
                self.prompt_area.setText(self.get_default_prompt())
        
        save_button = JButton("Save Configuration")
        save_button.setToolTipText("Save current configuration")
        reset_button = JButton("Reset to Defaults")
        reset_button.setToolTipText("Reset all settings to default values")
        
        save_button.addActionListener(save_config)
        reset_button.addActionListener(reset_config)
        
        button_panel.add(save_button)
        button_panel.add(Box.createHorizontalStrut(10))
        button_panel.add(reset_button)
        
        # Assemble panels
        main_constraints = GridBagConstraints()
        main_constraints.fill = GridBagConstraints.HORIZONTAL
        main_constraints.insets = Insets(5, 5, 5, 5)
        main_constraints.gridx = 0
        main_constraints.gridy = 0
        main_constraints.weightx = 1.0
        config_panel.add(api_panel, main_constraints)
        
        main_constraints.gridy = 1
        config_panel.add(proxy_panel, main_constraints)
        
        main_constraints.gridy = 2
        config_panel.add(limits_panel, main_constraints)
        
        main_constraints.gridy = 3
        main_constraints.weighty = 1.0
        main_constraints.fill = GridBagConstraints.BOTH
        config_panel.add(prompt_panel, main_constraints)
        
        main_constraints.gridy = 4
        main_constraints.weighty = 0.0
        main_constraints.fill = GridBagConstraints.HORIZONTAL
        config_panel.add(button_panel, main_constraints)
        
        return config_panel
        
    def get_default_prompt(self):
        return """Please analyze this HTTP request and response:

URL: {URL}
Method: {METHOD}

Request:
{REQUEST}

Response:
{RESPONSE}

Please identify any security issues and suggest fixes."""
        
    def create_results_panel(self):
        results_panel = JPanel(BorderLayout())
        
        split_pane = JSplitPane(JSplitPane.VERTICAL_SPLIT)
        split_pane.setDividerLocation(200)
        
        toolbar = JPanel()
        clear_button = JButton("Clear Results")
        search_field = JTextField(20)
        search_button = JButton("Search")
        export_button = JButton("Export Results")
        
        self.list_model = DefaultListModel()
        self.analysis_list = JList(self.list_model)
        analysis_scroll = JScrollPane(self.analysis_list)
        
        self.results_area = JTextArea()
        self.results_area.setEditable(False)
        results_scroll = JScrollPane(self.results_area)
        
        class SelectionListener(ListSelectionListener):
            def __init__(self, outer):
                self.outer = outer
            
            def valueChanged(self, event):
                if not event.getValueIsAdjusting():
                    index = self.outer.analysis_list.getSelectedIndex()
                    if index >= 0:
                        result = self.outer.analysis_results[index]
                        self.outer.display_result_details(result)
        
        self.analysis_list.addListSelectionListener(SelectionListener(self))
        self.analysis_results = []
        
        def clear_results(event):
            self.list_model.clear()
            self.analysis_results = []
            self.results_area.setText("")
        clear_button.addActionListener(clear_results)
        
        def search_results(event):
            search_text = search_field.getText().lower()
            if search_text:
                for i in range(self.list_model.size()):
                    if search_text in self.list_model.getElementAt(i).lower():
                        self.analysis_list.setSelectedIndex(i)
                        self.analysis_list.ensureIndexIsVisible(i)
                        break
        search_button.addActionListener(search_results)
        
        def export_results(event):
            try:
                if self.analysis_results:
                    timestamp = java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(java.util.Date())
                    filename = "gpt_analysis_{}.txt".format(timestamp)
                    from java.io import FileWriter, BufferedWriter
                    writer = BufferedWriter(FileWriter(filename))
                    try:
                        writer.write("Enhanced BurpGPT Analysis Report\n")
                        writer.write("=" * 50 + "\n")
                        writer.write("Export Time: {}\n".format(timestamp))
                        writer.write("Total Results: {}\n".format(len(self.analysis_results)))
                        writer.write("=" * 50 + "\n\n")
                        for index, result in enumerate(self.analysis_results, 1):
                            writer.write("Analysis #{}\n".format(index))
                            writer.write("-" * 30 + "\n")
                            writer.write("Time: {}\n".format(result.time))
                            writer.write("URL: {}\n".format(result.url))
                            writer.write("\nAnalysis Result:\n")
                            writer.write("-" * 30 + "\n")
                            writer.write(result.response)
                            writer.write("\n" + "=" * 50 + "\n\n")
                        JOptionPane.showMessageDialog(None, "Results exported to {}".format(filename))
                        self.log("[+] Exported {} analysis results to {}".format(len(self.analysis_results), filename))
                    finally:
                        writer.close()
            except Exception as e:
                error_msg = "Export failed: {}".format(str(e))
                JOptionPane.showMessageDialog(None, error_msg)
                self.log("[-] " + error_msg)
        export_button.addActionListener(export_results)
        
        toolbar.add(clear_button)
        toolbar.add(search_field)
        toolbar.add(search_button)
        toolbar.add(export_button)
        
        split_pane.setTopComponent(analysis_scroll)
        split_pane.setBottomComponent(results_scroll)
        
        results_panel.add(toolbar, BorderLayout.NORTH)
        results_panel.add(split_pane, BorderLayout.CENTER)
        
        return results_panel
        
    def send_to_gpt(self, invocation):
        try:
            current_time = System.currentTimeMillis()
            if hasattr(self, '_last_trigger_time') and (current_time - self._last_trigger_time < 1000):
                self.log("[*] Ignoring duplicate trigger")
                return
            self._last_trigger_time = current_time
            
            self.log("[+] Send to GPT method called")
            http_msgs = invocation.getSelectedMessages()
            self.log("[+] Selected messages: {}".format(len(http_msgs)))
            
            if http_msgs and len(http_msgs) == 1:
                msg = http_msgs[0]
                url = msg.getUrl().toString()
                self.log("[+] Processing URL: {}".format(url))
                
                progress_dialog = JDialog()
                progress_dialog.setTitle("Analyzing...")
                progress_dialog.setSize(300, 100)
                progress_dialog.setLocationRelativeTo(None)
                progress_dialog.setLayout(BorderLayout())
                
                progress_bar = JProgressBar()
                progress_bar.setIndeterminate(True)
                label = JLabel("GPT is analyzing the request/response...", SwingUtilities.CENTER)
                progress_dialog.add(label, BorderLayout.NORTH)
                progress_dialog.add(progress_bar, BorderLayout.CENTER)
                
                class AsyncTask(Runnable):
                    def __init__(self, outer):
                        self.outer = outer
                    
                    def run(self):
                        response_time = 0
                        try:
                            self.outer.log("[+] AsyncTask started")
                            SwingUtilities.invokeLater(lambda: progress_dialog.setVisible(True))
                            
                            self.outer.log("[+] Creating GPT request with truncation limits")
                            gpt_request = GPTRequest(
                                self.outer._helpers, 
                                msg, 
                                self.outer.model, 
                                self.outer.max_tokens,
                                self.outer.max_request_length,
                                self.outer.max_response_length,
                                self.outer.log
                            )
                            gpt_request.set_prompt(self.outer.prompt_area.getText())
                            
                            self.outer.log("[+] Sending request to API")
                            
                            # Measure response time
                            start_time = System.currentTimeMillis()
                            gpt_response = self.outer.call_gpt_api(gpt_request)
                            end_time = System.currentTimeMillis()
                            response_time = end_time - start_time
                            
                            self.outer.log("[+] Received response from API (took {:.2f}s)".format(response_time / 1000.0))
                            
                            def update_ui():
                                try:
                                    self.outer.log("[+] Updating UI")
                                    progress_dialog.dispose()
                                    
                                    if isinstance(gpt_response, dict):
                                        content = gpt_response.get("content", "")
                                        usage = gpt_response.get("usage", {})
                                        if content:
                                            self.outer.update_results(url, content, usage, response_time)
                                        else:
                                            self.outer.update_results(url, "No valid analysis result received.", {}, response_time)
                                    else:
                                        error_msg = "Error: {}".format(str(gpt_response))
                                        self.outer.update_results(url, error_msg, {}, response_time)
                                    
                                    self.outer.log("[+] UI updated successfully")
                                except Exception as e:
                                    self.outer.log("[-] Error in update_ui: {}".format(str(e)))
                            
                            SwingUtilities.invokeLater(update_ui)
                        except Exception as e:
                            self.outer.log("[-] Error in AsyncTask: {}".format(str(e)))
                            SwingUtilities.invokeLater(lambda: progress_dialog.dispose())
                            SwingUtilities.invokeLater(lambda: JOptionPane.showMessageDialog(None, "Error: {}".format(str(e))))
                
                Thread(AsyncTask(self)).start()
            else:
                self.log("[-] No message selected or multiple messages selected")
        except Exception as e:
            self.log("[-] Error in send_to_gpt: {}".format(str(e)))
            
    def call_gpt_api(self, gpt_request):
        """Call GPT API using the provider adapter"""
        try:
            self.log("[+] Using provider: {}".format(self.provider_type))
            
            provider = ProviderFactory.create_provider(
                self.provider_type,
                self.api_key,
                self.api_url,
                self.model,
                self.timeout_seconds
            )
            provider.set_log_callback(self.log)
            provider.disable_ssl = self.disable_ssl_verification
            
            # Configure proxy if enabled
            if self.enable_proxy:
                provider.set_proxy(
                    self.proxy_type,
                    self.proxy_host,
                    self.proxy_port,
                    self.proxy_username,
                    self.proxy_password
                )
            
            result = provider.send_request(gpt_request.prompt, self.max_tokens)
            return result
            
        except Exception as e:
            self.log("[-] Error calling API: {}".format(str(e)))
            raise Exception("Error calling API: {}".format(str(e)))
    
    def truncate_content(self, content, max_length):
        if not content:
            return ""
        try:
            content_str = self._helpers.bytesToString(content)
        except:
            content_str = str(content)
        
        if len(content_str) <= max_length:
            return content_str
        
        headers_end = content_str.find("\r\n\r\n")
        if headers_end == -1:
            return content_str[:max_length] + "\n... (content truncated)"
        
        headers = content_str[:headers_end]
        body = content_str[headers_end+4:]
        remaining_length = max_length - len(headers) - 50
        
        if remaining_length <= 0:
            return content_str[:max_length] + "\n... (content truncated)"
        
        truncated_body = body[:remaining_length]
        return "{}\r\n\r\n{}\n... (content truncated, total length: {})".format(headers, truncated_body, len(content_str))
    
    def getTabCaption(self):
        return "GPT Analysis"
        
    def getUiComponent(self):
        return self.tab

    def createMenuItems(self, invocation):
        menu_list = ArrayList()
        menu_item = JMenuItem("Send to GPT")
        menu_item.addActionListener(lambda x: self.send_to_gpt(invocation))
        menu_list.add(menu_item)
        return menu_list

    def log(self, message):
        self.stdout.println(message)
        if hasattr(self, 'log_area'):
            SwingUtilities.invokeLater(lambda: self.log_area.append(message + "\n"))
            SwingUtilities.invokeLater(lambda: self.log_area.setCaretPosition(self.log_area.getDocument().getLength()))

    def update_results(self, url, content, usage, response_time=0):
        timestamp = java.text.SimpleDateFormat("HH:mm:ss").format(java.util.Date())
        result = AnalysisResult(timestamp, url, content, usage, response_time)
        self.analysis_results.append(result)
        
        # Include token count in list display
        tokens_str = " [{}T]".format(result.total_tokens) if result.total_tokens > 0 else ""
        self.list_model.addElement("[{}] {}{}".format(timestamp, url, tokens_str))
        
        last_index = self.list_model.size() - 1
        self.analysis_list.setSelectedIndex(last_index)
        self.analysis_list.ensureIndexIsVisible(last_index)
        
        # Log token usage
        if result.total_tokens > 0:
            self.log("[+] Token usage: {} total ({} prompt + {} completion)".format(
                result.total_tokens, result.prompt_tokens, result.completion_tokens))

    def display_result_details(self, result):
        self.results_area.setText("")
        self.results_area.append("="*50 + "\n")
        self.results_area.append("Analysis Time: {}\n".format(result.time))
        self.results_area.append("Target URL: {}\n".format(result.url))
        
        # Display token usage statistics
        if result.total_tokens > 0:
            self.results_area.append("\n[Token Usage]\n")
            self.results_area.append("  Prompt Tokens: {}\n".format(result.prompt_tokens))
            self.results_area.append("  Completion Tokens: {}\n".format(result.completion_tokens))
            self.results_area.append("  Total Tokens: {}\n".format(result.total_tokens))
        
        # Display performance metrics
        if result.response_time > 0:
            self.results_area.append("  Response Time: {:.2f}s\n".format(result.response_time / 1000.0))
        
        self.results_area.append("-"*50 + "\n")
        self.results_area.append(result.response + "\n")
        self.results_area.append("="*50 + "\n")
        self.results_area.setCaretPosition(0)


class GPTRequest:
    def __init__(self, helpers, http_message, model, max_tokens, max_request_length=1000, max_response_length=2000, log_callback=None):
        try:
            request_info = helpers.analyzeRequest(http_message)
            self.url = str(http_message.getUrl())
            self.method = str(request_info.getMethod())
            
            # Get raw request and response
            request_bytes = http_message.getRequest()
            raw_request = helpers.bytesToString(request_bytes)
            response_bytes = http_message.getResponse()
            raw_response = helpers.bytesToString(response_bytes) if response_bytes else ""
            
            # Apply smart truncation to request and response
            self.request = ContentTruncator.smart_truncate(raw_request, max_request_length, "request")
            self.response = ContentTruncator.smart_truncate(raw_response, max_response_length, "response")
            
            # Log truncation info if callback provided
            if log_callback:
                if len(raw_request) > max_request_length:
                    log_callback("[*] Request truncated: {} chars -> {} chars".format(len(raw_request), max_request_length))
                if len(raw_response) > max_response_length:
                    log_callback("[*] Response truncated: {} chars -> {} chars".format(len(raw_response), max_response_length))
            
            self.model = model
            self.max_tokens = max_tokens
            self.prompt = None
            self.log_callback = log_callback
        except Exception as e:
            raise Exception("Error initializing GPTRequest: " + str(e))

    def set_prompt(self, prompt_template):
        """Build prompt from template with variable substitution"""
        try:
            prompt = prompt_template
            prompt = prompt.replace("{URL}", self.url)
            prompt = prompt.replace("{METHOD}", self.method)
            prompt = prompt.replace("{REQUEST}", self.request)
            prompt = prompt.replace("{RESPONSE}", self.response)
            
            # Store the final prompt
            self.prompt = prompt
            
            # Log final prompt size
            if self.log_callback:
                self.log_callback("[*] Final prompt size: {} characters".format(len(prompt)))
            
            return prompt
        except Exception as e:
            raise Exception("Error setting prompt: " + str(e))


class AnalysisResult:
    """Enhanced analysis result with token usage tracking"""
    
    def __init__(self, time, url, response, usage=None, response_time=0):
        self.time = time
        self.url = url
        self.response = response
        self.severity = "Information"
        self.notes = ""
        
        # Token usage statistics
        if usage is None:
            usage = {}
        self.prompt_tokens = usage.get("prompt_tokens", 0)
        self.completion_tokens = usage.get("completion_tokens", 0)
        self.total_tokens = usage.get("total_tokens", 0)
        
        # Performance metrics
        self.response_time = response_time  # in milliseconds
        
    def get_tokens_display(self):
        """Get formatted token usage string"""
        return "Tokens: {} prompt + {} completion = {} total".format(
            self.prompt_tokens,
            self.completion_tokens,
            self.total_tokens
        )
    
    def get_performance_display(self):
        """Get formatted performance string"""
        if self.response_time > 0:
            return "Response time: {:.2f}s".format(self.response_time / 1000.0)
        return ""
        
    def __str__(self):
        return "[{}] {}".format(self.time, self.url)