da

transformers/class_analyzer.py

@@ -1,9 +1,16 @@
+"""
+@file transformers/class_analyzer.py
+@brief Class analysis utilities for consistent renaming.
+@details Pre-analyzes classes, methods, attributes, and inheritance to build
+         consistent mappings used by obfuscation transformers.
+"""
 import ast
 from typing import Dict, Set, Tuple, List
 
-
 class ClassMethodMap:
-    """Stores method name mappings for all classes in the code."""
+    """
+    @brief Stores method and attribute mappings for all classes.
+    """
     def __init__(self):
         # Maps: original_class_name -> {original_method_name -> obfuscated_method_name}
         self.class_methods: Dict[str, Dict[str, str]] = {}
@@ -20,8 +27,8 @@ class ClassMethodMap:
 
 class ClassAnalyzer(ast.NodeVisitor):
     """
-    Pre-analyzes classes to ensure consistent renaming of methods and attributes.
-    This is crucial for making self.method() calls match def method() definitions.
+    @brief Analyze classes for consistent method/attribute renaming.
+    @details Ensures self.method() calls match def method() definitions and records inheritance.
     """
     def __init__(self, name_generator):
         self.name_generator = name_generator
@@ -35,14 +42,21 @@ class ClassAnalyzer(ast.NodeVisitor):
         self.method_calls: Dict[str, Set[str]] = {}
 
     def analyze(self, tree: ast.AST) -> ClassMethodMap:
-        """Analyzes the entire AST and returns populated method mappings."""
+        """
+        @brief Analyze the entire AST and produce mappings.
+        @param tree Parsed AST of the input module.
+        @return ClassMethodMap Populated mappings.
+        """
         self.visit(tree)
         self._resolve_inheritance()
         self._ensure_consistent_method_mapping()
         return self.method_map
     
     def visit_ClassDef(self, node: ast.ClassDef):
-        """Process a class definition and map its methods."""
+        """
+        @brief Process a class definition and map methods/attrs.
+        @param node ast.ClassDef node being visited.
+        """
         prev_class = self.current_class
         self.current_class = node.name
         
@@ -105,7 +119,10 @@ class ClassAnalyzer(ast.NodeVisitor):
         self.current_class = prev_class
     
     def visit_attribute_assign(self, node):
-        """Process attribute assignments like self.attr = value"""
+        """
+        @brief Process self.attr assignments within a class.
+        @param node ast.Assign node of the assignment.
+        """
         if not self.current_class:
             return
             
@@ -122,7 +139,10 @@ class ClassAnalyzer(ast.NodeVisitor):
         self.visit(node.value)
     
     def visit_Attribute(self, node):
-        """Track self.method references to ensure consistent naming"""
+        """
+        @brief Track self.method references for consistency checks.
+        @param node ast.Attribute being visited.
+        """
         if self.current_class:
             is_self_method, method_name = get_method_name(node)
             if is_self_method:
@@ -134,8 +154,7 @@ class ClassAnalyzer(ast.NodeVisitor):
     
     def _ensure_consistent_method_mapping(self):
         """
-        Make sure that all methods called via self.method() have a mapping,
-        even if they're not defined in the class.
+        @brief Ensure all self.method() calls have mappings, even if undefined in class.
         """
         for class_name, method_calls in self.method_calls.items():
             if class_name not in self.method_map.class_methods:
@@ -153,8 +172,8 @@ class ClassAnalyzer(ast.NodeVisitor):
     
     def _resolve_inheritance(self):
         """
-        Ensure child classes inherit method mappings from parent classes.
-        This ensures that overridden methods use the same obfuscated name.
+        @brief Propagate method mappings along inheritance hierarchies.
+        @details Ensures overridden methods reuse the same obfuscated name.
         """
         # Process inheritance depth-first to handle multi-level inheritance
         def process_inheritance(class_name):
@@ -181,8 +200,9 @@ class ClassAnalyzer(ast.NodeVisitor):
 
 def get_method_name(node: ast.Attribute) -> Tuple[bool, str]:
     """
-    Helper function to determine if an attribute is a self.method() call.
-    Returns (is_self_method, method_name)
+    @brief Determine if an attribute is a self.method() call.
+    @param node Attribute node to check.
+    @return Tuple[bool, str] (is_self_method, method_name)
     """
     if isinstance(node.value, ast.Name) and node.value.id == 'self':
         return True, node.attr
@@ -191,8 +211,9 @@ def get_method_name(node: ast.Attribute) -> Tuple[bool, str]:
 
 def update_obfuscator_with_class_mappings(obfuscator, class_map: ClassMethodMap):
     """
-    Updates the main obfuscator with class method and attribute mappings
-    to ensure consistent renaming across the codebase.
+    @brief Update obfuscator with class/method/attribute mappings.
+    @param obfuscator AdvancedObfuscator instance to update.
+    @param class_map ClassMethodMap mappings produced by analysis.
     """
     # Update class name mappings in global_var_renames
     for orig_name, obf_name in class_map.class_renames.items():