init

transformers/attribute_transformer.py

@@ -0,0 +1,53 @@
+import ast
+from typing import Dict, Optional
+
+class AttributeTransformer(ast.NodeTransformer):
+    """
+    Transforms attribute access expressions (obj.attr) consistently,
+    especially for class method calls and attribute accesses.
+    """
+    
+    def __init__(self, class_attr_mapping: Dict[str, Dict[str, str]], class_renames: Dict[str, str]):
+        """
+        Initialize with mapping dictionaries.
+        
+        Args:
+            class_attr_mapping: Maps class_name -> {attr_name -> obfuscated_attr_name}
+            class_renames: Maps original_class_name -> obfuscated_class_name
+        """
+        self.class_attr_mapping = class_attr_mapping
+        self.class_renames = class_renames
+        self.current_class: Optional[str] = None
+    
+    def visit_ClassDef(self, node):
+        """Keep track of the current class being processed"""
+        old_class = self.current_class
+        
+        # Get the obfuscated name for this class
+        class_name = node.name
+        if class_name in self.class_renames:
+            self.current_class = self.class_renames[class_name]
+        else:
+            self.current_class = class_name
+        
+        # Process the class body
+        node = self.generic_visit(node)
+        
+        # Restore the previous class context
+        self.current_class = old_class
+        return node
+    
+    def visit_Attribute(self, node):
+        """Transform attribute access like self.method to use consistent names"""
+        # First process any nested attributes
+        node = self.generic_visit(node)
+        
+        # Handle self.attr references within a class
+        if isinstance(node.value, ast.Name) and node.value.id == 'self' and self.current_class:
+            if self.current_class in self.class_attr_mapping:
+                attr_map = self.class_attr_mapping[self.current_class]
+                # Only substitute if node.attr is still in its original form (a mapping key)
+                if node.attr in attr_map and node.attr != attr_map[node.attr]:
+                    node.attr = attr_map[node.attr]
+        
+        return node