init
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
import ast
|
||||
from typing import Dict, Set, Tuple, List
|
||||
|
||||
|
||||
class ClassMethodMap:
|
||||
"""Stores method name mappings for all classes in the code."""
|
||||
def __init__(self):
|
||||
# Maps: original_class_name -> {original_method_name -> obfuscated_method_name}
|
||||
self.class_methods: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
# Maps: original_class_name -> {original_attr_name -> obfuscated_attr_name}
|
||||
self.class_attributes: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
# Maps: original_class_name -> obfuscated_class_name
|
||||
self.class_renames: Dict[str, str] = {}
|
||||
|
||||
# Track inheritance relationships: child_class -> [parent_classes]
|
||||
self.inheritance: Dict[str, List[str]] = {}
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
def __init__(self, name_generator):
|
||||
self.name_generator = name_generator
|
||||
self.method_map = ClassMethodMap()
|
||||
self.current_class = None
|
||||
|
||||
# To avoid duplicate scanning
|
||||
self.scanned_classes: Set[str] = set()
|
||||
|
||||
# Track method calls within each class
|
||||
self.method_calls: Dict[str, Set[str]] = {}
|
||||
|
||||
def analyze(self, tree: ast.AST) -> ClassMethodMap:
|
||||
"""Analyzes the entire AST and returns populated method 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."""
|
||||
prev_class = self.current_class
|
||||
self.current_class = node.name
|
||||
|
||||
# Skip if already processed this class
|
||||
if node.name in self.scanned_classes:
|
||||
self.current_class = prev_class
|
||||
return
|
||||
|
||||
# Initialize method calls tracking for this class
|
||||
self.method_calls[node.name] = set()
|
||||
|
||||
# Record class inheritance
|
||||
parent_classes = []
|
||||
for base in node.bases:
|
||||
if isinstance(base, ast.Name):
|
||||
parent_classes.append(base.id)
|
||||
|
||||
if parent_classes:
|
||||
self.method_map.inheritance[node.name] = parent_classes
|
||||
|
||||
# Initialize mappings for this class
|
||||
if node.name not in self.method_map.class_methods:
|
||||
self.method_map.class_methods[node.name] = {}
|
||||
|
||||
if node.name not in self.method_map.class_attributes:
|
||||
self.method_map.class_attributes[node.name] = {}
|
||||
|
||||
# Create a consistent obfuscated name for this class
|
||||
if node.name not in self.method_map.class_renames:
|
||||
new_name = self.name_generator.generate_name()
|
||||
self.method_map.class_renames[node.name] = new_name
|
||||
|
||||
# Process all method definitions in the class
|
||||
for item in node.body:
|
||||
# Methods
|
||||
if isinstance(item, ast.FunctionDef):
|
||||
# Skip dunder methods
|
||||
if not (item.name.startswith('__') and item.name.endswith('__')):
|
||||
# Generate a consistent obfuscated name for this method
|
||||
new_name = self.name_generator.generate_name()
|
||||
self.method_map.class_methods[node.name][item.name] = new_name
|
||||
|
||||
# Visit the method body to find self.method() calls
|
||||
self.visit(item)
|
||||
|
||||
# Attributes in assignments
|
||||
elif isinstance(item, ast.Assign):
|
||||
self.visit_attribute_assign(item)
|
||||
else:
|
||||
# Visit other nodes (like if statements that might contain self.method calls)
|
||||
self.visit(item)
|
||||
|
||||
self.scanned_classes.add(node.name)
|
||||
|
||||
# Visit any nested classes
|
||||
for item in node.body:
|
||||
if isinstance(item, ast.ClassDef):
|
||||
self.visit(item)
|
||||
|
||||
self.current_class = prev_class
|
||||
|
||||
def visit_attribute_assign(self, node):
|
||||
"""Process attribute assignments like self.attr = value"""
|
||||
if not self.current_class:
|
||||
return
|
||||
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name):
|
||||
if target.value.id == 'self':
|
||||
# This is a self.attribute assignment
|
||||
attr_name = target.attr
|
||||
if attr_name not in self.method_map.class_attributes[self.current_class]:
|
||||
new_name = self.name_generator.generate_name()
|
||||
self.method_map.class_attributes[self.current_class][attr_name] = new_name
|
||||
|
||||
# Visit the value part of the assignment for nested self.method() calls
|
||||
self.visit(node.value)
|
||||
|
||||
def visit_Attribute(self, node):
|
||||
"""Track self.method references to ensure consistent naming"""
|
||||
if self.current_class:
|
||||
is_self_method, method_name = get_method_name(node)
|
||||
if is_self_method:
|
||||
# Record this method call for consistency checks later
|
||||
self.method_calls[self.current_class].add(method_name)
|
||||
|
||||
# Continue traversing
|
||||
self.generic_visit(node)
|
||||
|
||||
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.
|
||||
"""
|
||||
for class_name, method_calls in self.method_calls.items():
|
||||
if class_name not in self.method_map.class_methods:
|
||||
continue
|
||||
|
||||
class_methods = self.method_map.class_methods[class_name]
|
||||
for method_name in method_calls:
|
||||
if method_name not in class_methods:
|
||||
# Skip dunder methods
|
||||
if method_name.startswith('__') and method_name.endswith('__'):
|
||||
continue
|
||||
# Existing check: mapping is generated only once.
|
||||
new_name = self.name_generator.generate_name()
|
||||
class_methods[method_name] = new_name
|
||||
|
||||
def _resolve_inheritance(self):
|
||||
"""
|
||||
Ensure child classes inherit method mappings from parent classes.
|
||||
This ensures that overridden methods use the same obfuscated name.
|
||||
"""
|
||||
# Process inheritance depth-first to handle multi-level inheritance
|
||||
def process_inheritance(class_name):
|
||||
if class_name not in self.method_map.inheritance:
|
||||
return
|
||||
|
||||
for parent in self.method_map.inheritance[class_name]:
|
||||
# Process parent's inheritance first
|
||||
process_inheritance(parent)
|
||||
|
||||
# Skip if parent isn't in our mappings (external class)
|
||||
if parent not in self.method_map.class_methods:
|
||||
continue
|
||||
|
||||
# Inherit parent's methods if not overridden
|
||||
for method_name, obf_name in self.method_map.class_methods[parent].items():
|
||||
if method_name not in self.method_map.class_methods[class_name]:
|
||||
self.method_map.class_methods[class_name][method_name] = obf_name
|
||||
|
||||
# Process inheritance for each class
|
||||
for class_name in list(self.method_map.class_methods.keys()):
|
||||
process_inheritance(class_name)
|
||||
|
||||
|
||||
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)
|
||||
"""
|
||||
if isinstance(node.value, ast.Name) and node.value.id == 'self':
|
||||
return True, node.attr
|
||||
return False, ""
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
# Update class name mappings in global_var_renames
|
||||
for orig_name, obf_name in class_map.class_renames.items():
|
||||
obfuscator.global_var_renames[orig_name] = obf_name
|
||||
|
||||
# Update class attr mapping with our analyzed data
|
||||
for class_name, class_obf_name in class_map.class_renames.items():
|
||||
# Initialize if needed
|
||||
if class_obf_name not in obfuscator.class_attr_mapping:
|
||||
obfuscator.class_attr_mapping[class_obf_name] = {}
|
||||
|
||||
# Copy method mappings
|
||||
if class_name in class_map.class_methods:
|
||||
for method, obf_method in class_map.class_methods[class_name].items():
|
||||
obfuscator.class_attr_mapping[class_obf_name][method] = obf_method
|
||||
|
||||
# Copy attribute mappings
|
||||
if class_name in class_map.class_attributes:
|
||||
for attr, obf_attr in class_map.class_attributes[class_name].items():
|
||||
obfuscator.class_attr_mapping[class_obf_name][attr] = obf_attr
|
||||
Reference in New Issue
Block a user