example_pipeline.py

#!/usr/bin/env python3
"""
End-to-End AutoCodeQL Pipeline Example

This script demonstrates the complete workflow from CVE investigation
to CodeQL query generation for ANY CVE.

The system is universal and works with different CVEs across various
projects, programming languages, and vulnerability types.

Usage:
    python example_pipeline.py CVE-YYYY-NNNNN
    
Examples:
    python example_pipeline.py CVE-2021-46143  # libexpat integer overflow
    python example_pipeline.py CVE-2019-9169   # glibc regex vulnerability
    python example_pipeline.py CVE-2021-3156   # sudo heap overflow
"""

import sys
import json
import logging
from pathlib import Path
import argparse

from codeql_generator.budget_tracker import get_budget_tracker

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

logger = logging.getLogger("example_pipeline")


def run_full_pipeline(cve_id: str, output_base_dir: str = "./pipeline_output", model: str = "gpt-4.1", max_iterations: int = 5, db_path: str = None, codebase_path: str = None, diff_link: str = None, investigation_file: str = None, auto_mode: bool = False, thinking_mode: bool = False):
    """Run the complete AutoCodeQL pipeline.
    
    Args:
        cve_id: CVE identifier (e.g., CVE-2021-46143)
        output_base_dir: Base directory for all outputs
        model: OpenAI model to use (default: gpt-4.1)
        max_iterations: Maximum refinement iterations (default: 5)
        db_path: Path to existing CodeQL database (optional)
        codebase_path: Path to source codebase (optional)
        diff_link: URL to the fix commit/PR diff (optional, manual mode)
        investigation_file: Path to pre-existing investigation JSON (skips CVE investigation)
        auto_mode: If True, use only generic templates without CVE-specific hints (for academic evaluation)
        thinking_mode: If True, use extended reasoning with high effort for GPT-5 models
    """
    logger.info("="*70)
    logger.info(f"AutoCodeQL Pipeline for {cve_id}")
    logger.info(f"Model: {model}, Max Iterations: {max_iterations}")
    if thinking_mode:
        logger.info("🧠 THINKING MODE: Extended reasoning enabled (high effort)")
    if diff_link:
        logger.info(f"Manual Diff Link: {diff_link}")
    if investigation_file:
        logger.info(f"Manual Investigation File: {investigation_file}")
    logger.info("="*70)
    
    # Create output directories
    output_dir = Path(output_base_dir) / cve_id
    output_dir.mkdir(parents=True, exist_ok=True)
    
    cvedigger_dir = output_dir / "cvedigger"
    builder_dir = output_dir / "builder"
    db_creator_dir = output_dir / "db_creator"
    codeql_generator_dir = output_dir / "codeql_generator"
    
    for d in [cvedigger_dir, builder_dir, db_creator_dir, codeql_generator_dir]:
        d.mkdir(exist_ok=True)
    
    # ========================================================================
    # STEP 1: CVE Investigation (CVE Digger)
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("STEP 1: CVE Investigation")
    logger.info("="*70)
    
    default_investigation_file = cvedigger_dir / f"{cve_id}_investigation.json"
    use_cache = False
    
    # Check if user provided a manual investigation file
    if investigation_file:
        manual_file = Path(investigation_file)
        if manual_file.exists():
            try:
                with open(manual_file, 'r') as f:
                    manual_data = json.load(f)
                logger.info(f"✓ Loading manual investigation file: {manual_file}")
                import copy
                investigation_state = manual_data
                use_cache = True
                
                # Ensure structure matches what's expected
                # Use deep copies to avoid circular references (investigation_state IS manual_data)
                if 'final_investigation' not in investigation_state:
                    investigation_state['final_investigation'] = copy.deepcopy(manual_data)
                if 'final_report' not in investigation_state:
                    investigation_state['final_report'] = copy.deepcopy(manual_data)
                
                # IMPORTANT: Enrich the investigation with auto-detected strategy AND ground truth
                # This adds proper DataFlow templates for UAF/double-free vulnerabilities
                # and extracts ground truth locations from diffs (if not manually provided)
                try:
                    from codeql_generator.automation import AutomationIntegrator
                    from codeql_generator.config import PipelineConfig
                    config = PipelineConfig(
                        auto_strategy=True, 
                        auto_ground_truth=True,
                        prefer_manual_when_available=True  # Use manual ground truth if provided
                    )
                    # Let the strategy detector decide whether DataFlow is needed
                    # based on vulnerability type (UAF/double-free → DataFlow, 
                    # buffer overflow/bounds check → pattern matching)
                    integrator = AutomationIntegrator(config=config, dataflow_mode=False, auto_mode=auto_mode)
                    
                    # Get the investigation data - make a deep copy to avoid circular references
                    inv_data = copy.deepcopy(investigation_state.get('final_investigation', investigation_state))
                    
                    # Use full enrich_investigation which adds BOTH strategy AND ground truth
                    enriched = integrator.enrich_investigation(inv_data, codebase_path=codebase_path)
                    
                    if enriched.get('query_strategy'):
                        if auto_mode:
                            logger.info("🔬 AUTONOMOUS MODE: Using generic DataFlow templates (no CVE-specific hints)")
                        else:
                            logger.info("🚨 Enriching investigation with DataFlow-based query strategy")
                        logger.info(f"✓ DataFlow strategy added to investigation")
                    
                    if enriched.get('ground_truth_locations'):
                        gt_count = len(enriched['ground_truth_locations'])
                        logger.info(f"✓ Ground truth extracted: {gt_count} locations")
                        for gt in enriched['ground_truth_locations'][:3]:
                            logger.info(f"   - {gt.get('file', 'unknown')}:{gt.get('line', '?')}")
                    else:
                        logger.warning("⚠ No ground truth locations extracted from diffs")
                    
                    investigation_state['final_investigation'] = enriched
                    investigation_state['final_report'] = enriched
                    
                    # Save ENRICHED investigation for other steps
                    # Use a separate enriched file to avoid corrupting the original
                    enriched_file = cvedigger_dir / f"{cve_id}_investigation_enriched.json"
                    try:
                        with open(enriched_file, 'w') as f:
                            json.dump(enriched, f, indent=2)
                        logger.info(f"  Saved enriched investigation to: {enriched_file}")
                    except Exception as write_err:
                        logger.warning(f"  Could not save enriched file: {write_err}")
                    
                    # Also save to default location if different from manual file
                    if manual_file.resolve() != default_investigation_file.resolve():
                        with open(default_investigation_file, 'w') as f:
                            json.dump(enriched, f, indent=2)
                        logger.info(f"  Saved enriched investigation to: {default_investigation_file}")
                except Exception as e:
                    import traceback
                    logger.warning(f"Failed to enrich investigation: {e}")
                    logger.debug(traceback.format_exc())
                    # Fall back to copying original file (only if different paths)
                    if manual_file.resolve() != default_investigation_file.resolve():
                        import shutil
                        shutil.copy(manual_file, default_investigation_file)
                        logger.info(f"  Copied original to: {default_investigation_file}")
                    else:
                        logger.info(f"  Investigation file already at default location, skipping copy")
            except Exception as e:
                logger.error(f"Failed to load manual investigation file: {e}")
                return False
        else:
            logger.error(f"Manual investigation file not found: {manual_file}")
            return False
    
    investigation_file = default_investigation_file
    
    if not use_cache and investigation_file.exists():
        try:
            with open(investigation_file, 'r') as f:
                cached_data = json.load(f)
            
            # Check if verified
            if cached_data.get("verification_status") == "verified":
                logger.info(f"✓ Loading verified cached investigation from {investigation_file}")
                investigation_state = cached_data
                use_cache = True
                
                # Ensure structure matches what's expected - always set final_investigation
                if 'final_investigation' not in investigation_state:
                    # The cached data IS the investigation, wrap it properly
                    investigation_state['final_investigation'] = cached_data
                if 'final_report' not in investigation_state:
                    investigation_state['final_report'] = cached_data
            else:
                logger.info(f"⚠ Cached investigation exists but verification status is '{cached_data.get('verification_status')}'. Re-running...")
        except Exception as e:
            logger.warning(f"Failed to read cache file: {e}")

    if not use_cache:
        try:
            from cvedigger.graph import run_cve_investigation
            
            logger.info(f"Investigating {cve_id}...")
            investigation_state = run_cve_investigation(cve_id, model=model, pr_url=diff_link)
            
            # Save investigation results
            with open(investigation_file, 'w') as f:
                # Save the entire state or just final_report if available
                report_data = investigation_state.get('final_report', investigation_state)
                json.dump(report_data, f, indent=2)
            
            # Ensure structure matches what's expected
            if 'final_report' in investigation_state and 'final_investigation' not in investigation_state:
                 investigation_state['final_investigation'] = investigation_state['final_report']
            elif 'final_investigation' not in investigation_state:
                 # Assume the whole state is the report if not nested
                 investigation_state['final_investigation'] = investigation_state
            
            logger.info(f"✓ Investigation complete")
            diffs_count = len(investigation_state.get('github_diffs', []))
            logger.info(f"  Found {diffs_count} diffs")
            logger.info(f"  Saved to: {investigation_file}")
            
            # Validate vulnerability classification
            try:
                from codeql_generator.analyzer import VulnerabilityAnalyzer
                analyzer = VulnerabilityAnalyzer()
                # Get the report data (handle nested structure)
                report_data = investigation_state.get('final_report', investigation_state)
                vuln_type = analyzer.classify_vulnerability_type(report_data)
                
                if vuln_type in ['improper validation', 'unknown']:
                    logger.warning(f"⚠ Vulnerability classified as '{vuln_type}'. This is too generic.")
                    logger.info("Re-running investigation with feedback to find specific vulnerability type...")
                    
                    # Create feedback state
                    feedback_state = investigation_state.copy()
                    feedback_state['verification_result'] = {
                        'is_valid': False,
                        'feedback': f"The vulnerability was classified as '{vuln_type}'. Please analyze the technical details and root cause more deeply to identify a specific vulnerability type (e.g., integer overflow, buffer overflow, use-after-free). 'Improper validation' is too generic.",
                        'reason': 'Generic vulnerability classification'
                    }
                    feedback_state['verification_status'] = 'failed'
                    # Reset retry count to allow extractor to try again
                    feedback_state['retry_count'] = 0 
                    
                    # Re-run investigation
                    investigation_state = run_cve_investigation(cve_id, initial_state=feedback_state)
                    
                    # Save updated investigation results
                    with open(investigation_file, 'w') as f:
                        report_data = investigation_state.get('final_report', investigation_state)
                        json.dump(report_data, f, indent=2)
                    
                    # Ensure structure matches what's expected
                    if 'final_report' in investigation_state and 'final_investigation' not in investigation_state:
                         investigation_state['final_investigation'] = investigation_state['final_report']
                    elif 'final_investigation' not in investigation_state:
                         investigation_state['final_investigation'] = investigation_state
                    
                    # Re-classify
                    vuln_type = analyzer.classify_vulnerability_type(report_data)
                    logger.info(f"New classification: {vuln_type}")
                    
            except Exception as e:
                logger.warning(f"Failed to validate vulnerability classification: {e}")
            
        except Exception as e:
            logger.error(f"✗ CVE Investigation failed: {e}")
            logger.info("Tip: Make sure SERPER_API_KEY and OPENAI_API_KEY are set in .env")
            return False
    
    # ========================================================================
    # STEP 2: Build Vulnerable Version (Builder)
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("STEP 2: Build Vulnerable Version")
    logger.info("="*70)
    
    # Skip build if database is provided
    if db_path:
        build_system = "cmake"  # Default assumption
        logger.info(f"Skipping build - using provided database: {db_path}")
        codebase_path = None  # Not needed when db_path is provided
    # Use provided codebase path if available
    elif codebase_path:
        build_system = "autotools"  # Default assumption
        logger.info(f"Using provided codebase path: {codebase_path}")
        logger.info(f"Using build system: {build_system}")
    else:
        try:
            from builder import run_builder
            
            logger.info(f"Building vulnerable version...")
            builder_state = run_builder(
                cve_id=cve_id,
                investigation=investigation_state.get('final_report', investigation_state),
                output_dir=str(builder_dir)
            )
            
            if not builder_state.get('build_success'):
                logger.error("✗ Build failed")
                return False
            
            codebase_path = builder_state['build_dir']
            build_system = builder_state.get('build_system')
            
            logger.info(f"✓ Build complete")
            logger.info(f"  Codebase: {codebase_path}")
            logger.info(f"  Build system: {build_system}")
            
        except Exception as e:
            logger.error(f"✗ Builder failed: {e}")
            logger.info("Tip: Check if git is installed and network is available")
            return False
    
    # ========================================================================
    # STEP 3: Create CodeQL Database (DB Creator)
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("STEP 3: Create CodeQL Database")
    logger.info("="*70)
    
    # Use provided database path if available
    if db_path:
        qlpack_path = None
        logger.info(f"Using provided database path: {db_path}")
    else:
        try:
            from db_creator import run_db_creator
            
            logger.info(f"Creating CodeQL database...")
            db_state = run_db_creator(
                cve_id=cve_id,
                codebase_path=codebase_path,
                build_system=build_system
            )
            
            if not db_state.get('db_creation_success'):
                logger.error("✗ Database creation failed")
                return False
            
            db_path = db_state['db_path']
            qlpack_path = db_state.get('qlpack_path')
            
            logger.info(f"✓ Database created")
            logger.info(f"  Database: {db_path}")
            logger.info(f"  qlpack.yml: {qlpack_path}")
            
        except Exception as e:
            logger.error(f"✗ DB Creator failed: {e}")
            logger.info("Tip: Make sure CodeQL CLI is installed and in PATH")
            return False
    
    # ========================================================================
    # STEP 4: Generate CodeQL Query (CodeQL Generator)
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("STEP 4: Generate CodeQL Query")
    logger.info("="*70)
    
    try:
        from codeql_generator import run_codeql_generator
        
        logger.info(f"Generating CodeQL query...")
        generator_state = run_codeql_generator(
            cve_id=cve_id,
            cve_investigation=investigation_state['final_investigation'],
            codebase_path=codebase_path,
            codeql_db_path=db_path,
            max_refinement_iterations=max_iterations,
            model=model,
            output_dir=str(codeql_generator_dir),
            thinking_mode=thinking_mode
        )
        
        final_query = generator_state.get('final_query')
        is_validated = generator_state.get('is_validated')
        metadata = generator_state.get('query_metadata', {})
        
        if final_query:
            # Save final query
            query_file = codeql_generator_dir / f"{cve_id}.ql"
            with open(query_file, 'w') as f:
                f.write(final_query)
            
            logger.info(f"✓ Query generation complete")
            logger.info(f"  Query saved: {query_file}")
            logger.info(f"  Validated: {is_validated}")
            logger.info(f"  Version: v{metadata.get('query_version', 0)}")
            logger.info(f"  F1 Score: {metadata.get('metrics', {}).get('f1_score', 0):.2f}")
        else:
            logger.error("✗ No query generated")
            return False
        
    except Exception as e:
        logger.error(f"✗ CodeQL Generator failed: {e}")
        logger.info("Tip: Make sure OPENAI_API_KEY is set and CodeQL database is valid")
        return False
    
    # ========================================================================
    # PIPELINE COMPLETE
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("✅ PIPELINE COMPLETE")
    logger.info("="*70)
    
    logger.info(f"\n📁 Output Directory: {output_dir}")
    logger.info(f"\n📊 Results:")
    logger.info(f"  1. Investigation: {investigation_file}")
    logger.info(f"  2. Codebase: {codebase_path}")
    logger.info(f"  3. Database: {db_path}")
    logger.info(f"  4. Query: {query_file}")
    logger.info(f"\n🎯 Final CodeQL Query:")
    logger.info(f"  {query_file}")
    logger.info(f"  F1 Score: {metadata.get('metrics', {}).get('f1_score', 0):.2f}")
    logger.info(f"  Validated: {'✅' if is_validated else '❌'}")
    
    # ========================================================================
    # BUDGET SUMMARY
    # ========================================================================
    logger.info("\n" + "="*70)
    logger.info("💰 TOKEN USAGE & COST SUMMARY")
    logger.info("="*70)
    
    budget_tracker = get_budget_tracker()
    budget_tracker.print_summary()
    
    return True


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='Run complete AutoCodeQL pipeline for any CVE',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Integer overflow in libexpat
  python example_pipeline.py CVE-2021-46143
  
  # Regex vulnerability in glibc
  python example_pipeline.py CVE-2019-9169
  
  # Heap overflow in sudo
  python example_pipeline.py CVE-2021-3156
  
  # Custom output directory
  python example_pipeline.py CVE-2021-46143 --output ./my_output

The system is universal and adapts to different:
  - Programming languages (C/C++, Java, Python, JavaScript, etc.)
  - Vulnerability types (overflow, use-after-free, injection, etc.)
  - Build systems (CMake, Autotools, Make, Maven, Gradle, etc.)
  - Project sizes (small libraries to large applications)
        """
    )
    
    parser.add_argument(
        'cve_id',
        help='CVE identifier (e.g., CVE-2021-46143)'
    )
    
    parser.add_argument(
        '--output',
        default='./pipeline_output',
        help='Base output directory (default: ./pipeline_output)'
    )

    parser.add_argument(
        '--model',
        default='gpt-4.1',
        help='OpenAI model to use (default: gpt-4.1)'
    )

    parser.add_argument(
        '--iterations',
        type=int,
        default=5,
        help='Maximum refinement iterations (default: 5)'
    )
    
    parser.add_argument(
        '--db-path',
        default=None,
        help='Path to existing CodeQL database (skips build and DB creation steps)'
    )
    
    parser.add_argument(
        '--codebase-path',
        default=None,
        help='Path to source codebase (skips build step)'
    )
    
    parser.add_argument(
        '--diff-link',
        default=None,
        help='URL to the fix commit/PR diff (manual mode)'
    )
    
    parser.add_argument(
        '--investigation-file',
        default=None,
        help='Path to pre-existing investigation JSON file (skips CVE investigation step)'
    )
    
    parser.add_argument(
        '--auto',
        action='store_true',
        default=False,
        help='Autonomous mode: Use only generic templates without CVE-specific hints. Required for academic evaluation.'
    )
    
    parser.add_argument(
        '--tm',
        action='store_true',
        default=False,
        help='Enable thinking mode (high reasoning effort) for GPT-5.x models'
    )
    
    args = parser.parse_args()
    
    # Validate CVE ID format
    if not args.cve_id.startswith('CVE-'):
        logger.error(f"Invalid CVE ID format: {args.cve_id}")
        logger.error("Expected format: CVE-YYYY-NNNNN (e.g., CVE-2021-46143)")
        sys.exit(1)
    
    # Log auto mode status
    if args.auto:
        logger.info("🔬 AUTONOMOUS MODE: Using only generic templates (no CVE-specific hints)")
    
    if args.tm:
        logger.info("🧠 THINKING MODE: Using high reasoning effort")
    
    # Run pipeline
    success = run_full_pipeline(args.cve_id, args.output, args.model, args.iterations, args.db_path, args.codebase_path, args.diff_link, args.investigation_file, args.auto, args.tm)
    
    if success:
        logger.info("\n✅ Pipeline executed successfully!")
        sys.exit(0)
    else:
        logger.error("\n❌ Pipeline failed. Check logs above for details.")
        sys.exit(1)


if __name__ == '__main__':
    main()