68 lines
2.6 KiB
Python
68 lines
2.6 KiB
Python
"""
|
|
@file transformers/attribute_transformer.py
|
|
@brief Attribute access transformer for consistent class member renaming.
|
|
@details Ensures references like self.method and self.attr are kept consistent
|
|
with class and attribute mappings produced by analysis.
|
|
"""
|
|
|
|
import ast
|
|
from typing import Dict, Optional
|
|
|
|
class AttributeTransformer(ast.NodeTransformer):
|
|
"""
|
|
@brief Transform attribute access expressions for consistency.
|
|
@details Handles self.method and self.attr translations based on provided
|
|
mappings to keep calls and attributes aligned with obfuscated names.
|
|
"""
|
|
|
|
def __init__(self, class_attr_mapping: Dict[str, Dict[str, str]], class_renames: Dict[str, str]):
|
|
"""
|
|
@brief Initialize transformer with mapping dictionaries.
|
|
@param class_attr_mapping Maps class_name -> {attr_name -> obfuscated_attr_name}
|
|
@param 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):
|
|
"""
|
|
@brief Track the current class and process its body.
|
|
@param node ast.ClassDef being visited.
|
|
@return ast.ClassDef Potentially modified node.
|
|
"""
|
|
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):
|
|
"""
|
|
@brief Transform attribute access for consistency within a class.
|
|
@param node ast.Attribute node to transform.
|
|
@return ast.AST Updated attribute node.
|
|
"""
|
|
# 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
|