da

transformers/class_mapper.py

@@ -1,3 +1,10 @@
+"""
+@file transformers/class_mapper.py
+@brief Class, method, and attribute mapping analyzer and transformer.
+@details Builds consistent rename mappings across classes, resolves inheritance,
+         and applies the mappings back to the AST for coherent obfuscation.
+"""
+
 import ast
 from typing import Dict, Set, List, Tuple, Optional
 import logging
@@ -7,7 +14,11 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("ClassMapper")
 
 class ClassMapping:
-    """Stores all mappings related to classes in a centralized way."""
+    """
+    @brief Stores all mappings related to classes in a centralized way.
+    @details Tracks class renames, method and attribute mappings, inheritance structure,
+             and seen method calls to ensure complete coverage.
+    """
     
     def __init__(self):
         # Original class name -> obfuscated class name
@@ -49,7 +60,11 @@ class ClassMapping:
 
 
 class ClassMapAnalyzer(ast.NodeVisitor):
-    """Analyzes the AST to create a complete class mapping."""
+    """
+    @brief Analyze the AST to create a complete class mapping.
+    @details Performs multi-pass analysis to collect classes, methods, attributes,
+             inheritance, and method-call references and builds consistent mappings.
+    """
     
     def __init__(self, name_generator):
         self.name_generator = name_generator
@@ -59,7 +74,11 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         self.processed_classes: Set[str] = set()
     
     def analyze(self, tree: ast.AST) -> ClassMapping:
-        """Perform a complete analysis of the AST."""
+        """
+        @brief Perform a complete analysis of the AST.
+        @param tree Parsed AST to analyze.
+        @return ClassMapping Aggregated mappings for class renaming and members.
+        """
         # First pass: collect all class definitions, methods, and inheritance
         self.visit(tree)
         
@@ -73,7 +92,10 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         return self.mapping
     
     def visit_ClassDef(self, node: ast.ClassDef):
-        """Process a class definition."""
+        """
+        @brief Process a class definition.
+        @param node ast.ClassDef node.
+        """
         prev_class = self.current_class
         self.current_class = node.name
         
@@ -125,7 +147,10 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         self.current_class = prev_class
     
     def visit_method_def(self, node: ast.FunctionDef):
-        """Process a method definition in a class."""
+        """
+        @brief Process a method definition in a class.
+        @param node ast.FunctionDef node.
+        """
         if not self.current_class:
             return
         
@@ -147,7 +172,10 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         self.current_method = prev_method
     
     def visit_assign_in_class(self, node: ast.Assign):
-        """Process assignments in class body or methods."""
+        """
+        @brief Process assignments in class body or methods.
+        @param node ast.Assign possibly containing self.attr writes.
+        """
         if not self.current_class:
             return
         
@@ -166,7 +194,10 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         self.visit(node.value)
     
     def visit_Attribute(self, node: ast.Attribute):
-        """Process attribute access like self.method or self.attr."""
+        """
+        @brief Process attribute access like self.method or self.attr.
+        @param node ast.Attribute node.
+        """
         if self.current_class and isinstance(node.value, ast.Name) and node.value.id == 'self':
             # Record this access for later processing
             method_name = node.attr
@@ -177,7 +208,10 @@ class ClassMapAnalyzer(ast.NodeVisitor):
         self.generic_visit(node)
     
     def visit_Assign(self, node: ast.Assign):
-        """Process assignments that might contain self.attr references."""
+        """
+        @brief Process assignments that might contain self.attr references.
+        @param node ast.Assign node.
+        """
         # Visit both sides of the assignment
         for target in node.targets:
             self.visit(target)
@@ -185,7 +219,8 @@ class ClassMapAnalyzer(ast.NodeVisitor):
     
     def _resolve_inheritance(self):
         """
-        Ensure child classes inherit method mappings from parent classes.
+        @brief Ensure child classes inherit method mappings from parent classes.
+        @details Copies parent method mappings into children when not overridden.
         """
         def process_inheritance(class_name):
             if class_name not in self.mapping.inheritance:
@@ -211,8 +246,8 @@ class ClassMapAnalyzer(ast.NodeVisitor):
     
     def _ensure_complete_method_mapping(self):
         """
-        Make sure all method calls have corresponding mappings.
-        This handles methods called but not defined in the class.
+        @brief Ensure all method calls have corresponding mappings.
+        @details Handles methods called but not defined in the class.
         """
         for class_name, method_calls in self.mapping.seen_method_calls.items():
             if class_name not in self.mapping.class_methods:
@@ -231,14 +266,22 @@ class ClassMapAnalyzer(ast.NodeVisitor):
 
 
 class ClassTransformer(ast.NodeTransformer):
-    """Transforms class-related nodes using the mapping."""
+    """
+    @brief Transform class-related nodes using the mapping.
+    @details Renames class names, methods, and self.attr/self.method references
+             according to the analyzed mappings.
+    """
     
     def __init__(self, mapping: ClassMapping):
         self.mapping = mapping
         self.current_class: Optional[str] = None
     
     def visit_ClassDef(self, node: ast.ClassDef):
-        """Transform class name and process its body."""
+        """
+        @brief Transform class name and process its body.
+        @param node ast.ClassDef node.
+        @return ast.ClassDef Transformed class node.
+        """
         prev_class = self.current_class
         orig_name = node.name
         self.current_class = orig_name
@@ -255,7 +298,11 @@ class ClassTransformer(ast.NodeTransformer):
         return node
     
     def visit_FunctionDef(self, node: ast.FunctionDef):
-        """Transform method name."""
+        """
+        @brief Transform method name.
+        @param node ast.FunctionDef node.
+        @return ast.FunctionDef Transformed method node.
+        """
         if self.current_class and node.name in self.mapping.class_methods.get(self.current_class, {}):
             orig_name = node.name
             node.name = self.mapping.class_methods[self.current_class][node.name]
@@ -266,7 +313,11 @@ class ClassTransformer(ast.NodeTransformer):
         return node
     
     def visit_Attribute(self, node: ast.Attribute):
-        """Transform self.method and self.attr references."""
+        """
+        @brief Transform self.method and self.attr references.
+        @param node ast.Attribute node.
+        @return ast.Attribute Transformed attribute node.
+        """
         # Process any child nodes first (for nested attributes)
         node.value = self.visit(node.value)
         
@@ -288,7 +339,12 @@ class ClassTransformer(ast.NodeTransformer):
 
 # Helper function to apply the class mapping transformation
 def apply_class_mapping(tree: ast.AST, name_generator) -> ast.AST:
-    """Analyze and transform classes consistently."""
+    """
+    @brief Analyze and transform classes consistently.
+    @param tree Input AST.
+    @param name_generator Name generator used to create obfuscated identifiers.
+    @return Tuple[ast.AST, ClassMapping] Transformed AST and the mapping produced.
+    """
     # First pass: analyze all classes
     analyzer = ClassMapAnalyzer(name_generator)
     mapping = analyzer.analyze(tree)