init
This commit is contained in:
@@ -0,0 +1,300 @@
|
||||
import ast
|
||||
from typing import Dict, Set, List, Tuple, Optional
|
||||
import logging
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger("ClassMapper")
|
||||
|
||||
class ClassMapping:
|
||||
"""Stores all mappings related to classes in a centralized way."""
|
||||
|
||||
def __init__(self):
|
||||
# Original class name -> obfuscated class name
|
||||
self.class_names: Dict[str, str] = {}
|
||||
|
||||
# Original class name -> {original method name -> obfuscated method name}
|
||||
self.class_methods: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
# Original class name -> {original attr name -> obfuscated attr name}
|
||||
self.class_attributes: Dict[str, Dict[str, str]] = {}
|
||||
|
||||
# Child class -> list of parent classes (original names)
|
||||
self.inheritance: Dict[str, List[str]] = {}
|
||||
|
||||
# Track all seen method calls to ensure complete coverage
|
||||
self.seen_method_calls: Dict[str, Set[str]] = {}
|
||||
|
||||
def debug_info(self) -> str:
|
||||
"""Return debug information about mappings."""
|
||||
info = []
|
||||
info.append(f"Class mappings: {len(self.class_names)} classes")
|
||||
|
||||
for cls_name, obf_name in self.class_names.items():
|
||||
info.append(f" {cls_name} -> {obf_name}")
|
||||
|
||||
if cls_name in self.class_methods:
|
||||
methods = self.class_methods[cls_name]
|
||||
info.append(f" Methods: {len(methods)}")
|
||||
for method, obf_method in methods.items():
|
||||
info.append(f" {method} -> {obf_method}")
|
||||
|
||||
if cls_name in self.class_attributes:
|
||||
attrs = self.class_attributes[cls_name]
|
||||
info.append(f" Attributes: {len(attrs)}")
|
||||
for attr, obf_attr in attrs.items():
|
||||
info.append(f" {attr} -> {obf_attr}")
|
||||
|
||||
return "\n".join(info)
|
||||
|
||||
|
||||
class ClassMapAnalyzer(ast.NodeVisitor):
|
||||
"""Analyzes the AST to create a complete class mapping."""
|
||||
|
||||
def __init__(self, name_generator):
|
||||
self.name_generator = name_generator
|
||||
self.mapping = ClassMapping()
|
||||
self.current_class: Optional[str] = None
|
||||
self.current_method: Optional[str] = None
|
||||
self.processed_classes: Set[str] = set()
|
||||
|
||||
def analyze(self, tree: ast.AST) -> ClassMapping:
|
||||
"""Perform a complete analysis of the AST."""
|
||||
# First pass: collect all class definitions, methods, and inheritance
|
||||
self.visit(tree)
|
||||
|
||||
# Second pass: resolve inheritance and method mappings
|
||||
self._resolve_inheritance()
|
||||
self._ensure_complete_method_mapping()
|
||||
|
||||
logger.info(f"Class analysis complete: {len(self.mapping.class_names)} classes processed")
|
||||
logger.debug(self.mapping.debug_info())
|
||||
|
||||
return self.mapping
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef):
|
||||
"""Process a class definition."""
|
||||
prev_class = self.current_class
|
||||
self.current_class = node.name
|
||||
|
||||
# Skip if already processed
|
||||
if node.name in self.processed_classes:
|
||||
self.current_class = prev_class
|
||||
return
|
||||
|
||||
# Add class name mapping
|
||||
if node.name not in self.mapping.class_names:
|
||||
self.mapping.class_names[node.name] = self.name_generator.generate_name()
|
||||
|
||||
# Initialize dictionaries
|
||||
if node.name not in self.mapping.class_methods:
|
||||
self.mapping.class_methods[node.name] = {}
|
||||
|
||||
if node.name not in self.mapping.class_attributes:
|
||||
self.mapping.class_attributes[node.name] = {}
|
||||
|
||||
if node.name not in self.mapping.seen_method_calls:
|
||||
self.mapping.seen_method_calls[node.name] = set()
|
||||
|
||||
# Record inheritance
|
||||
parent_classes = []
|
||||
for base in node.bases:
|
||||
if isinstance(base, ast.Name):
|
||||
parent_classes.append(base.id)
|
||||
|
||||
if parent_classes:
|
||||
self.mapping.inheritance[node.name] = parent_classes
|
||||
|
||||
# Process class body
|
||||
for item in node.body:
|
||||
if isinstance(item, ast.FunctionDef):
|
||||
self.visit_method_def(item)
|
||||
elif isinstance(item, ast.Assign):
|
||||
self.visit_assign_in_class(item)
|
||||
elif isinstance(item, ast.Expr):
|
||||
# Could contain calls to self.methods
|
||||
self.visit(item)
|
||||
elif isinstance(item, ast.ClassDef):
|
||||
# Nested class
|
||||
self.visit(item)
|
||||
else:
|
||||
# Other nodes that might contain self.method calls
|
||||
self.visit(item)
|
||||
|
||||
self.processed_classes.add(node.name)
|
||||
self.current_class = prev_class
|
||||
|
||||
def visit_method_def(self, node: ast.FunctionDef):
|
||||
"""Process a method definition in a class."""
|
||||
if not self.current_class:
|
||||
return
|
||||
|
||||
prev_method = self.current_method
|
||||
self.current_method = node.name
|
||||
|
||||
# Skip dunder methods from obfuscation
|
||||
if not (node.name.startswith('__') and node.name.endswith('__')):
|
||||
# Map method name if not already mapped
|
||||
if node.name not in self.mapping.class_methods[self.current_class]:
|
||||
obf_name = self.name_generator.generate_name()
|
||||
self.mapping.class_methods[self.current_class][node.name] = obf_name
|
||||
logger.debug(f"Mapped method {self.current_class}.{node.name} to {obf_name}")
|
||||
|
||||
# Visit method body to find self.method calls and self.attr assignments
|
||||
for item in node.body:
|
||||
self.visit(item)
|
||||
|
||||
self.current_method = prev_method
|
||||
|
||||
def visit_assign_in_class(self, node: ast.Assign):
|
||||
"""Process assignments in class body or methods."""
|
||||
if not self.current_class:
|
||||
return
|
||||
|
||||
# Check for self.attribute assignments
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == 'self':
|
||||
attr_name = target.attr
|
||||
|
||||
# Map attribute name if not already mapped
|
||||
if attr_name not in self.mapping.class_attributes[self.current_class]:
|
||||
obf_name = self.name_generator.generate_name()
|
||||
self.mapping.class_attributes[self.current_class][attr_name] = obf_name
|
||||
logger.debug(f"Mapped attribute {self.current_class}.{attr_name} to {obf_name}")
|
||||
|
||||
# Visit the value to find nested self.method calls
|
||||
self.visit(node.value)
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute):
|
||||
"""Process attribute access like self.method or self.attr."""
|
||||
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
|
||||
self.mapping.seen_method_calls[self.current_class].add(method_name)
|
||||
logger.debug(f"Recorded method call: {self.current_class}.{method_name}")
|
||||
|
||||
# Continue traversal
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Assign(self, node: ast.Assign):
|
||||
"""Process assignments that might contain self.attr references."""
|
||||
# Visit both sides of the assignment
|
||||
for target in node.targets:
|
||||
self.visit(target)
|
||||
self.visit(node.value)
|
||||
|
||||
def _resolve_inheritance(self):
|
||||
"""
|
||||
Ensure child classes inherit method mappings from parent classes.
|
||||
"""
|
||||
def process_inheritance(class_name):
|
||||
if class_name not in self.mapping.inheritance:
|
||||
return
|
||||
|
||||
for parent in self.mapping.inheritance[class_name]:
|
||||
# Process parent's inheritance first
|
||||
process_inheritance(parent)
|
||||
|
||||
# Skip if parent isn't in our mappings
|
||||
if parent not in self.mapping.class_methods:
|
||||
continue
|
||||
|
||||
# Copy parent's method mappings to child if not overridden
|
||||
for method_name, obf_name in self.mapping.class_methods[parent].items():
|
||||
if method_name not in self.mapping.class_methods[class_name]:
|
||||
self.mapping.class_methods[class_name][method_name] = obf_name
|
||||
logger.debug(f"Inherited method {class_name}.{method_name} from {parent}")
|
||||
|
||||
# Process all classes
|
||||
for class_name in list(self.mapping.class_methods.keys()):
|
||||
process_inheritance(class_name)
|
||||
|
||||
def _ensure_complete_method_mapping(self):
|
||||
"""
|
||||
Make sure all method calls have corresponding mappings.
|
||||
This 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:
|
||||
continue
|
||||
|
||||
for method_name in method_calls:
|
||||
# Skip dunder methods
|
||||
if method_name.startswith('__') and method_name.endswith('__'):
|
||||
continue
|
||||
|
||||
# Add mapping if method was called but not defined
|
||||
if method_name not in self.mapping.class_methods[class_name]:
|
||||
obf_name = self.name_generator.generate_name()
|
||||
self.mapping.class_methods[class_name][method_name] = obf_name
|
||||
logger.debug(f"Added mapping for called method {class_name}.{method_name} -> {obf_name}")
|
||||
|
||||
|
||||
class ClassTransformer(ast.NodeTransformer):
|
||||
"""Transforms class-related nodes using the mapping."""
|
||||
|
||||
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."""
|
||||
prev_class = self.current_class
|
||||
orig_name = node.name
|
||||
self.current_class = orig_name
|
||||
|
||||
# Rename class if it's in our mapping
|
||||
if node.name in self.mapping.class_names:
|
||||
node.name = self.mapping.class_names[node.name]
|
||||
logger.debug(f"Transformed class {orig_name} -> {node.name}")
|
||||
|
||||
# Process class body
|
||||
node.body = [self.visit(item) for item in node.body]
|
||||
|
||||
self.current_class = prev_class
|
||||
return node
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef):
|
||||
"""Transform method name."""
|
||||
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]
|
||||
logger.debug(f"Transformed method {self.current_class}.{orig_name} -> {node.name}")
|
||||
|
||||
# Visit the method body
|
||||
node.body = [self.visit(item) for item in node.body]
|
||||
return node
|
||||
|
||||
def visit_Attribute(self, node: ast.Attribute):
|
||||
"""Transform self.method and self.attr references."""
|
||||
# Process any child nodes first (for nested attributes)
|
||||
node.value = self.visit(node.value)
|
||||
|
||||
# Check if this is a self.attr or self.method reference
|
||||
if self.current_class and isinstance(node.value, ast.Name) and node.value.id == 'self':
|
||||
orig_name = node.attr
|
||||
|
||||
# Check in method mappings first
|
||||
if self.current_class in self.mapping.class_methods and node.attr in self.mapping.class_methods[self.current_class]:
|
||||
node.attr = self.mapping.class_methods[self.current_class][node.attr]
|
||||
logger.debug(f"Transformed self.method {self.current_class}.{orig_name} -> {node.attr}")
|
||||
|
||||
# Then check attribute mappings
|
||||
elif self.current_class in self.mapping.class_attributes and node.attr in self.mapping.class_attributes[self.current_class]:
|
||||
node.attr = self.mapping.class_attributes[self.current_class][node.attr]
|
||||
logger.debug(f"Transformed self.attr {self.current_class}.{orig_name} -> {node.attr}")
|
||||
|
||||
return node
|
||||
|
||||
# Helper function to apply the class mapping transformation
|
||||
def apply_class_mapping(tree: ast.AST, name_generator) -> ast.AST:
|
||||
"""Analyze and transform classes consistently."""
|
||||
# First pass: analyze all classes
|
||||
analyzer = ClassMapAnalyzer(name_generator)
|
||||
mapping = analyzer.analyze(tree)
|
||||
|
||||
# Second pass: transform using the mapping
|
||||
transformer = ClassTransformer(mapping)
|
||||
transformed = transformer.visit(tree)
|
||||
|
||||
return transformed, mapping
|
||||
Reference in New Issue
Block a user