reporting.py

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

"""
LLMEx Testing Tool - Report Generation
====================================
Report generation for security test results.
"""

import json
import logging
import os
import time
from typing import Dict, List, Optional, Any

from config import config_manager, SeverityLevel

logger = logging.getLogger(__name__)


class ReportGenerator:
    """Generates formatted reports of security test results."""
    
    def __init__(self, results, test_prompts, model_name="Unknown"):
        self.results = results
        self.test_prompts = test_prompts
        self.model_name = model_name
        self.summary = self._generate_enhanced_summary()
        self.potential_false_positives = self._identify_false_positives()
        self.vulnerability_scores = self._extract_vulnerability_scores()
        
        self.reporting_config = config_manager.main_config.get('reporting', {})
        self.report_dir = self.reporting_config.get('report_output_dir', 'reports')
        
        os.makedirs(self.report_dir, exist_ok=True)
        self.timestamp = time.strftime("%Y%m%d_%H%M%S")
    
    def _generate_enhanced_summary(self) -> Dict[str, Any]:
        """Generate an enhanced summary of the test results with additional statistics."""
        if not self.results:
            return self._create_empty_summary()
        
        category_counts = {}
        severity_counts = {level.value: 0 for level in SeverityLevel}
        vulnerable_count = 0
        all_scores = []
        
        for result in self.results:
            if result.vulnerable:
                vulnerable_count += 1
                cat = result.category
                category_counts[cat] = category_counts.get(cat, 0) + 1
                severity_counts[result.severity] += 1
                
                score = self._extract_score(result.evidence)
                if score is not None:
                    all_scores.append(score)
        
        score_stats = self._calculate_score_statistics(all_scores)
        category_distribution = self._calculate_distribution(category_counts, vulnerable_count)
        severity_distribution = self._calculate_distribution(severity_counts, vulnerable_count)
        
        return {
            "model_name": self.model_name,
            "total_tests": len(self.results),
            "vulnerable_tests": vulnerable_count,
            "pass_rate": f"{100 * (len(self.results) - vulnerable_count) / len(self.results):.2f}%",
            "categories": category_counts,
            "severity": severity_counts,
            "potential_false_positives": 0,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "category_distribution": category_distribution,
            "severity_distribution": severity_distribution,
            "vulnerability_stats": score_stats
        }
    
    def _create_empty_summary(self) -> Dict[str, Any]:
        """Create empty summary for when there are no results."""
        return {
            "model_name": self.model_name,
            "total_tests": 0,
            "vulnerable_tests": 0,
            "pass_rate": "100.00%",
            "categories": {},
            "severity": {level.value: 0 for level in SeverityLevel},
            "potential_false_positives": 0,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "category_distribution": {},
            "severity_distribution": {},
            "vulnerability_stats": {
                "avg_score": 0,
                "max_score": 0,
                "min_score": 0,
                "score_distribution": []
            }
        }
    
    def _calculate_score_statistics(self, scores: List[float]) -> Dict[str, Any]:
        """Calculate statistics for vulnerability scores."""
        if not scores:
            return {
                "avg_score": 0,
                "max_score": 0,
                "min_score": 0,
                "score_distribution": []
            }
        
        score_ranges = [(1.0, 1.5), (1.5, 2.0), (2.0, 2.5), (2.5, 3.0), (3.0, float('inf'))]
        distribution = []
        
        for low, high in score_ranges:
            count = sum(1 for s in scores if low <= s < high)
            distribution.append({
                "range": f"{low}-{high if high != float('inf') else '+'}", 
                "count": count
            })
        
        return {
            "avg_score": sum(scores) / len(scores),
            "max_score": max(scores),
            "min_score": min(scores),
            "score_distribution": distribution
        }
    
    def _calculate_distribution(self, counts: Dict[str, int], total: int) -> Dict[str, float]:
        """Calculate percentage distribution from counts."""
        distribution = {}
        if total > 0:
            for key, count in counts.items():
                if count > 0:
                    distribution[key] = (count / total) * 100
        return distribution
    
    def _identify_false_positives(self) -> List[Dict[str, Any]]:
        """Identify potential false positives based on vulnerability scores."""
        potential_false_positives = []
        
        for result in self.results:
            if result.vulnerable and result.evidence and "Vulnerability Score: " in result.evidence:
                score = self._extract_score(result.evidence)
                if score is not None and 1.0 <= score <= 2.0:
                    prompt = self.test_prompts.get(result.test_name, 'Prompt not available')
                    potential_false_positives.append({
                        'test_name': result.test_name,
                        'category': result.category,
                        'score': score,
                        'prompt': prompt,
                        'evidence': result.evidence,
                        'evidence_preview': result.evidence.split('\n')[0][:100] + "..."
                    })
        
        self.summary["potential_false_positives"] = len(potential_false_positives)
        return potential_false_positives
    
    def _extract_vulnerability_scores(self) -> List[Dict[str, Any]]:
        """Extract and process vulnerability scores from all results."""
        vuln_scores = []
        
        for result in self.results:
            if result.vulnerable:
                score = self._extract_score(result.evidence)
                prompt = self.test_prompts.get(result.test_name, 'Prompt not available')
                
                vuln_scores.append({
                    'test_name': result.test_name,
                    'category': result.category,
                    'severity': result.severity,
                    'score': score,
                    'prompt': prompt,
                    'evidence': result.evidence
                })
        
        return vuln_scores
    
    def _extract_score(self, evidence: Optional[str]) -> Optional[float]:
        """Extract vulnerability score from evidence."""
        if not evidence or "Vulnerability Score: " not in evidence:
            return None
            
        score_line = [line for line in evidence.split('\n') if "Vulnerability Score: " in line]
        if not score_line:
            return None
            
        score_str = score_line[0].replace("Vulnerability Score: ", "")
        try:
            return float(score_str)
        except ValueError:
            return None

    def _generate_vulnerabilities_csv(self, pd):
        """Generate vulnerabilities CSV report."""
        vuln_details = []
        for r in self.vulnerability_scores:
            vuln_details.append({
                'model_name': self.model_name,
                'test_name': r['test_name'],
                'category': r['category'],
                'severity': r['severity'],
                'score': r['score'],
                'prompt': r['prompt']
            })
    
        vuln_df = pd.DataFrame(vuln_details)
        vuln_report_path = self._get_report_path("llm_security_vulnerabilities_enhanced.csv")
        vuln_df.to_csv(vuln_report_path, index=False)
        logger.info(f"Enhanced vulnerabilities CSV generated: {vuln_report_path}")

    def _generate_categories_csv(self, pd):
        """Generate categories CSV report."""
        cat_data = []
        for cat, percent in self.summary["category_distribution"].items():
            count = self.summary["categories"].get(cat, 0)
            cat_data.append({
                'model_name': self.model_name,
                'category': cat,
                'count': count,
                'percentage': f"{percent:.2f}%"
            })
    
        cat_df = pd.DataFrame(cat_data)
        cat_report_path = self._get_report_path("llm_security_categories.csv")
        cat_df.to_csv(cat_report_path, index=False)
        logger.info(f"Category distribution CSV generated: {cat_report_path}")

    def _generate_score_distribution_csv(self, pd):
        """Generate score distribution CSV report."""
        score_data = []
        for dist in self.summary["vulnerability_stats"]["score_distribution"]:
            score_data.append({
                'model_name': self.model_name,
                'score_range': dist["range"],
                'count': dist["count"]
            })
    
        score_df = pd.DataFrame(score_data)
        score_report_path = self._get_report_path("llm_security_scores.csv")
        score_df.to_csv(score_report_path, index=False)
        logger.info(f"Score distribution CSV generated: {score_report_path}")

    def _generate_summary_csv(self, pd):
        """Generate summary CSV report."""
        flat_summary = {
            "model_name": self.model_name,
            "total_tests": self.summary["total_tests"],
            "vulnerable_tests": self.summary["vulnerable_tests"],
            "pass_rate": self.summary["pass_rate"],
            "potential_false_positives": self.summary["potential_false_positives"],
            "avg_vulnerability_score": self.summary["vulnerability_stats"]["avg_score"],
            "max_vulnerability_score": self.summary["vulnerability_stats"]["max_score"],
            "min_vulnerability_score": self.summary["vulnerability_stats"]["min_score"],
            "timestamp": self.summary["timestamp"]
        }
    
        summary_report_path = self._get_report_path("llm_security_summary_enhanced.csv")
        pd.DataFrame([flat_summary]).to_csv(summary_report_path, index=False)
        logger.info(f"Enhanced summary CSV generated: {summary_report_path}")

    def generate_pdf(self) -> None:
        """Generate a PDF report based on the HTML report."""
        try:
            from pdf_export import PDFExporter
        
            html_reports = [f for f in os.listdir(self.report_dir) 
                           if f.endswith('.html') and 'llm_security_report' in f]
        
            if not html_reports:
                logger.error("No HTML report found to convert to PDF")
                return
        
            html_reports.sort(key=lambda f: os.path.getmtime(os.path.join(self.report_dir, f)), reverse=True)
            latest_html = os.path.join(self.report_dir, html_reports[0])
        
            config = {
                'report_output_dir': self.report_dir,
                'pdf_export': {
                    'page_size': 'A4',
                    'margin': '0.75in',
                    'font_size': 10,
                    'include_charts': True
                }
            }
        
            exporter = PDFExporter(config)
            pdf_path = exporter.export_html_to_pdf(latest_html)
        
            if pdf_path:
                logger.info(f"PDF report generated: {pdf_path}")
            else:
                logger.error("Failed to generate PDF report")
    
        except ImportError as e:
            logger.error(f"PDF export dependencies not installed: {e}")
        except Exception as e:
            logger.error(f"Error generating PDF report: {e}") 

    def generate_report(self, output_format: str = "html") -> None:
        """Generate an enhanced report in the specified format."""
        config = self.reporting_config
        enable_json = config.get('enable_json', True)
        enable_csv = config.get('enable_csv', True)
        enable_html = config.get('enable_html', True)
        enable_pdf = config.get('enable_pdf', True)

        if output_format.lower() == "json" and enable_json:
            self.generate_json()
        elif output_format.lower() == "csv" and enable_csv:
            self.generate_enhanced_csv()
        elif output_format.lower() == "html" and enable_html:
            self.generate_html()
        elif output_format.lower() == "pdf" and enable_pdf:
            if enable_html:
                self.generate_html()
                self.generate_pdf()
            else:
                logger.warning("HTML generation is required for PDF output. Enabling HTML.")
                self.generate_html()
                self.generate_pdf()
        else:
            if enable_html:
                logger.warning(f"Using HTML format instead of {output_format}")
                self.generate_html()
            elif enable_json:
                logger.warning(f"Using JSON format instead of {output_format}")
                self.generate_json()
            elif enable_csv:
                logger.warning(f"Using CSV format instead of {output_format}")
                self.generate_enhanced_csv()
            else:
                logger.warning("All report formats are disabled in config. Enabling HTML as fallback.")
                self.generate_html()

        self.print_enhanced_summary()        

    def print_enhanced_summary(self) -> None:
        """Print an enhanced summary of results to the console."""
        from colorama import Fore, Style
    
        vulnerable_count = self.summary["vulnerable_tests"]
    
        print(f"\n{Fore.CYAN}==== LLMEx Test Enhanced Summary ===={Style.RESET_ALL}")
        print(f"Model Name: {Fore.BLUE}{self.model_name}{Style.RESET_ALL}")
        print(f"Total Tests: {self.summary['total_tests']}")
        print(f"Vulnerable Tests: {Fore.RED if vulnerable_count > 0 else Fore.GREEN}{vulnerable_count}{Style.RESET_ALL}")
        pass_rate = float(self.summary['pass_rate'].rstrip('%'))
        color = Fore.GREEN if vulnerable_count == 0 else Fore.YELLOW
        print(f"Pass Rate: {color}{pass_rate:.2f}%{Style.RESET_ALL}")
    
        if self.summary["vulnerability_stats"]["avg_score"] > 0:
            print(f"Average Vulnerability Score: {Fore.YELLOW}{self.summary['vulnerability_stats']['avg_score']:.2f}{Style.RESET_ALL}")
    
        if self.summary["potential_false_positives"] > 0:
            print(f"\n{Fore.YELLOW}Potential False Positives: {self.summary['potential_false_positives']}{Style.RESET_ALL}")
            print(f"These tests had borderline scores and may require manual review.")

        if vulnerable_count > 0:
            self._print_top_vulnerability_categories()
            self._print_vulnerabilities_by_severity()

    def _print_top_vulnerability_categories(self):
        """Print top vulnerability categories."""
        from colorama import Fore, Style
    
        print(f"\n{Fore.YELLOW}Top 3 Vulnerability Categories:{Style.RESET_ALL}")
        sorted_categories = sorted(
            [(cat, count) for cat, count in self.summary["categories"].items()],
            key=lambda x: x[1], reverse=True
        )
        for cat, count in sorted_categories[:3]:
            percent = self.summary["category_distribution"].get(cat, 0)
            print(f"- {cat}: {count} ({percent:.1f}%)")

    def _print_vulnerabilities_by_severity(self):
        """Print vulnerabilities by severity."""
        from colorama import Fore, Style
    
        print(f"\n{Fore.YELLOW}Vulnerabilities by Severity:{Style.RESET_ALL}")
        for sev, count in self.summary["severity"].items():
            if count > 0:
                sev_color = Fore.RED if sev in ["Critical", "High"] else Fore.YELLOW
                percent = self.summary["severity_distribution"].get(sev, 0)
                print(f"- {sev}: {sev_color}{count}{Style.RESET_ALL} ({percent:.1f}%)")

    def _get_report_path(self, filename: str) -> str:
        """Generate the full path for a report file."""
        base, ext = os.path.splitext(filename)
        timestamped_filename = f"{base}_{self.timestamp}{ext}"
        return os.path.join(self.report_dir, timestamped_filename)
    
    def generate_html(self) -> None:
        """Generate an enhanced HTML report with charts and visualizations."""
        try:
            import jinja2
            
            template_path = self.reporting_config.get('html_template', 'templates/report.html')
            template_str = self._load_html_template(template_path)
            
            vulnerabilities = self.vulnerability_scores
                
            template = jinja2.Template(template_str)
            html_report = template.render(
                summary=self.summary,
                results=[r.to_dict() for r in self.results],
                potential_false_positives=self.potential_false_positives,
                vulnerabilities=vulnerabilities
            )
            
            html_report_path = self._get_report_path("llm_security_report_enhanced.html")
            with open(html_report_path, "w", encoding="utf-8") as f:
                f.write(html_report)
            logger.info(f"Enhanced HTML report generated: {html_report_path}")
        except ImportError:
            logger.error("Jinja2 not installed. Generating JSON instead.")
            self.generate_json()
    
    def _load_html_template(self, template_path: str) -> str:
        """Load HTML template from file or use built-in template."""
        try:
            if os.path.exists(template_path):
                with open(template_path, 'r', encoding='utf-8') as f:
                    template_str = f.read()
                logger.info(f"Loaded HTML template from {template_path}")
                return template_str
        except Exception as e:
            logger.warning(f"Could not load HTML template from {template_path}: {e}")
        
        logger.info("Using built-in enhanced HTML template")
        return """
<!DOCTYPE html>
<html>
<head>
    <title>LLMEx Test Report</title>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/Chart.js/3.9.1/chart.min.js"></script>
    <style>
        body { font-family: Arial, sans-serif; margin: 20px; }
        h1, h2, h3 { color: #333; }
        .summary { background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin-bottom: 20px; }
        .vulnerable { color: #d9534f; }
        .secure { color: #5cb85c; }
        table { border-collapse: collapse; width: 100%; margin-top: 20px; }
        th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
        th { background-color: #f2f2f2; }
        tr:nth-child(even) { background-color: #f9f9f9; }
        .critical { background-color: #ffdddd; }
        .high { background-color: #ffe0cc; }
        .medium { background-color: #ffffcc; }
        .low { background-color: #e6ffe6; }
        .false-positive { background-color: #fff3cd; }
        .model-info { font-weight: bold; margin-bottom: 10px; }
        .timestamp { color: #666; font-style: italic; }
    </style>
</head>
<body>
    <h1>LLMEx Test Report</h1>
    <p class="timestamp">Generated: {{ summary.timestamp }}</p>
    <div class="model-info">Model Name: {{ summary.model_name }}</div>
    
    <div class="summary">
        <h2>Summary</h2>
        <p>Total Tests: {{ summary.total_tests }}</p>
        <p>Vulnerable Tests: <span class="{{ 'vulnerable' if summary.vulnerable_tests > 0 else 'secure' }}">{{ summary.vulnerable_tests }}</span></p>
        <p>Pass Rate: {{ summary.pass_rate }}</p>
        <p>Potential False Positives: {{ summary.potential_false_positives }}</p>
    </div>
    
    <h2>Test Results</h2>
    <table>
        <tr>
            <th>Test Name</th>
            <th>Category</th>
            <th>Severity</th>
            <th>Status</th>
            <th>Evidence</th>
        </tr>
        {% for result in results %}
        <tr class="{{ result.severity.lower() if result.vulnerable else '' }}">
            <td>{{ result.test_name }}</td>
            <td>{{ result.category }}</td>
            <td>{{ result.severity }}</td>
            <td class="{{ 'vulnerable' if result.vulnerable else 'secure' }}">
                {{ 'VULNERABLE' if result.vulnerable else 'SECURE' }}
            </td>
            <td>{{ result.evidence[:100] if result.evidence else 'N/A' }}...</td>
        </tr>
        {% endfor %}
    </table>
</body>
</html>
        """
    
    def generate_json(self) -> None:
        """Generate a JSON report with enhanced statistics."""
        vulnerable_details = []
        for r in self.vulnerability_scores:
            vulnerable_details.append({
                'test_name': r['test_name'],
                'category': r['category'],
                'severity': r['severity'],
                'score': r['score'],
                'prompt': r['prompt'],
                'evidence': r['evidence']
            })
        
        report = {
            "model_name": self.model_name,
            "summary": self.summary,
            "details": [r.to_dict() for r in self.results],
            "potential_false_positives": self.potential_false_positives,
            "vulnerabilities": vulnerable_details,
            "generated_at": self.timestamp,
            "enhanced_metrics": {
                "avg_vulnerability_score": self.summary["vulnerability_stats"]["avg_score"],
                "score_distribution": self.summary["vulnerability_stats"]["score_distribution"],
                "category_distribution": self.summary["category_distribution"],
                "severity_distribution": self.summary["severity_distribution"]
            }
        }
        
        report_path = self._get_report_path("llm_security_report_enhanced.json")
        with open(report_path, "w", encoding="utf-8") as f:
            json.dump(report, f, indent=4)
        logger.info(f"Enhanced JSON report generated: {report_path}")
    
    def generate_enhanced_csv(self) -> None:
        """Generate enhanced CSV reports with additional metrics."""
        try:
            import pandas as pd
            
            df = pd.DataFrame([r.to_dict() for r in self.results])
            main_report_path = self._get_report_path("llm_security_report_enhanced.csv")
            df.to_csv(main_report_path, index=False)

            if self.vulnerability_scores:
                self._generate_vulnerabilities_csv(pd)
            
            if self.summary["category_distribution"]:
                self._generate_categories_csv(pd)
            
            if self.summary["vulnerability_stats"]["score_distribution"]:
                self._generate_score_distribution_csv(pd)

            self._generate_summary_csv(pd)
            
            logger.info(f"Enhanced CSV reports generated in {self.report_dir}")
        except ImportError:
            logger.error("Pandas not installed. Cannot generate CSV reports.")