collect_fuzz_python.py

"""
Script for Python project fuzzing and test template conversion
usage: PYTHONPATH=. python3 fuzz/collect_fuzz_python.py --pipeline all
"""

from pathlib import Path
import ast
import astunparse
import logging
from typing import Optional
import fire
import os
from UniTSyn.frontend.util import wrap_repo, parallel_subprocess
import subprocess
from os.path import join as pjoin, abspath
from tqdm import tqdm
from pathos.multiprocessing import ProcessingPool
import random
from difflib import SequenceMatcher
from itertools import islice
from datetime import datetime
import re

# Import AST-related functionality
from ast_utils import (
    TestFunctionTransformer,
    TestGenTransformer,
    generate_test_template,
)


def build_image(repos: list[str], jobs: int):
    """
    Build Docker images for OSS-Fuzz projects corresponding to each repository

    Args:
        repos (list[str]): List of repository paths
        jobs (int): Number of parallel tasks
    """
    logging.info(f"Building Docker images for {len(repos)} OSS-Fuzz projects")
    log_dir = os.path.abspath("fuzz_pipeline_log")
    os.makedirs(log_dir, exist_ok=True)

    def _build_cmd(path: str):
        project_name = os.path.basename(path.rstrip("/"))
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = os.path.join(log_dir, f"{project_name}_{timestamp}.log")

        logging.info(f"Start building {project_name}, logging to {log_file}")
        return subprocess.Popen(
            f"yes | python3 infra/helper.py build_image {project_name}",
            cwd=os.path.abspath(os.path.join(path, "../../")),
            stdout=open(log_file, "w"),
            stderr=subprocess.STDOUT,
            shell=True,
        )

    _ = parallel_subprocess(repos, jobs, _build_cmd, on_exit=None)


def build_fuzzer(repos: list[str], jobs: int):
    """
    Build fuzzers in parallel for successfully built projects

    Args:
        repos (list[str]): List of repository paths
        jobs (int): Number of parallel tasks
    """
    logging.info(f"Building fuzzers for {len(repos)} OSS-Fuzz projects")
    log_dir = os.path.abspath("fuzz_pipeline_log")
    os.makedirs(log_dir, exist_ok=True)

    def _build_cmd(path: str):
        project_name = os.path.basename(path.rstrip("/"))
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = os.path.join(log_dir, f"{project_name}_fuzzer_{timestamp}.log")

        logging.info(
            f"Start building fuzzers for {project_name}, logging to {log_file}"
        )
        return subprocess.Popen(
            f"python3 infra/helper.py build_fuzzers --sanitizer address {project_name}",
            cwd=os.path.abspath(os.path.join(path, "../../")),
            stdout=open(log_file, "w"),
            stderr=subprocess.STDOUT,
            shell=True,
        )

    _ = parallel_subprocess(repos, jobs, _build_cmd, on_exit=None)


def discover_targets(project_name: str, oss_fuzz_dir: Path) -> list[str]:
    """
    Discover fuzzing targets
    """
    out_dir = oss_fuzz_dir / "build" / "out" / project_name
    targets: list[str] = []

    logging.debug(f"Searching fuzz targets in: {out_dir}")

    if not out_dir.is_dir():
        logging.warning(f"Build output directory for {project_name} does not exist")
        return targets

    # 常见非 fuzz target 的工具列表
    non_target_tools = {
        "llvm-symbolizer",
        "asan_symbolize",
        "msan_symbolize",
        "tsan_symbolize",
        "ubsan_symbolize",
        "clang",
        "clang++",
        "llvm-ar",
        "llvm-nm",
        "llvm-objcopy",
        "llvm-objdump",
        "llvm-ranlib",
        "llvm-readelf",
        "llvm-readobj",
        "llvm-size",
        "llvm-strings",
        "llvm-strip",
        "ld",
        "ld.lld",
        "lld",
        "lld-link",
    }

    try:
        for f in out_dir.iterdir():
            if (
                f.is_file()
                and "." not in f.name
                and os.access(f, os.X_OK)
                and f.name not in non_target_tools  # 排除已知的工具
            ):
                targets.append(f.name)
        logging.info(
            f"Discovered {len(targets)} fuzz targets in {project_name}: {targets}"
        )
    except Exception as e:
        logging.error(
            f"Error discovering targets for {project_name}: {e}", exc_info=True
        )

    return targets


def fuzz_one_target(target: tuple[str, str], timeout: int):
    """
    Perform fuzzing on a single fuzzing target
    """
    repo_path, target_name = target
    project_name = os.path.basename(repo_path)
    oss_fuzz_root = os.path.dirname(os.path.dirname(repo_path))

    input_file_path = pjoin(repo_path, "fuzz_inputs", target_name)
    os.makedirs(os.path.dirname(input_file_path), exist_ok=True)

    logging.info(
        f"Starting fuzzing: project={project_name}, target={target_name}, timeout={timeout}s"
    )
    logging.debug(f"Fuzz output will be saved to: {input_file_path}")

    try:
        with open(input_file_path, "w") as input_file:
            return subprocess.Popen(
                [
                    "bash",
                    "-c",
                    f"python3 infra/helper.py run_fuzzer {project_name} {target_name} -- -max_total_time={timeout}",
                ],
                cwd=oss_fuzz_root,
                stdout=input_file,
                stderr=subprocess.DEVNULL,
            )
    except Exception as e:
        logging.error(
            f"Error starting fuzzer for {project_name}/{target_name}: {e}",
            exc_info=True,
        )
        return None


def fuzz_repos(repos: list[str], jobs: int, timeout: int):
    """
    Perform fuzzing on a set of repositories
    """
    logging.info(f"Discovering fuzz targets for {len(repos)} repositories...")

    targets_list = []
    for repo in repos:
        project_name = os.path.basename(repo)
        oss_fuzz_dir = Path(repo).parent.parent
        targets = discover_targets(project_name, oss_fuzz_dir)
        targets_list.append(targets)

    target_map = {repo: targets for repo, targets in zip(repos, targets_list)}
    all_targets: list[tuple[str, str]] = [
        (k, v) for k, vs in target_map.items() for v in vs
    ]

    logging.info(f"Total fuzz targets discovered: {len(all_targets)}")
    for repo, targets in target_map.items():
        logging.info(f"{os.path.basename(repo)}: {len(targets)} targets")

    for repo in repos:
        os.makedirs(pjoin(repo, "fuzz_inputs"), exist_ok=True)

    logging.info(
        f"Starting parallel fuzzing with {jobs} jobs, timeout={timeout}s per target"
    )
    parallel_subprocess(
        all_targets, jobs, lambda p: fuzz_one_target(p, timeout), on_exit=None
    )


def transform_repos(repos: list[str], jobs: int):
    """
    Generate test templates for all targets

    Args:
        repos (list[str]): List of repository paths
        jobs (int): Number of parallel tasks
    """
    logging.info("Generating test templates")

    def _transform_repo(repo: str):
        project_name = os.path.basename(repo)
        oss_fuzz_dir = Path(repo).parent.parent
        raw_targets = discover_targets(project_name, oss_fuzz_dir)

        # Simply remove "_print1" from target names, don't add any new suffix
        transformed_targets = [t.replace("_print1", "") for t in raw_targets]

        # Remove duplicates
        targets = list(set(transformed_targets))

        # Pass simple target names to generate_test_template
        return [generate_test_template(t, repo) for t in targets]

    with ProcessingPool(jobs) as p:
        return list(p.map(_transform_repo, repos))


def substitute_one_repo(
    repo: str,
    targets: list[tuple[str, str]],  # Each element is (transformed_target, raw_target)
    n_fuzz: int,
    strategy: str,
    max_len: int,
    sim_thresh: float,
):
    """
    Copy files from fuzz target template and generate multiple testgen files based on fuzz inputs
    using AST transformations
    """

    input_dir = pjoin(repo, "fuzz_inputs")
    template_dir = pjoin(repo, "tests-gen")
    os.makedirs(template_dir, exist_ok=True)

    for transformed_target, raw_target in targets:
        # Build template file path using transformed target name
        source_file = pjoin(template_dir, transformed_target + ".py")

        # Build input file path using raw target name
        input_path = pjoin(input_dir, raw_target)

        # Ensure source file exists
        if not os.path.exists(source_file):
            logging.warning(f"Source file not found: {source_file}")
            continue
        if not os.path.exists(input_path):
            logging.warning(f"Input file not found: {input_path}")
            continue

        # Read all valid input data
        valid_inputs = []
        with open(input_path, "rb") as f_input:
            lines = f_input.readlines()
        # File is closed, now process data
        for line in lines:
            # Use errors='replace' to ensure decoding doesn't fail
            decoded = line.decode("utf-8", errors="replace")

            # Only process lines starting with b' or b"
            if decoded.startswith(("b'", 'b"')):
                if decoded.startswith("b'") and decoded.endswith("'\n"):
                    byte_data = line[2:-2]
                elif decoded.startswith('b"') and decoded.endswith('"\n'):
                    byte_data = line[2:-2]
                else:
                    continue

                if 0 < len(byte_data) <= max_len:
                    valid_inputs.append(byte_data)
            # For other lines, if length is within range and doesn't start with b' or b", also consider adding
            elif 0 < len(line) <= max_len:
                valid_inputs.append(line)

        if not valid_inputs:
            # Use transformed_target instead of target_name
            logging.warning(f"No valid inputs found for {transformed_target}")
            continue

        # Use transformed_target instead of target_name
        logging.info(f"Loaded {len(valid_inputs)} inputs for {transformed_target}")
        # Strategy for selecting inputs
        if strategy == "shuffle":
            random.shuffle(valid_inputs)
            inputs = valid_inputs[:n_fuzz]
        elif strategy == "reverse":
            inputs = list(reversed(valid_inputs))[:n_fuzz]
        else:
            inputs = valid_inputs[:n_fuzz]

        # Generate a separate file for each fuzz input (using AST)
        for idx, fuzz_input in enumerate(inputs, start=1):
            with open(source_file, "r") as f_src:
                code = f_src.read()

            try:
                # Parse into AST
                tree = ast.parse(code)

                # Apply transformer
                transformer = TestGenTransformer(idx, fuzz_input)
                new_tree = transformer.visit(tree)
                ast.fix_missing_locations(new_tree)

                # Ensure test function was found and processed
                if not transformer.found_test_function:
                    logging.warning(f"No test_ function found in {source_file}")
                    continue

                # Generate new code
                new_code = astunparse.unparse(new_tree)

                # Use transformed_target instead of target_name
                out_path = pjoin(template_dir, f"{transformed_target}.testgen_{idx}.py")
                with open(out_path, "w") as f_out:
                    f_out.write(new_code)

                # Format code
                try:
                    subprocess.run(["black", out_path], check=False)
                except FileNotFoundError:
                    logging.warning("Black formatter not found, skipping formatting")

            except SyntaxError as e:
                logging.error(f"Syntax error when processing {source_file}: {e}")
            except Exception as e:
                # Use transformed_target instead of target_name
                logging.error(
                    f"Error generating test case for {transformed_target}: {e}"
                )


def testgen_repos(
    repos: list[str],
    jobs: int,
    n_fuzz: int = 100,
    strategy: str = "shuffle",
    max_len: int = 100,
    sim_thresh: float = 0.8,
):
    """
    Generate test cases from fuzzing inputs

    Args:
        repos (list[str]): List of repository paths
        jobs (int): Number of parallel tasks
        n_fuzz (int): Number of inputs to use
        strategy (str): Selection strategy
        max_len (int): Maximum length
        sim_thresh (float): Similarity threshold
    """
    # First get all targets and apply transformation
    target_map = {}
    for repo in repos:
        project_name = os.path.basename(repo)
        oss_fuzz_dir = Path(repo).parent.parent
        raw_targets = discover_targets(project_name, oss_fuzz_dir)

        # Save original target names and transformed target names
        transformed_targets = [t.replace("_print1", "") for t in raw_targets]
        targets = list(zip(transformed_targets, raw_targets))  # (transformed, raw)
        target_map[repo] = targets

    # Process each repository in parallel
    with ProcessingPool(jobs) as p:
        list(
            p.map(
                lambda item: substitute_one_repo(
                    item[0],  # repo path
                    item[1],  # list of (transformed, raw) targets
                    n_fuzz,
                    strategy,
                    max_len,
                    sim_thresh,
                ),
                target_map.items(),
            )
        )


def main(
    repo_id: str = "data/valid_projects.txt",
    repo_root: str = "fuzz/oss-fuzz/projects/",
    timeout: int = 30,
    jobs: int = 8,
    pipeline: str = "all",
    n_fuzz: int = 100,
    strategy: str = "shuffle",
    max_len: int = 100,
    sim_thresh: float = 0.8,
):
    """
    Main function, controlling the entire fuzzing process

    Args:
        repo_id (str): Project ID file path
        repo_root (str): Project root directory
        timeout (int): Timeout duration
        jobs (int): Number of parallel tasks
        pipeline (str): Pipeline type
        n_fuzz (int): Number of inputs to use
        strategy (str): Selection strategy
        max_len (int): Maximum length
        sim_thresh (float): Similarity threshold
    """
    try:
        with open(repo_id, "r") as f:
            repo_id_list = [line.strip() for line in f if line.strip()]
    except FileNotFoundError:
        repo_id_list = [repo_id]

    # Collect repository paths
    repos = []
    for repo_id in repo_id_list:
        repo_path = abspath(os.path.join(repo_root, repo_id))
        if os.path.isdir(repo_path):
            repos.append(repo_path)

    # Execute specified pipeline
    if pipeline == "build_image":
        build_image(repos, jobs)
    elif pipeline == "build_fuzzer":
        build_fuzzer(repos, jobs)
    elif pipeline == "fuzz":
        fuzz_repos(repos, jobs, timeout)
    elif pipeline == "testgen":
        testgen_repos(repos, jobs, n_fuzz, strategy, max_len, sim_thresh)
    elif pipeline == "transform":
        transform_repos(repos, jobs)
    elif pipeline == "all":
        build_image(repos, jobs)
        build_fuzzer(repos, jobs)
        transform_repos(repos, jobs)  # Generate test templates
        fuzz_repos(repos, jobs, timeout)
        testgen_repos(repos, jobs, n_fuzz, strategy, max_len, sim_thresh)
    else:
        logging.error(f"Unknown pipeline: {pipeline}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    fire.Fire(main)