423 lines
17 KiB
Python
423 lines
17 KiB
Python
import ast
|
|
from utils.encryption import StringEncryptor
|
|
from utils.name_gen import NameGenerator
|
|
import logging
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("RenameTransformer")
|
|
|
|
class RenameTransformer(ast.NodeTransformer):
|
|
def __init__(self, name_generator, global_var_renames, class_attr_mapping,
|
|
primary_key, secondary_key, salt, debug_mode=False):
|
|
self.name_generator = name_generator
|
|
self.global_var_renames = global_var_renames
|
|
self.class_attr_mapping = class_attr_mapping
|
|
|
|
# String encryption
|
|
self.encryptor = StringEncryptor(primary_key, secondary_key, salt)
|
|
|
|
# We'll collect any code needed for key setup (from encrypt_string calls).
|
|
self.key_setup_code = []
|
|
|
|
# Each element is a dict that maps old_name -> new_name for the local scope
|
|
self.scope_stack = [{}]
|
|
|
|
# For class/attribute rename support
|
|
self.in_class = False
|
|
self.current_class_name = None
|
|
# Add this to track methods in each class
|
|
self.class_method_mapping = {}
|
|
# Add a first-pass flag to help with method detection
|
|
self.first_pass = True
|
|
|
|
# Debug data if debug mode is enabled
|
|
self.debug_mode = debug_mode
|
|
if self.debug_mode:
|
|
self.debug_data = {
|
|
"variable_mappings": {},
|
|
"string_encryption": [],
|
|
"renamed_nodes": 0,
|
|
"issues": []
|
|
}
|
|
|
|
def log_debug(self, category: str, data: any):
|
|
"""Log debugging information if debug mode is on."""
|
|
if self.debug_mode:
|
|
if category not in self.debug_data:
|
|
self.debug_data[category] = []
|
|
self.debug_data[category].append(data)
|
|
|
|
def visit_Module(self, node: ast.Module) -> ast.Module:
|
|
"""
|
|
Enhanced module visitor to properly handle global variables
|
|
AND top-level function definitions for consistent renaming.
|
|
|
|
Now includes a two-pass strategy for classes to ensure method consistency.
|
|
"""
|
|
# First pass: detect methods in classes
|
|
if self.first_pass:
|
|
self.first_pass = False
|
|
# First scan all classes and methods recursively
|
|
for stmt in ast.walk(node):
|
|
if isinstance(stmt, ast.ClassDef):
|
|
self.scan_class_methods(stmt)
|
|
|
|
# Rest of the code remains the same...
|
|
# First pass: detect top-level assignments, globals, AND function definitions
|
|
for stmt in node.body:
|
|
# 1. If a global statement
|
|
if isinstance(stmt, ast.Global):
|
|
for name in stmt.names:
|
|
if name not in self.global_var_renames:
|
|
self.global_var_renames[name] = self.name_generator.generate_name()
|
|
|
|
# 2. If a top-level assignment
|
|
elif isinstance(stmt, ast.Assign):
|
|
for target in stmt.targets:
|
|
if isinstance(target, ast.Name):
|
|
if target.id not in self.global_var_renames:
|
|
self.global_var_renames[target.id] = self.name_generator.generate_name()
|
|
|
|
# 3. If a top-level function definition
|
|
elif isinstance(stmt, ast.FunctionDef):
|
|
# Skip special (dunder) methods to avoid messing up e.g. __init__
|
|
if not (stmt.name.startswith('__') and stmt.name.endswith('__')):
|
|
if stmt.name not in self.global_var_renames:
|
|
self.global_var_renames[stmt.name] = self.name_generator.generate_name()
|
|
|
|
# Create a new body that starts with any needed imports
|
|
new_body = [
|
|
ast.parse("import base64").body[0],
|
|
ast.parse("import random").body[0]
|
|
]
|
|
|
|
# Transform the module body
|
|
transformed_body = []
|
|
for item in node.body:
|
|
visited = self.visit(item)
|
|
if isinstance(visited, list):
|
|
transformed_body.extend(visited)
|
|
else:
|
|
transformed_body.append(visited)
|
|
|
|
# If there's any key setup code, parse & insert that
|
|
if self.key_setup_code:
|
|
setup_nodes = ast.parse('\n'.join(self.key_setup_code)).body
|
|
new_body.extend(setup_nodes)
|
|
|
|
new_body.extend(transformed_body)
|
|
node.body = new_body
|
|
return node
|
|
|
|
def scan_class_methods(self, node):
|
|
"""
|
|
Pre-scan a class to identify and map all its methods before actual renaming.
|
|
"""
|
|
class_name = node.name
|
|
# Create a mapping entry for this class if not exists
|
|
if class_name not in self.class_method_mapping:
|
|
self.class_method_mapping[class_name] = {}
|
|
|
|
# Scan all method definitions in the class
|
|
for item in node.body:
|
|
if isinstance(item, ast.FunctionDef):
|
|
method_name = item.name
|
|
# Skip dunder methods
|
|
if not (method_name.startswith('__') and method_name.endswith('__')):
|
|
# Generate a consistent obfuscated name for this method
|
|
new_name = self.name_generator.generate_name()
|
|
self.class_method_mapping[class_name][method_name] = new_name
|
|
|
|
def _push_scope(self):
|
|
self.scope_stack.append({})
|
|
|
|
def _pop_scope(self):
|
|
self.scope_stack.pop()
|
|
|
|
def visit_Global(self, node: ast.Global) -> ast.Global:
|
|
"""
|
|
Handle global statement declarations by:
|
|
1. Adding the variable names to global_var_renames if not already there
|
|
2. Adding them to the current scope to mark them as global
|
|
"""
|
|
for name in node.names:
|
|
# If this global name hasn't been seen before, generate a new obfuscated name
|
|
if name not in self.global_var_renames:
|
|
self.global_var_renames[name] = self.name_generator.generate_name()
|
|
|
|
# Mark this name as global in the current scope
|
|
self.scope_stack[-1][name] = self.global_var_renames[name]
|
|
|
|
# Update the global statement with obfuscated names
|
|
node.names = [self.global_var_renames[name] for name in node.names]
|
|
return node
|
|
|
|
def visit_ListComp(self, node: ast.ListComp) -> ast.ListComp:
|
|
self._push_scope()
|
|
for gen in node.generators:
|
|
gen = self.visit(gen)
|
|
node.elt = self.visit(node.elt)
|
|
self._pop_scope()
|
|
return node
|
|
|
|
def visit_SetComp(self, node: ast.SetComp) -> ast.SetComp:
|
|
self._push_scope()
|
|
for gen in node.generators:
|
|
gen = self.visit(gen)
|
|
node.elt = self.visit(node.elt)
|
|
self._pop_scope()
|
|
return node
|
|
|
|
def visit_DictComp(self, node: ast.DictComp) -> ast.DictComp:
|
|
self._push_scope()
|
|
for gen in node.generators:
|
|
gen = self.visit(gen)
|
|
node.key = self.visit(node.key)
|
|
node.value = self.visit(node.value)
|
|
self._pop_scope()
|
|
return node
|
|
|
|
def visit_GeneratorExp(self, node: ast.GeneratorExp) -> ast.GeneratorExp:
|
|
self._push_scope()
|
|
for gen in node.generators:
|
|
gen = self.visit(gen)
|
|
node.elt = self.visit(node.elt)
|
|
self._pop_scope()
|
|
return node
|
|
|
|
def visit_comprehension(self, node: ast.comprehension) -> ast.comprehension:
|
|
if isinstance(node.target, ast.Name):
|
|
if node.target.id not in self.global_var_renames:
|
|
new_name = self.name_generator.generate_name()
|
|
self.scope_stack[-1][node.target.id] = new_name
|
|
node.target.id = new_name
|
|
else:
|
|
node.target = self.visit(node.target)
|
|
|
|
node.iter = self.visit(node.iter)
|
|
node.ifs = [self.visit(i) for i in node.ifs]
|
|
return node
|
|
|
|
def visit_Name(self, node: ast.Name) -> ast.AST:
|
|
"""
|
|
Handle variable names and function names in calls.
|
|
1. If it's a known global/ top-level function, use global_var_renames.
|
|
2. Otherwise, handle locally within scope_stack.
|
|
"""
|
|
if node.id in self.global_var_renames:
|
|
node.id = self.global_var_renames[node.id]
|
|
return node
|
|
|
|
# Check if this name is marked as global in any scope
|
|
for scope in self.scope_stack:
|
|
if node.id in scope and node.id in self.global_var_renames:
|
|
node.id = self.global_var_renames[node.id]
|
|
return node
|
|
|
|
# Otherwise, handle local variables
|
|
if isinstance(node.ctx, ast.Store):
|
|
if self.in_class and not isinstance(node.ctx, ast.Param):
|
|
# Class attribute
|
|
if node.id not in self.scope_stack[-1]:
|
|
new_name = self.name_generator.generate_name()
|
|
self.scope_stack[-1][node.id] = new_name
|
|
if self.current_class_name:
|
|
self.class_attr_mapping[self.current_class_name][node.id] = new_name
|
|
node.id = self.scope_stack[-1][node.id]
|
|
else:
|
|
# Regular variable assignment
|
|
if node.id not in self.scope_stack[-1]:
|
|
self.scope_stack[-1][node.id] = self.name_generator.generate_name()
|
|
node.id = self.scope_stack[-1][node.id]
|
|
else:
|
|
# Load context
|
|
for scope in reversed(self.scope_stack):
|
|
if node.id in scope:
|
|
node.id = scope[node.id]
|
|
break
|
|
|
|
return node
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
|
|
has_bases = len(node.bases) > 0
|
|
self.scope_stack[-1]['_has_bases'] = has_bases
|
|
prev_in_class = self.in_class
|
|
prev_class_name = self.current_class_name
|
|
|
|
self.in_class = True
|
|
new_class_name = self.name_generator.generate_name()
|
|
self.scope_stack[-1][node.name] = new_class_name
|
|
self.current_class_name = new_class_name
|
|
self.class_attr_mapping[new_class_name] = {}
|
|
|
|
# Transfer method mappings from the original class name to the renamed one
|
|
if node.name in self.class_method_mapping:
|
|
self.class_method_mapping[new_class_name] = self.class_method_mapping[node.name]
|
|
del self.class_method_mapping[node.name]
|
|
|
|
self.scope_stack.append({})
|
|
node.bases = [self.visit(base) for base in node.bases]
|
|
node.body = [self.visit(b) for b in node.body]
|
|
|
|
self.scope_stack[-2][node.name] = new_class_name
|
|
node.name = new_class_name
|
|
|
|
self.in_class = prev_in_class
|
|
self.current_class_name = prev_class_name
|
|
self.scope_stack.pop()
|
|
return node
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
|
|
"""
|
|
Make sure top-level function definitions also get renamed
|
|
if they're in global_var_renames.
|
|
"""
|
|
is_special = node.name.startswith('__') and node.name.endswith('__')
|
|
has_bases = any(('_has_bases' in scope and scope['_has_bases']) for scope in self.scope_stack)
|
|
|
|
# Save original name for later
|
|
orig_name = node.name
|
|
|
|
self._push_scope()
|
|
|
|
# Handle arguments (skip renaming for self if method):
|
|
if self.in_class and len(node.args.args) > 0:
|
|
self.scope_stack[-1][node.args.args[0].arg] = node.args.args[0].arg
|
|
for arg in node.args.args[1:]:
|
|
if arg.arg not in self.global_var_renames:
|
|
new_arg_name = self.name_generator.generate_name()
|
|
self.scope_stack[-1][arg.arg] = new_arg_name
|
|
arg.arg = new_arg_name
|
|
else:
|
|
for arg in node.args.args:
|
|
if arg.arg not in self.global_var_renames:
|
|
new_arg_name = self.name_generator.generate_name()
|
|
self.scope_stack[-1][arg.arg] = new_arg_name
|
|
arg.arg = new_arg_name
|
|
|
|
# Visit body
|
|
node.body = [self.visit(n) for n in node.body]
|
|
|
|
# --------------------------------
|
|
# NEW LOGIC: Actually rename the func if it's top-level
|
|
# or if we want to rename it anyway (and not dunder).
|
|
# --------------------------------
|
|
if not is_special:
|
|
# If the function was declared top-level and recognized in global_var_renames,
|
|
# then rename using that. Otherwise generate a brand new obfuscated name.
|
|
if node.name in self.global_var_renames:
|
|
node.name = self.global_var_renames[node.name]
|
|
elif self.in_class and not has_bases:
|
|
# For class methods, use the name we saved earlier
|
|
if self.current_class_name and orig_name in self.class_method_mapping.get(self.current_class_name, {}):
|
|
node.name = self.class_method_mapping[self.current_class_name][orig_name]
|
|
else:
|
|
new_fn_name = self.name_generator.generate_name()
|
|
self.scope_stack[-2][node.name] = new_fn_name
|
|
node.name = new_fn_name
|
|
|
|
self._pop_scope()
|
|
return node
|
|
|
|
def visit_Attribute(self, node: ast.Attribute) -> ast.AST:
|
|
"""
|
|
Handle attribute access for renaming:
|
|
1. Map method calls on self to the corresponding renamed methods
|
|
2. Preserve external/inherited method names
|
|
3. Rename self.attribute for class-defined attributes
|
|
"""
|
|
# First visit any nested expressions
|
|
node = self.generic_visit(node)
|
|
|
|
# Check if this is a self.something attribute access
|
|
if isinstance(node.value, ast.Name) and node.value.id == 'self':
|
|
# Check all class methods across all classes (more robust)
|
|
for class_name, methods in self.class_method_mapping.items():
|
|
if node.attr in methods:
|
|
# We found a match in our method mapping
|
|
node.attr = methods[node.attr]
|
|
return node
|
|
|
|
# If we're in a class context, apply class-specific logic
|
|
if self.current_class_name:
|
|
# Case 1: Is it a call to one of our renamed methods?
|
|
class_methods = self.class_method_mapping.get(self.current_class_name, {})
|
|
if node.attr in class_methods:
|
|
node.attr = class_methods[node.attr]
|
|
return node
|
|
|
|
# Case 2: Is it an external method call with Qt-style naming?
|
|
is_external_method = (
|
|
node.attr[0].islower() and
|
|
any(c.isupper() for c in node.attr) and
|
|
not node.attr.startswith('__')
|
|
)
|
|
if is_external_method:
|
|
return node
|
|
|
|
# Case 3: Handle normal class attributes
|
|
attr_map = self.class_attr_mapping.get(self.current_class_name, {})
|
|
if node.attr not in attr_map:
|
|
attr_map[node.attr] = self.name_generator.generate_name()
|
|
node.attr = attr_map[node.attr]
|
|
|
|
return node
|
|
|
|
def visit_Subscript(self, node: ast.Subscript) -> ast.AST:
|
|
node.value = self.visit(node.value)
|
|
if isinstance(node.slice, ast.AST):
|
|
self.visit(node.slice)
|
|
return node
|
|
|
|
def visit_Constant(self, node: ast.Constant) -> ast.AST:
|
|
"""
|
|
Encrypt string literals into a multi-step XOR, then base85 decode at runtime.
|
|
"""
|
|
if isinstance(node.value, str):
|
|
encoded, key_setup, modifier = self.encryptor.encrypt_string(node.value)
|
|
if key_setup not in self.key_setup_code:
|
|
self.key_setup_code.append(key_setup)
|
|
|
|
decrypt_str = (
|
|
f"bytes(("
|
|
f"k2^k1^m for k1,k2,m in zip("
|
|
f"bytes(c^k for c,k in zip(base64.b85decode('{encoded}'),"
|
|
f"_sk*((len(base64.b85decode('{encoded}'))//8)+1))),"
|
|
f"_pk*((len(base64.b85decode('{encoded}'))//8)+1),"
|
|
f"bytes.fromhex('{modifier}')*((len(base64.b85decode('{encoded}'))//8)+1)))"
|
|
f").decode()"
|
|
)
|
|
return ast.parse(decrypt_str).body[0].value
|
|
|
|
return node
|
|
|
|
def visit_Import(self, node: ast.Import) -> ast.AST:
|
|
return self.generic_visit(node)
|
|
|
|
def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST:
|
|
return self.generic_visit(node)
|
|
|
|
def visit_Call(self, node: ast.Call) -> ast.Call:
|
|
"""
|
|
Handle super() calls specifically to ensure class names are properly updated.
|
|
"""
|
|
# First visit all arguments and the function itself
|
|
node = self.generic_visit(node)
|
|
|
|
# Check if this is a super() call
|
|
if isinstance(node.func, ast.Name) and node.func.id == 'super':
|
|
# For super() with no args in Python 3
|
|
if not node.args:
|
|
return node
|
|
|
|
# For super(Class, self) style calls
|
|
if len(node.args) >= 1 and isinstance(node.args[0], ast.Name):
|
|
class_name = node.args[0].id
|
|
# Look for the renamed class in all scopes
|
|
for scope in reversed(self.scope_stack):
|
|
if class_name in scope:
|
|
node.args[0].id = scope[class_name]
|
|
break
|
|
|
|
return node |