core.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
LLMEx Testing Tool - Core Security Testing Engine
===============================================
Main security testing engine and coordination.
"""

import logging
import sys
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional, Any
from colorama import Fore, Style, Back, init

from config import config_manager
from clients import LLMClient
from testing import create_test_case, TestResult
from reporting import ReportGenerator
from database import DatabaseManager

init()
logger = logging.getLogger(__name__)


class SecurityTester:
    """Main class for running security tests on LLMs with false positive reduction."""
    
    def __init__(self, llm_client: LLMClient) -> None:
        self.llm_client = llm_client
        self.test_cases = []
        self.test_prompts = {}
        
        config_manager.get_refusal_patterns()
        self._load_test_cases()
        
    def _load_test_cases(self) -> None:
        """Load and filter test cases from configuration files."""
        if not config_manager.category_configs:
            loaded = config_manager.load_all_configs()
            if not loaded:
                logger.error("Failed to load configuration files")
                return

        enabled_categories = config_manager.main_config.get('categories', [])
        
        all_tests = config_manager.get_all_tests()
        for test_config in all_tests:
            category_name = test_config.get('category', '').upper()
            
            if enabled_categories and category_name not in enabled_categories:
                logger.debug(f"Skipping test {test_config.get('name', 'Unknown')} as its category {category_name} is not in enabled categories")
                continue
                
            self._add_false_positive_config(test_config)
            
            test_case = create_test_case(self.llm_client, test_config)
            self.test_cases.append(test_case)
            
            self._store_prompt_for_reporting(test_case.name, test_config)
            
        logger.info(f"Loaded {len(self.test_cases)} test cases")
    
    def _add_false_positive_config(self, test_config: Dict[str, Any]) -> None:
        """Add false positive reduction configuration to test config."""
        fp_config = config_manager.main_config.get('security_testing', {}).get('false_positive_reduction', {})
        test_config['false_positive_threshold'] = fp_config.get('false_positive_threshold', 1.0)
        test_config['context_aware_matching'] = fp_config.get('context_aware_matching', True)
    
    def _store_prompt_for_reporting(self, test_name: str, test_config: Dict[str, Any]) -> None:
        """Store prompt for reporting purposes."""
        if 'prompt' in test_config:
            self.test_prompts[test_name] = test_config['prompt']
        elif 'requests' in test_config and test_config['requests']:
            self.test_prompts[test_name] = [req.get('prompt', '') for req in test_config['requests']]
        elif 'conversation' in test_config:
            self.test_prompts[test_name] = [msg.get('content', '') for msg in test_config['conversation'] 
                                            if msg.get('role') == 'user']
    
    def log_configuration_details(self) -> None:
        """Log all configuration parameters being used for the security tests."""
        logger.info("=" * 80)
        logger.info("LLMEx TESTING CONFIGURATION DETAILS")
        logger.info("=" * 80)
    
        self._log_main_config()
        self._log_rate_limiting_config()
        self._log_security_testing_config()
        self._log_api_provider_config()
        self._log_categories_config()
        self._log_reporting_config()
        self._log_logging_config()
        self._log_test_case_config()
        self._log_client_connection_info()
    
        logger.info("=" * 80)
        logger.info("STARTING SECURITY TESTS")
        logger.info("=" * 80)
    
    def _log_main_config(self) -> None:
        """Log main configuration parameters."""
        logger.info("MAIN CONFIGURATION:")
        logger.info(f"- Version: {config_manager.main_config.get('version', 'Not specified')}")
        logger.info(f"- Description: {config_manager.main_config.get('description', 'Not specified')}")
        logger.info(f"- Default Model: {config_manager.main_config.get('default_model', 'Not specified')}")
        logger.info(f"- Default Output Format: {config_manager.main_config.get('default_output_format', 'Not specified')}")
        logger.info(f"- Max Workers: {config_manager.main_config.get('max_workers', 4)}")
        logger.info(f"- Timeout: {config_manager.main_config.get('timeout', 200)}")
    
    def _log_rate_limiting_config(self) -> None:
        """Log rate limiting configuration."""
        rate_limiting = config_manager.main_config.get('rate_limiting', {})
        logger.info("RATE LIMITING CONFIGURATION:")
        logger.info(f"- Enabled: {rate_limiting.get('enabled', False)}")
        logger.info(f"- Max Requests Per Minute: {rate_limiting.get('max_requests_per_minute', 60)}")
    
    def _log_security_testing_config(self) -> None:
        """Log security testing configuration."""
        security_testing = config_manager.main_config.get('security_testing', {})
        logger.info("SECURITY TESTING CONFIGURATION:")
    
        rate_limit_protection = security_testing.get('rate_limit_protection', {})
        logger.info("- Rate Limit Protection:")
        logger.info(f"  - Max Requests Per Minute: {rate_limit_protection.get('max_requests_per_minute', 30)}")
        logger.info(f"  - Cooldown Period: {rate_limit_protection.get('cooldown_period', 5)}")
    
        response_sanitization = security_testing.get('response_sanitization', {})
        logger.info("- Response Sanitization:")
        logger.info(f"  - Max Response Length: {response_sanitization.get('max_response_length', 2000)}")
        logger.info(f"  - Truncate If Exceed: {response_sanitization.get('truncate_if_exceed', False)}")
    
        fp_reduction = security_testing.get('false_positive_reduction', {})
        logger.info("- False Positive Reduction:")
        logger.info(f"  - Enabled: {fp_reduction.get('enabled', True)}")
        logger.info(f"  - False Positive Threshold: {fp_reduction.get('false_positive_threshold', 1.5)}")
        logger.info(f"  - Context Aware Matching: {fp_reduction.get('context_aware_matching', True)}")
        logger.info(f"  - Refusal Detection: {fp_reduction.get('refusal_detection', True)}")
        logger.info(f"  - Refusal Reduction Factor: {fp_reduction.get('refusal_reduction_factor', 0.8)}")
        logger.info(f"  - Refusal Context Reduction Factor: {fp_reduction.get('refusal_context_reduction_factor', 0.5)}")
        logger.info(f"  - Instructional Format Bonus: {fp_reduction.get('instructional_format_bonus', 2.0)}")
        logger.info(f"  - Detail Threshold: {fp_reduction.get('detail_threshold', 500)}")
        logger.info(f"  - Detail Bonus: {fp_reduction.get('detail_bonus', 1.0)}")
    
        instructional_patterns = fp_reduction.get('instructional_patterns', [])
        logger.info("  - Instructional Patterns:")
        for pattern in instructional_patterns:
            logger.info(f"    - {pattern}")
    
    def _log_api_provider_config(self) -> None:
        """Log API provider configuration."""
        api_providers = config_manager.main_config.get('api_providers', {})
        logger.info("API PROVIDERS CONFIGURATION:")
    
        openai_config = api_providers.get('openai', {})
        if openai_config:
            logger.info("- OpenAI:")
            logger.info(f"  - Base URL: {openai_config.get('base_url', 'https://api.openai.com/v1/chat/completions')}")
            logger.info("  - Models:")
            for model in openai_config.get('models', []):
                logger.info(f"    - {model}")
        
            default_params = openai_config.get('default_params', {})
            logger.info("  - Default Parameters:")
            for key, value in default_params.items():
                logger.info(f"    - {key}: {value}")
    
        custom_api = api_providers.get('custom_api', {})
        if custom_api:
            logger.info("- Custom API:")
            headers = custom_api.get('headers', {})
            logger.info("  - Headers:")
            for key, value in headers.items():
                logger.info(f"    - {key}: {value}")
        
            default_params = custom_api.get('default_params', {})
            logger.info("  - Default Parameters:")
            for key, value in default_params.items():
                logger.info(f"    - {key}: {value}")
    
    def _log_categories_config(self) -> None:
        """Log categories configuration."""
        categories = config_manager.main_config.get('categories', [])
        logger.info("CATEGORIES TO TEST:")
        if categories:
            for category in categories:
                logger.info(f"- {category}")
        else:
            logger.info("- All categories (no filter specified)")
    
    def _log_reporting_config(self) -> None:
        """Log reporting configuration."""
        reporting = config_manager.main_config.get('reporting', {})
        logger.info("REPORTING CONFIGURATION:")
        logger.info(f"- Enable HTML: {reporting.get('enable_html', True)}")
        logger.info(f"- Enable CSV: {reporting.get('enable_csv', True)}")
        logger.info(f"- Enable JSON: {reporting.get('enable_json', True)}")
        logger.info(f"- HTML Template: {reporting.get('html_template', 'templates/report.html')}")
        logger.info(f"- Report Output Directory: {reporting.get('report_output_dir', 'reports')}")
        pdf_config = config_manager.main_config.get('reporting', {}).get('pdf_export', {})
        logger.info("PDF REPORTING CONFIGURATION:")
        logger.info(f"- Enabled: {config_manager.main_config.get('reporting', {}).get('enable_pdf', True)}")
        logger.info(f"- Page Size: {pdf_config.get('page_size', 'A4')}")
        logger.info(f"- Margin: {pdf_config.get('margin', '0.75in')}")
        logger.info(f"- Font Size: {pdf_config.get('font_size', 10)}")
        logger.info(f"- Include Charts: {pdf_config.get('include_charts', True)}")
    
    def _log_logging_config(self) -> None:
        """Log logging configuration."""
        logging_config = config_manager.main_config.get('logging', {})
        logger.info("LOGGING CONFIGURATION:")
        logger.info(f"- Level: {logging_config.get('level', 'INFO')}")
        logger.info(f"- File: {logging_config.get('file', 'llm_security_tests.log')}")
        logger.info(f"- Format: {logging_config.get('format', '%(asctime)s - %(name)s - %(levelname)s - %(message)s')}")
    
    def _log_test_case_config(self) -> None:
        """Log test case configuration."""
        logger.info("TEST CASE CONFIGURATIONS:")
        logger.info(f"- Total Test Cases: {len(self.test_cases)}")
    
        category_counts = {}
        for test in self.test_cases:
            cat = test.category.value
            category_counts[cat] = category_counts.get(cat, 0) + 1
    
        for category, count in category_counts.items():
            logger.info(f"- {category}: {count} tests")
    
    def _log_client_connection_info(self) -> None:
        """Log client connection information."""
        client_info = "Using "
        if hasattr(self.llm_client, 'model'):
            client_info += f"OpenAI API with model: {self.llm_client.model}"
        else:
            client_info += f"Custom API"
        logger.info(f"CLIENT CONNECTION: {client_info}")
    
    def log_test_details(self, test_case) -> None:
        """Log detailed information about a specific test case before execution."""
        logger.info("-" * 60)
        logger.info(f"EXECUTING TEST: {test_case.name}")
        logger.info(f"- Category: {test_case.category.value}")
        logger.info(f"- Severity: {test_case.severity.value}")
        logger.info(f"- Description: {test_case.description}")
    
        if 'prompt' in test_case.config:
            logger.info(f"- Test Type: Single Prompt")
            logger.info(f"- Prompt Length: {len(test_case.config['prompt'])} characters")
        elif 'conversation' in test_case.config:
            logger.info(f"- Test Type: Multi-Turn Conversation")
            logger.info(f"- Conversation Turns: {len(test_case.config['conversation'])}")
        elif 'requests' in test_case.config:
            logger.info(f"- Test Type: Multiple Requests")
            logger.info(f"- Request Count: {len(test_case.config['requests'])}")
            logger.info(f"- Concurrent Execution: {test_case.config.get('concurrent', False)}")
    
        logger.info(f"- False Positive Threshold: {test_case.analyzer.fp_threshold}")
        logger.info(f"- Context-Aware Matching: {test_case.analyzer.context_aware}")
    
        logger.info("-" * 60)
    
    def run_all_tests(self) -> List[TestResult]:
        """Run all loaded test cases in parallel and gather their results."""
        self.log_configuration_details()
    
        logger.info("Starting LLMEx testing...")
        results: List[TestResult] = []
        max_workers = config_manager.main_config.get('max_workers', 4)

        if not self.llm_client.check_connection():
            logger.error("Failed to connect to the LLM. Please check your API key and connection settings.")
            return results
    
        logger.info(f"Using {max_workers} worker threads for parallel test execution")

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_test = {}
        
            for test in self.test_cases:
                self.log_test_details(test)
                fut = executor.submit(test.run)
                future_to_test[fut] = test
            
            return self._process_test_results(future_to_test)
    
    def _process_test_results(self, future_to_test) -> List[TestResult]:
        """Process test results with progress bar."""
        try:
            from tqdm import tqdm
        except ImportError:
            tqdm = lambda x, desc="": x
            
        results = []
        for future in tqdm(future_to_test, desc="Running security tests"):
            test = future_to_test[future]
            try:
                result = future.result()
                results.append(result)
                
                self._log_test_result(test, result)
            except Exception as e:
                logger.error(f"Error running test {test.name}: {e}")
                import traceback
                logger.error(traceback.format_exc())
    
        vulnerable_count = sum(1 for r in results if r.vulnerable)
        logger.info("=" * 80)
        logger.info(f"TESTING COMPLETE: {len(results)} tests executed, {vulnerable_count} vulnerabilities found")
        logger.info("=" * 80)
    
        return results
    
    def _log_test_result(self, test, result):
        """Log test result information."""
        if result.vulnerable:
            if hasattr(test, 'vulnerability_score'):
                score_info = f" (Score: {test.vulnerability_score:.2f})"
            else:
                score_info = ""
                
            evidence_preview = ""
            if result.evidence:
                first_line = result.evidence.split('\n')[0][:100]
                evidence_preview = first_line
                
            logger.warning(f"{Back.RED}Vulnerability found in {test.name}{score_info}: "
                           f"{evidence_preview}{Style.RESET_ALL}")
        else:
            logger.info(f"{Back.GREEN}Test passed: {test.name}{Style.RESET_ALL}")
    
    def run_category_tests(self, category: str) -> List[TestResult]:
        """Run tests for a specific category."""
        self.log_configuration_details()
    
        logger.info(f"Starting LLMEx testing for category: {category}...")
        results: List[TestResult] = []

        category_tests = [
            test for test in self.test_cases
            if test.category.name.lower() == category.lower()
        ]
    
        if not category_tests:
            logger.error(f"No tests found for category: {category}")
            return results

        logger.info(f"Found {len(category_tests)} tests for category: {category}")
    
        if not self.llm_client.check_connection():
            logger.error("Failed to connect to the LLM. Please check your API key and connection settings.")
            return results
    
        max_workers = config_manager.main_config.get('max_workers', 4)
        logger.info(f"Using {max_workers} worker threads for parallel test execution")

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_test = {}
        
            for test in category_tests:
                self.log_test_details(test)
                fut = executor.submit(test.run)
                future_to_test[fut] = test
            
            results = self._process_category_test_results(future_to_test, category)

        vulnerable_count = sum(1 for r in results if r.vulnerable)
        logger.info("=" * 80)
        logger.info(f"TESTING COMPLETE FOR CATEGORY {category.upper()}: {len(results)} tests executed, {vulnerable_count} vulnerabilities found")
        logger.info("=" * 80)
    
        return results
    
    def _process_category_test_results(self, future_to_test, category: str) -> List[TestResult]:
        """Process category-specific test results with progress bar."""
        try:
            from tqdm import tqdm
        except ImportError:
            tqdm = lambda x, desc="": x
            
        results = []
        for future in tqdm(future_to_test, desc=f"Running {category} tests"):
            test = future_to_test[future]
            try:
                result = future.result()
                results.append(result)
                
                self._log_test_result(test, result)
            except Exception as e:
                logger.error(f"Error running test {test.name}: {e}")
                import traceback
                logger.error(traceback.format_exc())
        
        return results


def generate_reports(results: List[TestResult], tester: SecurityTester, 
                    output_format: str, model_name: str) -> None:
    """Generate reports based on test results."""
    reporter = ReportGenerator(results, tester.test_prompts, model_name)
    reporter.generate_report(output_format)


def store_results_in_database(db_path: str, results: List[TestResult], 
                             tester: SecurityTester, model_name: str) -> None:
    """Store test results in database."""
    logger.info("Storing test results in database...")
    db = DatabaseManager(db_path)

    vulnerability_scores = extract_vulnerability_scores(results, tester.test_prompts)
    reporter = ReportGenerator(results, tester.test_prompts, model_name)

    run_id = db.store_results(
        results=[r.to_dict() for r in results], 
        summary=reporter.summary,
        vulnerability_scores=vulnerability_scores,
        potential_false_positives=reporter.potential_false_positives,
        model_name=model_name,
        config_hash=None,
        notes=f"Test run with {len(results)} tests"
    )

    logger.info(f"Results stored in database with run ID: {run_id}")

def count_tests_in_categories(categories_filter: Optional[str] = None) -> None:
    """Count and display the number of tests in specified categories."""
    config_manager.load_all_configs()
    
    categories_to_count = []
    if categories_filter:
        if categories_filter.strip():
            categories_to_count = [cat.strip().upper() for cat in categories_filter.split(',')]
            print(f"Counting tests in categories: {categories_filter}")
        else:
            print("No categories specified. Counting tests in ALL categories.")
    else:
        print("No categories specified. Counting tests in ALL categories.")
    
    all_tests = config_manager.get_all_tests()
    
    category_counts = {}
    total_count = 0
    
    for test in all_tests:
        test_category = test.get('category', '').upper()
        if not categories_to_count or test_category in categories_to_count:
            category_counts[test_category] = category_counts.get(test_category, 0) + 1
            total_count += 1
    
    print("\nTest Count Results:")
    print("-" * 50)
    print(f"{'Category':<30} | {'Count':<10}")
    print("-" * 50)
    
    for category, count in sorted(category_counts.items()):
        print(f"{category:<30} | {count:<10}")
    
    print("-" * 50)
    print(f"{'TOTAL':<30} | {total_count:<10}")
    print("-" * 50)
    
    print("\nNote: Counts reflect tests defined in the current configuration directory.")
    print(f"Configuration directory: {config_manager.config_dir}")

def exit_with_results(results: List[TestResult]) -> None:
    """Exit with appropriate code based on test results."""
    vulnerable_count = sum(1 for result in results if result.vulnerable)
    if vulnerable_count > 0:
        logger.warning(f"{vulnerable_count} vulnerabilities found")
        sys.exit(vulnerable_count)
    else:
        logger.info(f"{Back.GREEN}No vulnerabilities found{Style.RESET_ALL}")  
        sys.exit(0)


def extract_vulnerability_scores(results: List[TestResult], 
                               test_prompts: Dict[str, str]) -> List[Dict[str, Any]]:
    """Extract vulnerability scores from test results."""
    vulnerability_scores = []
    for result in results:
        if result.vulnerable:
            score = None
            if result.evidence and "Vulnerability Score: " in result.evidence:
                score_line = [line for line in result.evidence.split('\n') if "Vulnerability Score: " in line]
                if score_line:
                    try:
                        score = float(score_line[0].replace("Vulnerability Score: ", ""))
                    except ValueError:
                        pass
        
            vulnerability_scores.append({
                'test_name': result.test_name,
                'category': result.category,
                'severity': result.severity,
                'score': score,
                'prompt': test_prompts.get(result.test_name, 'Prompt not available'),
                'evidence': result.evidence
            })
    
    return vulnerability_scores