run.py

"""
Python-C Cross-Language Vulnerability Analyzer - Quick Run Script

运行方式:
    python run.py                        # 默认模式：多Agent分析（包含Fuzz测试和运行时插桩）
    python run.py --mode single          # 单Agent模式
    python run.py --help                 # 显示帮助信息
"""

import sys
import os
import argparse

# 添加源码路径到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

# 加载 .env 文件
from dotenv import load_dotenv
load_dotenv()


def run_analysis(mode=None, max_workers=None, request_interval=None, 
                 dataset_path=None, output_path=None):
    """运行漏洞分析"""
    from tqdm import tqdm
    from pyc_bug_analyzer import Config
    from pyc_bug_analyzer.analyzers import SingleAgentAnalyzer, MultiAgentAnalyzer
    
    # 创建Config获取默认值（从环境变量读取）
    config = Config()
    
    # 使用命令行参数或Config默认值，默认使用多Agent模式
    actual_mode = mode if mode is not None else "multi"
    actual_max_workers = max_workers if max_workers is not None else config.max_workers
    actual_request_interval = request_interval if request_interval is not None else config.request_interval
    actual_dataset_path = dataset_path if dataset_path is not None else os.path.join(os.path.dirname(__file__), "datasets", "PyCBench")
    actual_output_path = output_path if output_path is not None else os.path.join(os.path.dirname(__file__), "output")
    
    # 更新config
    config.max_workers = actual_max_workers
    
    print(f"Python-C Cross-Language Vulnerability Analyzer")
    print(f"Mode: {actual_mode}")
    print(f"Workers: {actual_max_workers}")
    print(f"Request Interval: {actual_request_interval}s")
    print(f"Dataset: {actual_dataset_path}")
    print(f"Output: {actual_output_path}\n")
    
    if actual_mode == "single":
        analyzer = SingleAgentAnalyzer(config)
        print("Using Single Agent Mode (StaticAnalysisAgent)")
    else:
        analyzer = MultiAgentAnalyzer(config)
        print("Using Multi Agent Mode (4Agent + Runtime Instrumentation)")
    
    analyzer.output_base_dir = actual_output_path
    os.makedirs(actual_output_path, exist_ok=True)
    analyzer.request_interval = actual_request_interval
    
    case_paths = analyzer.get_all_cases(actual_dataset_path)
    total_cases = len(case_paths)
    print(f"Total cases to analyze: {total_cases}")
    
    if total_cases == 0:
        print("\nNo cases found! Please check your DATASET_PATH.")
        return
    
    summary_file = os.path.join(actual_output_path, "all_agent_responses.json")
    
    with tqdm(total=total_cases, desc="Analyzing cases", unit="case") as pbar:
        analyzer.analyze_cases_parallel(
            case_paths,
            progress_callback=lambda: pbar.update(1),
            summary_file=summary_file
        )


def main():
    parser = argparse.ArgumentParser(
        description="Python-C Cross-Language Vulnerability Analyzer",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run.py                        # 多Agent模式分析（默认，包含Fuzz测试和运行时插桩）
  python run.py --mode single          # 单Agent模式分析
  python run.py --dataset /path/to/dataset  # 指定数据集路径
  python run.py --output /path/to/output   # 指定输出目录
  python run.py --workers 8           # 设置并行线程数
        """
    )
    
    parser.add_argument(
        "--mode", 
        choices=["single", "multi"], 
        default=None,
        help="分析模式：single（单Agent）或 multi（多Agent），默认multi（包含Fuzz测试和运行时插桩）"
    )
    parser.add_argument(
        "--dataset", 
        type=str,
        help="数据集路径"
    )
    parser.add_argument(
        "--output", 
        type=str,
        help="输出目录路径"
    )
    parser.add_argument(
        "--workers", 
        type=int, 
        default=None,
        help=f"并行线程数（默认从Config读取，通常为{16}）"
    )
    parser.add_argument(
        "--interval", 
        type=float, 
        default=None,
        help=f"请求间隔（秒，默认从Config读取，通常为{1.0}）"
    )
    
    args = parser.parse_args()
    
    run_analysis(
        mode=args.mode,
        max_workers=args.workers,
        request_interval=args.interval,
        dataset_path=args.dataset,
        output_path=args.output
    )


if __name__ == "__main__":
    main()