54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
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
|