use ASTfor transform and testgen

Author: joyguoguo

fuzz/collect_fuzz_python.py

@@ -3,8 +3,10 @@
 usage: PYTHONPATH=. python3 fuzz/collect_fuzz_python.py --pipeline all
 """
 from pathlib import Path
+import ast
+import astunparse 
 import logging
-from typing import Optional, List, Tuple
+from typing import Optional
 import fire
 import os
 from UniTSyn.frontend.util import wrap_repo, parallel_subprocess
@@ -159,10 +161,28 @@ def fuzz_repos(repos: list[str], jobs: int, timeout: int):
     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
+        targets = discover_targets(project_name, oss_fuzz_dir)
+        return [generate_test_template(t, repo) for t in targets]
+    
+    with ProcessingPool(jobs) as p:
+        return list(p.map(_transform_repo, repos))
+
 def generate_test_template(target_name: str, repo_path: str):
     """
-    Generate Python test template for a single target by stripping license header,
-    main() block, and converting TestInput/TestOneInput to test_ with data=b"".
+    Generate Python test template using AST for more precise code transformations
     """
     src_file = pjoin(repo_path, target_name + ".py")
     if not os.path.exists(src_file):
@@ -184,68 +204,22 @@ def generate_test_template(target_name: str, repo_path: str):
     )
     code_no_license = re.sub(license_pattern, "", original_code, count=1)
 
-    # --- 2. Remove main function and if __name__ == '__main__' ---
-    code_no_main = re.sub(
-        r"\n?def\s+main\([\s\S]*?(?=^if\s+__name__\s*==\s*['\"]__main__['\"]:)",
-        "",
-        code_no_license,
-        flags=re.MULTILINE
-    )
-    code_no_main = re.sub(
-        r"\n?if\s+__name__\s*==\s*['\"]__main__['\"]:\s*main\(\s*\)\s*",
-        "",
-        code_no_main,
-        flags=re.MULTILINE
-    )
+    # --- 2. Parse code to AST ---
+    try:
+        tree = ast.parse(code_no_license)
+    except SyntaxError as e:
+        logging.error(f"Syntax error in {src_file}: {e}")
+        return None
 
-    # --- 3. Convert test functions ---
-    def process_test_function(match):
-        # Extract the complete function definition and body
-        func_str = match.group(0)
-        
-        # 1. Remove print(data) statements
-        func_str = re.sub(r'print\s*\(\s*data\s*\)\s*', '', func_str)
-        
-        # 2. Change TestInput/TestOneInput to test_()
-        func_str = re.sub(r'def\s+(TestInput|TestOneInput)\s*\(data\)', 'def test_()', func_str)
-        
-        # 3. Insert data = b"" before the first executable line in the function body
-        # Find the first non-empty line (ignoring empty lines and comments)
-        lines = func_str.splitlines()
-        if len(lines) < 2:
-            return func_str
-            
-        # Find the first non-empty, non-comment line after the function definition
-        insert_idx = None
-        for i in range(1, len(lines)):
-            line = lines[i].strip()
-            if line and not line.startswith('#'):
-                insert_idx = i
-                break
-        
-        if insert_idx is None:
-            return func_str
-            
-        # Get the indentation level of that line
-        indent_match = re.match(r'^(\s*)', lines[insert_idx])
-        if not indent_match:
-            return func_str
-            
-        indent = indent_match.group(1)
-        
-        # Insert data = b""
-        lines.insert(insert_idx, f"{indent}data = b\"\"")
-        
-        return "\n".join(lines)
-    
-    cleaned_code = re.sub(
-        r"def\s+(TestInput|TestOneInput)\s*\(data\):[\s\S]*?(?=\n\w|\Z)",
-        process_test_function,
-        code_no_main,
-        flags=re.MULTILINE
-    )
+    # --- 3. AST transformation ---
+    transformer = TestFunctionTransformer()
+    new_tree = transformer.visit(tree)
+    ast.fix_missing_locations(new_tree)
+
+    # --- 4. Generate cleaned code ---
+    cleaned_code = astunparse.unparse(new_tree)
 
-    # --- 4. Output to tests-gen directory ---
+    # --- 5. Output to tests-gen directory ---
     template_dir = pjoin(repo_path, "tests-gen")
     os.makedirs(template_dir, exist_ok=True)
 
@@ -260,26 +234,144 @@ def process_test_function(match):
 
     logging.info(f"Generated cleaned template: {template_path}")
     return template_path
-
-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")
+class TestFunctionTransformer(ast.NodeTransformer):
+    """AST transformer for test function conversion"""
 
-    def _transform_repo(repo: str):
-        project_name = os.path.basename(repo)
-        oss_fuzz_dir = Path(repo).parent.parent
-        targets = discover_targets(project_name, oss_fuzz_dir)
-        return [generate_test_template(t, repo) for t in targets]
+    def visit_FunctionDef(self, node):
+        # 首先处理 main 函数（移除）
+        if node.name == "main":
+            return None
+        
+        # 处理 TestInput/TestOneInput 函数
+        if node.name in ["TestInput", "TestOneInput"]:
+            # a. 记录参数名称（假设只有一个参数）
+            param_name = None
+            if node.args.args:
+                param_name = node.args.args[0].arg
+            
+            # b. 将函数名改为 test_
+            node.name = "test_"
+            
+            # c. 移除参数（将参数列表设为空）
+            node.args = ast.arguments(
+                posonlyargs=[],
+                args=[],
+                vararg=None,
+                kwonlyargs=[],
+                kw_defaults=[],
+                kwarg=None,
+                defaults=[]
+            )
+            
+            # d. 在函数体开头插入 原参数名 = b""
+            if param_name:
+                self.add_param_assignment(node, param_name)
+            
+            # f. 删除所有 print(原参数名) 的语句
+            if param_name:
+                self.remove_print_param(node, param_name)
+        
+        # 确保继续遍历子节点
+        self.generic_visit(node)
+        return node
 
-    with ProcessingPool(jobs) as p:
-        return list(p.map(_transform_repo, repos))
-
+    def add_param_assignment(self, node, param_name):
+        """Add param_name = b"" at the beginning of the function body"""
+        # 创建赋值节点
+        assign_node = ast.Assign(
+            targets=[ast.Name(id=param_name, ctx=ast.Store())],
+            value=ast.Constant(value=b"")
+        )
+        
+        # 如果有文档字符串，插入在文档字符串之后
+        if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Str):
+            node.body.insert(1, assign_node)
+        else:
+            node.body.insert(0, assign_node)
+    
+    def remove_print_param(self, node, param_name):
+        """Remove print statements for the specific parameter"""
+        new_body = []
+        for stmt in node.body:
+            # 跳过 print(param_name) 调用
+            if (isinstance(stmt, ast.Expr) and 
+                isinstance(stmt.value, ast.Call) and
+                isinstance(stmt.value.func, ast.Name) and
+                stmt.value.func.id == "print" and
+                any(isinstance(arg, ast.Name) and arg.id == param_name 
+                    for arg in stmt.value.args)):
+                continue
+            new_body.append(stmt)
+        node.body = new_body
+    
+    def visit_If(self, node):
+        """Remove if __name__ == '__main__' blocks"""
+        # 检查是否是主函数保护
+        if (isinstance(node.test, ast.Compare) and
+            isinstance(node.test.left, ast.Name) and
+            node.test.left.id == "__name__" and
+            isinstance(node.test.ops[0], ast.Eq) and
+            isinstance(node.test.comparators[0], ast.Constant) and
+            node.test.comparators[0].value == "__main__"):
+            
+            # 移除整个 if 块
+            return None
+        
+        # 确保继续遍历子节点
+        self.generic_visit(node)
+        return node
+class TestGenTransformer(ast.NodeTransformer):
+    """AST transformer for generating test cases from fuzzing inputs"""
+    
+    def __init__(self, idx: int, fuzz_input: bytes):
+        self.idx = idx
+        self.fuzz_input = fuzz_input
+        self.found_test_function = False
+    
+    def visit_FunctionDef(self, node):
+        # 只处理名为 test_ 的函数
+        if node.name == "test_":
+            self.found_test_function = True
+            
+            # 1. 将函数名改为 test_{idx}
+            node.name = f"test_{self.idx}"
+            
+            # 2. 找到并替换 data = b"" 赋值语句
+            self.replace_data_assignment(node)
+        
+        return node
+    
+    def replace_data_assignment(self, node):
+        """Replace data assignment with fuzz input"""
+        for i, stmt in enumerate(node.body):
+            # 查找赋值语句
+            if isinstance(stmt, ast.Assign):
+                # 检查是否是 data = b"" 格式的赋值
+                if (len(stmt.targets) == 1 and 
+                    isinstance(stmt.targets[0], ast.Name) and 
+                    isinstance(stmt.value, ast.Constant) and 
+                    stmt.value.value == b""):
+                    
+                    # 替换为新的输入数据
+                    node.body[i] = ast.Assign(
+                        targets=[stmt.targets[0]],
+                        value=ast.Constant(value=self.fuzz_input)
+                    )
+                    return
+                
+                # 检查是否是 data = b'' 格式的赋值
+                if (len(stmt.targets) == 1 and 
+                    isinstance(stmt.targets[0], ast.Name) and 
+                    isinstance(stmt.value, ast.Constant) and 
+                    stmt.value.value == b''):
+                    
+                    # 替换为新的输入数据
+                    node.body[i] = ast.Assign(
+                        targets=[stmt.targets[0]],
+                        value=ast.Constant(value=self.fuzz_input)
+                    )
+                    return
 
 def substitute_one_repo(
     repo: str,
@@ -291,6 +383,7 @@ def substitute_one_repo(
 ):
     """
     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")
@@ -307,15 +400,15 @@ def substitute_one_repo(
             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:
             for line in f_input:
                 try:
-                    # Attempt to decode the line to check content
+                    # 尝试解码行以检查内容
                     decoded = line.decode('utf-8', errors='replace')
 
-                    # Only process lines starting with b' or b"
+                    # 只处理以 b' 或 b" 开头的行
                     if decoded.startswith(("b'", 'b"')):
                         if decoded.startswith("b'") and decoded.endswith("'\n"):
                             byte_data = line[2:-2]
@@ -336,7 +429,7 @@ def substitute_one_repo(
 
         logging.info(f"Loaded {len(valid_inputs)} inputs for {target_name}")
 
-        # Strategy for selecting inputs
+        # 策略选择输入
         if strategy == "shuffle":
             random.shuffle(valid_inputs)
             inputs = valid_inputs[:n_fuzz]
@@ -345,27 +438,42 @@ def substitute_one_repo(
         else:
             inputs = valid_inputs[:n_fuzz]
 
-        # Generate a separate file for each fuzz input
+        # 每个 fuzz input 生成一个单独的文件（使用 AST）
         for idx, fuzz_input in enumerate(inputs, start=1):
             with open(source_file, "r") as f_src:
                 code = f_src.read()
 
-            # 1. Change function name from test_ to test_{idx}
-            code = re.sub(r'def\s+test_', f'def test_{idx}', code)
-            
-            # 2. Replace data = b"" with input data
-            input_repr = repr(fuzz_input)
-            code = code.replace('data = b""', f'data = {input_repr}')
-            
-            out_path = pjoin(template_dir, f"{target_name}.testgen_{idx}.py")
-            with open(out_path, "w") as f_out:
-                f_out.write(code)
-
-            # Format code
             try:
-                subprocess.run(["black", out_path], check=False)
-            except FileNotFoundError:
-                logging.warning("Black formatter not found, skipping formatting")
+                # 解析为 AST
+                tree = ast.parse(code)
+                
+                # 应用转换器
+                transformer = TestGenTransformer(idx, fuzz_input)
+                new_tree = transformer.visit(tree)
+                ast.fix_missing_locations(new_tree)
+                
+                # 确保找到并处理了测试函数
+                if not transformer.found_test_function:
+                    logging.warning(f"No test_ function found in {source_file}")
+                    continue
+                
+                # 生成新代码
+                new_code = astunparse.unparse(new_tree)
+                
+                out_path = pjoin(template_dir, f"{target_name}.testgen_{idx}.py")
+                with open(out_path, "w") as f_out:
+                    f_out.write(new_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:
+                logging.error(f"Error generating test case for {target_name}: {e}")
 def testgen_repos(
     repos: list[str],
     jobs: int,
@@ -405,7 +513,7 @@ def testgen_repos(
         ))
 
 def main(
-    repo_id: str = "data/valid_projects3.txt",
+    repo_id: str = "data/valid_projects.txt",
     repo_root: str = "fuzz/oss-fuzz/projects/",
     timeout: int = 30,
     jobs: int = 8,