remove the class outside of the function

Author: joyguoguo

fuzz/modify_fuzz_files.py

@@ -3,39 +3,40 @@
 import ast
 import fire
 
+
+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_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=[]
+                )
+            )
+            # 添加空body检查
+            if not node.body:
+                node.body.append(print_stmt)
+            else:
+                # 增强重复检查逻辑
+                first_stmt = node.body[0]
+                if not (isinstance(first_stmt, ast.Expr) 
+                        and isinstance(first_stmt.value, ast.Call)
+                        and hasattr(first_stmt.value.func, 'id')
+                        and first_stmt.value.func.id == 'print'):
+                    node.body.insert(0, print_stmt)
+        return node
+
 def add_print_to_testoneinput(file_path):
     with open(file_path, 'r') as f:
         content = f.read()
 
-    # 解析 AST
     tree = ast.parse(content)
-
-    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
@@ -44,31 +45,21 @@ def main(
     projects_path="fuzz/oss-fuzz/projects",
     valid_projects_file="data/valid_projects.txt"
 ):
-    """
-    给 fuzz target 的 TestOneInput / TestInput 函数开头插入 print(参数名)
-
-    Args:
-        projects_path (str): OSS-Fuzz 项目的根目录
-        valid_projects_file (str): 包含有效项目名的文件路径
-    """
+    """为fuzz target添加打印语句"""
     with open(valid_projects_file, 'r') as f:
         projects = [line.strip() for line in f if line.strip()]
 
     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)
-
-                        # 保存修改后的文件
                         new_file_path = file_path.rsplit('.', 1)[0] + '_print1.py'
                         with open(new_file_path, 'w') as f:
                             f.write(new_content)
@@ -78,4 +69,4 @@ def main(
                         print(f"Error processing {file_path}: {str(e)}")
 
 if __name__ == "__main__":
-    fire.Fire(main)
+    fire.Fire(main)