use AST

Author: joyguoguo

fuzz/modify_fuzz_files.py

@@ -1,34 +1,41 @@
 import os
-import re
+import ast
 
 def add_print_to_testoneinput(file_path):
     with open(file_path, 'r') as f:
         content = f.read()
 
-    # 正则表达式匹配TestOneInput或TestInput函数定义及其函数体
-    pattern = r'(\bdef\s+(TestOneInput|TestInput)\(data\):\s*\n)((?:[ \t]+.*\n|\s*\n)*)'
-    matches = re.finditer(pattern, content, re.MULTILINE)
+    # 解析 AST
+    tree = ast.parse(content)
 
-    new_content = content
-    for match in reversed(list(matches)):
-        function_def = match.group(1)
-        function_body = match.group(3)
-        
-        # 在函数体开头添加print(data)语句
-        new_function_body = re.sub(
-            r'^([ \t]*)(.*\n)', 
-            r'\g<1>\2\g<1>print(data)\n', 
-            function_body, 
-            count=1
-        )
-        
-        # 只有在函数体非空且未添加过print时才替换
-        if new_function_body != function_body:
-            new_content = (
-                new_content[:match.start(3)] + 
-                new_function_body + 
-                new_content[match.end(3):]
-            )
+    class InsertPrintTransformer(ast.NodeTransformer):
+        def visit_FunctionDef(self, node):
+            if node.name in ("TestOneInput", "TestInput") and node.args.args:
+                first_arg_name = node.args.args[0].arg
+                # 创建 print(参数名) 语句
+                print_stmt = ast.Expr(
+                    value=ast.Call(
+                        func=ast.Name(id='print', ctx=ast.Load()),
+                        args=[ast.Name(id=first_arg_name, ctx=ast.Load())],
+                        keywords=[]
+                    )
+                )
+                # 确保没有重复插入
+                if not (
+                    isinstance(node.body[0], ast.Expr)
+                    and isinstance(node.body[0].value, ast.Call)
+                    and getattr(node.body[0].value.func, "id", None) == "print"
+                ):
+                    node.body.insert(0, print_stmt)
+            return node
+
+    transformer = InsertPrintTransformer()
+    new_tree = transformer.visit(tree)
+    ast.fix_missing_locations(new_tree)
+
+    # 转回代码
+    import astor
+    new_content = astor.to_source(new_tree)
 
     return new_content
 
@@ -41,26 +48,26 @@ def main():
 
     for project in projects:
         project_dir = os.path.join(projects_path, project)
-        
+
         if not os.path.isdir(project_dir):
             continue
 
         for root, _, files in os.walk(project_dir):
             for file in files:
                 if file.startswith('fuzz_') and file.endswith('.py'):
                     file_path = os.path.join(root, file)
-                    
+
                     try:
                         new_content = add_print_to_testoneinput(file_path)
-                        
-                        # 保存修改后的文件（添加_print后缀）
+
+                        # 保存修改后的文件
                         new_file_path = file_path.rsplit('.', 1)[0] + '_print1.py'
                         with open(new_file_path, 'w') as f:
                             f.write(new_content)
                         print(f"Processed: {file_path} -> {new_file_path}")
-                            
+
                     except Exception as e:
                         print(f"Error processing {file_path}: {str(e)}")
 
 if __name__ == "__main__":
-    main()
+    main()