modify_fuzz_files.py

#!/usr/bin/env python3
import os
import re
import shutil
import ast
import astunparse


def process_projects():
    # Read valid projects from data/valid_projects.txt
    with open("data/valid_projects.txt", "r") as f:
        projects = [line.strip() for line in f.readlines() if line.strip()]

    # Base directory containing projects
    base_dir = "fuzz/oss-fuzz/projects"

    for project in projects:
        project_dir = os.path.join(base_dir, project)

        if not os.path.exists(project_dir):
            print(f"Project directory not found: {project_dir}")
            continue

        print(f"Processing project: {project}")

        # Remove _print1.py and .bak files
        for root, dirs, files in os.walk(project_dir):
            for file in files:
                if file.endswith("_print1.py") or file.endswith(".bak.py"):
                    file_path = os.path.join(root, file)
                    os.remove(file_path)
                    print(f"Removed: {file_path}")

        # Find all remaining .py files
        py_files = []
        for root, dirs, files in os.walk(project_dir):
            for file in files:
                if file.endswith(".py"):
                    py_files.append(os.path.join(root, file))

        # Process each .py file
        for py_file in py_files:
            process_py_file(py_file)


class FunctionVisitor(ast.NodeVisitor):
    def __init__(self, target_func):
        self.target_func = target_func
        self.found_node = None
        self.first_param = None

    def visit_FunctionDef(self, node):
        if node.name == self.target_func:
            self.found_node = node
            if node.args.args:
                self.first_param = node.args.args[0].arg
        self.generic_visit(node)


def process_py_file(file_path):
    print(f"Processing file: {file_path}")

    # Create backup with .bak.py suffix
    base_name = os.path.splitext(file_path)[0]  # Remove .py extension
    backup_path = base_name + ".bak.py"
    shutil.copy2(file_path, backup_path)
    print(f"Created backup: {backup_path}")

    # Read file content
    with open(file_path, "r") as f:
        content = f.read()

    # Find atheris.Setup() call and extract function signature (only the second parameter)
    # 改进的正则表达式，只匹配第二个参数
    setup_pattern = r"atheris\.Setup\([^,]*,\s*([^,)]*)"
    match = re.search(setup_pattern, content)

    if not match:
        print(f"No atheris.Setup() found in {file_path}")
        return

    function_signature = match.group(1).strip()
    print(f"Found function signature: {function_signature}")

    # Parse AST to find target function
    try:
        tree = ast.parse(content)
    except SyntaxError as e:
        print(f"Syntax error in {file_path}: {e}")
        return

    visitor = FunctionVisitor(function_signature)
    visitor.visit(tree)

    if not visitor.found_node:
        print(f"Function {function_signature} not found in {file_path}")
        return

    if not visitor.first_param:
        print(f"No parameters found in function {function_signature}")
        return

    # Create print statement node
    print_stmt = ast.Expr(
        value=ast.Call(
            func=ast.Name(id="print", ctx=ast.Load()),
            args=[ast.Name(id=visitor.first_param, ctx=ast.Load())],
            keywords=[],
        )
    )

    # Insert print statement at the beginning of function body
    if visitor.found_node.body:
        # Preserve docstring if present
        first_item = visitor.found_node.body[0]
        if isinstance(first_item, ast.Expr) and isinstance(first_item.value, ast.Str):
            # Insert after docstring
            visitor.found_node.body.insert(1, print_stmt)
        else:
            # Insert at the very beginning
            visitor.found_node.body.insert(0, print_stmt)
    else:
        visitor.found_node.body = [print_stmt]

    # Generate modified source code
    modified_content = astunparse.unparse(tree)

    # Write modified content back to file
    with open(file_path, "w") as f:
        f.write(modified_content)

    print(f"Added print({visitor.first_param}) to function {function_signature}")


if __name__ == "__main__":
    process_projects()